"""MobileNetV3-Small classification inference on AX650."""
import numpy as np
import axengine

class MobileNetV3Classifier:
    def __init__(self, model_path: str):
        self.sess = axengine.InferenceSession(model_path)
        self.input_name = self.sess.get_inputs()[0].name
        
    def classify(self, image: np.ndarray) -> np.ndarray:
        """image: float32 NCHW [1,3,224,224], range [0,1]"""
        return self.sess.run(None, {self.input_name: image})[0]