File size: 569 Bytes
cdf6b61 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 | 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)
|