import numpy as np
import axengine
from .preprocess import preprocess
from .postprocess import postprocess

class ResNet50Classifier:
    def __init__(self, model_path):
        self.session = axengine.InferenceSession(model_path)
        self.input_name = self.session.get_inputs()[0].name
        self.output_name = self.session.get_outputs()[0].name
        
    def classify(self, image_path, top_k=5):
        input_data = preprocess(image_path)
        outputs = self.session.run(None, {self.input_name: input_data})
        return postprocess(outputs[0], top_k)