import onnxruntime as ort
from utils.config import IMAGE_ENCODER
from utils.preprocessing import ImageProcess


class ImageEncoder:
    def __init__(self):
        self.image_model = ort.InferenceSession(IMAGE_ENCODER)
        self.preprocess = ImageProcess()
        self.input_name = self.image_model.get_inputs()[0].name

    def get_image_embedding(self, path):
        image_tensor = self.preprocess.process(path)

        image_embeddings = self.image_model.run(
            None, {self.input_name: image_tensor})[0]
        return image_embeddings.squeeze(0)