Update handler.py

handler.py

@@ -1,46 +1,32 @@
 from transformers import pipeline
 import torch
 from PIL import Image
-from io import BytesIO
 import base64
-from typing import Dict, List, Any
+from io import BytesIO
 
-class EndpointHandler():
+class EndpointHandler:
     def __init__(self, model_path=""):
-        # Check if a GPU is available
-        if torch.cuda.is_available():
-            device = 0
-            gpu_info = torch.cuda.get_device_name(device)
-            print(f"Using GPU: {gpu_info}")
-        else:
-            device = -1
-            cpu_info = torch.get_num_threads()
-            print(f"Using CPU with {cpu_info} threads")
-            
-        # Initialize the pipeline with the specified model and set the device to GPU
-        self.pipeline = pipeline(task="zero-shot-object-detection", model=model_path, device=device)
+        # Dynamically assign computing device based on availability.
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        print(f"Using {'GPU: ' + torch.cuda.get_device_name(0) if self.device == 'cuda' else 'CPU'}")
+
+        # Initialize model with the capability to automatically adjust to GPU or CPU.
+        self.pipeline = pipeline("zero-shot-object-detection", model=model_path, device=0 if self.device == 'cuda' else -1)
 
-    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
+    def __call__(self, data):
         """
-        Process an incoming request for zero-shot object detection.
-        
+        Decode image, run zero-shot object detection, and return results.
+
         Args:
-            data (Dict[str, Any]): The input data containing an encoded image and candidate labels.
-            
+            data (dict): Contains base64-encoded image and candidate labels.
+
         Returns:
-            A list of dictionaries, each containing a label and its corresponding score.
+            list[dict]: Each dict contains a label and its score from object detection.
         """
-        # Correctly accessing the 'inputs' key and fixing the typo in 'candidates'
-        inputs = data.get("inputs", {})
-        
-        # Decode the base64 image to a PIL image
-        image = Image.open(BytesIO(base64.b64decode(inputs['image'])))
+        # Decode the base64 image to PIL format.
+        image = Image.open(BytesIO(base64.b64decode(data['inputs']['image'])))
 
-        # Get candidate labels
-        candidate_labels=inputs["candidates"]
+        # Run detection and obtain results.
+        results = self.pipeline(image=image, candidate_labels=data['inputs']['candidates'])
 
-        # Correctly passing the image and candidate labels to the pipeline
-        detection_results = self.pipeline(image=image, candidate_labels=inputs["candidates"], threshold = 0)
-        
-        # Adjusting the return statement to match the expected output structure
-        return detection_results
+        return results