Update handler.py

handler.py

@@ -5,6 +5,33 @@ from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
 # check for GPU
 device = 0 if torch.cuda.is_available() else -1
 
+
+from PIL import Image
+from torchvision.transforms import Compose, ConvertImageDtype, Normalize, PILToTensor, Resize
+from torchvision.transforms.functional import InterpolationMode
+from pyrovision.models import model_from_hf_hub
+
+# model = model_from_hf_hub("pyronear/mobilenet_v3_small").eval()
+
+# img = Image.open(path_to_an_image).convert("RGB")
+
+# # Preprocessing
+# config = model.default_cfg
+# transform = Compose([
+#     Resize(config['input_shape'][1:], interpolation=InterpolationMode.BILINEAR),
+#     PILToTensor(),
+#     ConvertImageDtype(torch.float32),
+#     Normalize(config['mean'], config['std'])
+# ])
+
+# input_tensor = transform(img).unsqueeze(0)
+
+# # Inference
+# with torch.inference_mode():
+#     output = model(input_tensor)
+# probs = output.squeeze(0).softmax(dim=0)
+
+
 # multi-model list
 multi_model_list = [
     {"model_id": "distilbert-base-uncased-finetuned-sst-2-english", "task": "text-classification"},
@@ -32,9 +59,10 @@ class EndpointHandler():
             raise ValueError(f"model_id: {model_id} is not valid. Available models are: {list(self.multi_model.keys())}")
 
         # pass inputs with all kwargs in data
-        if parameters is not None:
-            prediction = self.multi_model[model_id](inputs, **parameters)
-        else:
-            prediction = self.multi_model[model_id](inputs)
-        # postprocess the prediction
+        prediction = {'output':'test'}
+        # if parameters is not None:
+        #     prediction = self.multi_model[model_id](inputs, **parameters)
+        # else:
+        #     prediction = self.multi_model[model_id](inputs)
+        # # postprocess the prediction
         return prediction