Updated handler with EndpointHandler class

handler.py

@@ -5,7 +5,7 @@ import torchvision.transforms as transforms
 from PIL import Image
 import json
 
-# Hugging Face model URL
+# Hugging Face Model Information
 HF_REPO = "your-username/your-model-name"
 MODEL_FILENAME = "model_scripted_efficientnet.pt"
 MODEL_PATH = f"/repository/{MODEL_FILENAME}"
@@ -54,20 +54,28 @@ supported_issues = [
     "Brow Asymmetry"
 ]
 
-def predict(image_path):
-    image = Image.open(image_path).convert("RGB")
-    image = transform(image).unsqueeze(0).to(device)
+class EndpointHandler:
+    def __init__(self, model_dir=None):
+        """Initialize the inference model."""
+        self.model = model
+        self.device = device
+        self.model.to(self.device)
+        self.model.eval()
+        print("Model loaded successfully.")
 
-    with torch.no_grad():
-        outputs = model(image)
+    def __call__(self, data):
+        """Perform inference on an image."""
+        if "image" not in data:
+            return {"error": "No image provided"}
 
-    predictions = outputs.squeeze().tolist()
-    output_labels = [label for label, prob in zip(supported_issues, predictions) if prob > 0.5]
+        image_data = data["image"]
+        image = Image.open(image_data).convert("RGB")
+        image = transform(image).unsqueeze(0).to(self.device)
 
-    return {"predictions": output_labels}
+        with torch.no_grad():
+            outputs = self.model(image)
 
-# Hugging Face API Inference Handler
-def handle_request(data):
-    image_path = data["image_path"]
-    results = predict(image_path)
-    return json.dumps(results)
+        predictions = outputs.squeeze().tolist()
+        output_labels = [label for label, prob in zip(supported_issues, predictions) if prob > 0.5]
+
+        return {"predictions": output_labels}