Update handler.py

handler.py

@@ -1,48 +1,79 @@
-from typing import  Dict, List, Any
-from PIL import Image
 import torch
-import os
+from transformers import AutoFeatureExtractor, AutoModelForImageClassification
+from PIL import Image
+import json
+import base64
 from io import BytesIO
-from transformers import BlipForConditionalGeneration, BlipProcessor
-# -
-
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
-class EndpointHandler():
-    def __init__(self, path=""):
-        # load the optimized model
+# Load the model and feature extractor when the handler is initialized
+class VisionModelHandler:
+    def __init__(self, model_name_or_path="https://huggingface.co/VishalD1234/Florence-metere1"):
+        self.model_name_or_path = model_name_or_path
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         
-        self.processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base") 
-        self.model = BlipForConditionalGeneration.from_pretrained(
-            "Salesforce/blip-image-captioning-base"
-        ).to(device)
-        self.model.eval()
-        self.model = self.model.to(device)
+        # Load model and feature extractor
+        self.model = AutoModelForImageClassification.from_pretrained(self.model_name_or_path)
+        self.feature_extractor = AutoFeatureExtractor.from_pretrained(self.model_name_or_path)
         
-
-
-    def __call__(self, data: Any) -> Dict[str, Any]:
+        # Move the model to the appropriate device (GPU/CPU)
+        self.model.to(self.device)
+        self.model.eval()  # Set model to evaluation mode
+    
+    def preprocess_image(self, image_data):
+        """
+        Preprocess the image for the model. Convert it from base64 and apply the necessary transformations.
+        """
+        image = Image.open(BytesIO(base64.b64decode(image_data)))
+        inputs = self.feature_extractor(images=image, return_tensors="pt").to(self.device)
+        return inputs
+    
+    def predict(self, inputs):
+        """
+        Perform inference and return the model's predictions.
+        """
+        with torch.no_grad():
+            outputs = self.model(**inputs)
+            logits = outputs.logits
+            predicted_class_idx = logits.argmax(-1).item()  # Get the index of the highest score
+        return predicted_class_idx
+    
+    def handle(self, event, context):
         """
+        Entry point for the inference request. This will be called by the inference endpoint.
+        
         Args:
-            data (:obj:):
-                includes the input data and the parameters for the inference.
-        Return:
-            A :obj:`dict`:. The object returned should be a dict of one list like {"captions": ["A hugging face at the office"]} containing :
-                - "caption": A string corresponding to the generated caption.
+            event: This will contain the input data, usually in the form of a JSON with an image in base64.
+            context: Optional, can contain metadata about the request (not used here).
+        
+        Returns:
+            A JSON response with the prediction result.
         """
-        inputs = data.pop("inputs", data)
-        parameters = data.pop("parameters", {})
- 
-        raw_images = [Image.open(BytesIO(_img)) for _img in inputs]
-                                     
-        processed_image = self.processor(images=raw_images, return_tensors="pt") 
-        processed_image["pixel_values"] = processed_image["pixel_values"].to(device)
-        processed_image = {**processed_image, **parameters}
+        # Extract image data from the request body
+        body = json.loads(event.get("body", "{}"))
+        image_data = body.get("image_base64", None)
         
-        with torch.no_grad():
-            out = self.model.generate(
-                **processed_image
-            )
-        captions = self.processor.batch_decode(out, skip_special_tokens=True)
-        # postprocess the prediction
-        return {"captions": captions}
+        if image_data is None:
+            return {
+                "statusCode": 400,
+                "body": json.dumps({"error": "No image data found in the request"})
+            }
+        
+        # Preprocess the image and make predictions
+        inputs = self.preprocess_image(image_data)
+        prediction = self.predict(inputs)
+        
+        # You can add more details to this mapping depending on your use case
+        response = {
+            "statusCode": 200,
+            "body": json.dumps({"prediction": prediction})
+        }
+        
+        return response
+
+
+# Instantiate the handler
+vision_model_handler = VisionModelHandler()
+
+# If running in an environment like AWS Lambda or Sagemaker, ensure this is exposed
+def lambda_handler(event, context):
+    return vision_model_handler.handle(event, context)