model added

handler.py

@@ -0,0 +1,55 @@
+from typing import Dict, List, Any
+from transformers import AutoModelForImageClassification, AutoFeatureExtractor
+import torch
+from PIL import Image
+import io
+
+class EndpointHandler:
+    def __init__(self):
+        # Initialize model and feature extractor
+        model_id = "alexionby/ainoai"
+        self.model = AutoModelForImageClassification.from_pretrained(model_id)
+        self.feature_extractor = AutoFeatureExtractor.from_pretrained(model_id)
+
+    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
+        # Convert bytes to PIL Image
+        image = data.pop('image', data)
+        image = Image.open(io.BytesIO(image))
+
+        # Preprocess the image
+        inputs = self.feature_extractor(images=image, return_tensors="pt")
+
+        # Run the model
+        with torch.no_grad():
+            outputs = self.model(**inputs)
+
+        # Post-process the model outputs as needed
+        logits = outputs.logits
+        probabilities = torch.nn.functional.softmax(logits, dim=-1)
+        predictions = probabilities.argmax(-1)
+
+        # Convert predictions to JSON-serializable format
+        return {"label": str(predictions.item())}
+
+
+# import torch
+# from PIL import Image
+# import io
+
+# class EndpointHandler():
+#     def __init__(self, path=""):
+#         # Preload all the elements you are going to need at inference.
+#         # pseudo:
+#         self.model= load_model(path)
+
+#     def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
+#         """
+#        data args:
+#             inputs (:obj: `str` | `PIL.Image` | `np.array`)
+#             kwargs
+#       Return:
+#             A :obj:`list` | `dict`: will be serialized and returned
+#         """
+
+#         # pseudo
+#         # self.model(input)

test.py

@@ -0,0 +1,9 @@
+from handler import EndpointHandler
+
+# init handler
+my_handler = EndpointHandler()
+
+from PIL import Image
+image_bytes = image.open("imgs/optimus.jpg")
+
+print(my_handler(image_bytes))