revise handler based on deepseek attempt

handler.py

@@ -1,81 +1,64 @@
-import os
 import torch
-import requests
 import torchvision.transforms as transforms
 from PIL import Image
-import json
+import io
 
-# Hugging Face Model Information
-HF_REPO = "your-username/your-model-name"
-MODEL_FILENAME = "model_scripted_efficientnet.pt"
-MODEL_PATH = f"/repository/{MODEL_FILENAME}"
+# Load the model
+def model_fn(model_dir):
+    # Load the scripted PyTorch model
+    model = torch.jit.load(f"{model_dir}/model_scripted_efficientnet.pt")
+    model.eval()
+    return model
 
-def download_model():
-    """Download model file from Hugging Face if it does not exist."""
-    if not os.path.exists(MODEL_PATH):
-        print("Model file not found! Downloading from Hugging Face...")
-        model_url = f"https://huggingface.co/{HF_REPO}/resolve/main/{MODEL_FILENAME}"
-        response = requests.get(model_url, stream=True)
+# Preprocess the input image
+def input_fn(request_body, request_content_type):
+    if request_content_type == "image/jpeg" or request_content_type == "image/png":
+        # Load the image from the request body
+        image = Image.open(io.BytesIO(request_body)).convert("RGB")
+        
+        # Define the image transformation
+        transform = transforms.Compose([
+            transforms.Resize((224, 224)),
+            transforms.ToTensor(),
+            transforms.Normalize(
+                mean=[0.485, 0.456, 0.406],  # Mean for ImageNet
+                std=[0.229, 0.224, 0.225]    # Std for ImageNet
+            )
+        ])
+        
+        # Apply the transformation and add a batch dimension
+        image = transform(image).unsqueeze(0)
+        return image
+    else:
+        raise ValueError(f"Unsupported content type: {request_content_type}")
 
-        if response.status_code == 200:
-            with open(MODEL_PATH, "wb") as f:
-                for chunk in response.iter_content(chunk_size=8192):
-                    f.write(chunk)
-            print("Model downloaded successfully!")
-        else:
-            raise RuntimeError(f"Failed to download model, status code: {response.status_code}")
+# Run inference
+def predict_fn(input_data, model):
+    with torch.no_grad():
+        # Get the model predictions
+        outputs = model(input_data)
+    # Convert the outputs to probabilities using softmax (if needed)
+    probabilities = torch.nn.functional.softmax(outputs, dim=1)
+    return probabilities.squeeze().tolist()
 
-# Ensure model is downloaded before loading
-download_model()
-
-# Load Model
-device = "cuda" if torch.cuda.is_available() else "cpu"
-model = torch.jit.load(MODEL_PATH, map_location=device)
-model.eval()
-
-# Define Transform
-transform = transforms.Compose([
-    transforms.Resize((224, 224)),
-    transforms.ToTensor(),
-    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
-])
-
-# Supported labels
-supported_issues = [
-    "Dark Spots",
-    "Dry Lips",
-    "Forehead Wrinkles",
-    "Jowls",
-    "Nasolabial Folds",
-    "Prejowl Sulcus",
-    "Thin Lips",
-    "Under Eye Hollow",
-    "Under Eye Wrinkles",
-    "Brow Asymmetry"
-]
-
-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.")
-
-    def __call__(self, data):
-        """Perform inference on an image."""
-        if "image" not in data:
-            return {"error": "No image provided"}
-
-        image_data = data["image"]
-        image = Image.open(image_data).convert("RGB")
-        image = transform(image).unsqueeze(0).to(self.device)
-
-        with torch.no_grad():
-            outputs = self.model(image)
-
-        predictions = outputs.squeeze().tolist()
-        output_labels = [label for label, prob in zip(supported_issues, predictions) if prob > 0.5]
-
-        return {"predictions": output_labels}
+# Postprocess the output
+def output_fn(predictions, content_type):
+    # Define the supported issues and their corresponding labels
+    supported_issues = [
+        "Dark Spots",
+        "Dry Lips",
+        "Forehead Wrinkles",
+        "Jowls",
+        "Nasolabial Folds",
+        "Prejowl Sulcus",
+        "Thin Lips",
+        "Under Eye Hollow",
+        "Under Eye Wrinkles",
+        "Brow Asymmetry"
+    ]
+    
+    # Filter issues with probability > 0.5
+    output = [issue for issue, prob in zip(supported_issues, predictions) if prob > 0.5]
+    
+    # Return the output as a JSON response
+    return {"predictions": output}