Update handler.py

handler.py

@@ -1,20 +1,56 @@
 from ultralytics import YOLO
-from huggingface_inference_toolkit.handler import BaseHandler
+import json
 
-class CustomYOLOHandler(BaseHandler):
-    def __init__(self, model_dir, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        self.model = YOLO(f"{model_dir}/model.pt")
+class YOLOHandler:
+    def __init__(self, model_path):
+        """
+        Initialize the YOLO handler with the given model path.
+        """
+        self.model = YOLO(model_path)
 
     def preprocess(self, inputs):
-        # Preprocess inputs for YOLO
-        return inputs
+        """
+        Preprocess the input data for YOLO inference.
+        :param inputs: JSON input data
+        :return: Preprocessed data
+        """
+        # Assuming inputs is a dictionary with a key 'image' containing the image path or base64-encoded image
+        image_path = inputs.get("image")
+        if not image_path:
+            raise ValueError("Input must contain an 'image' key with a valid path.")
+        return image_path
 
-    def inference(self, inputs):
-        # Perform inference
-        results = self.model(inputs)
+    def predict(self, inputs):
+        """
+        Run inference on the YOLO model.
+        :param inputs: Preprocessed input data
+        :return: YOLO model results
+        """
+        results = self.model(inputs)  # Run the YOLO model
         return results
 
     def postprocess(self, outputs):
-        # Convert YOLO results to HuggingFace pipeline outputs
-        return outputs
+        """
+        Postprocess YOLO model results for returning JSON response.
+        :param outputs: Raw model output
+        :return: Processed results as a dictionary
+        """
+        detections = []
+        for result in outputs:
+            for box in result.boxes:
+                detections.append({
+                    "class": self.model.names[int(box.cls)],  # Get class name
+                    "confidence": box.conf.tolist(),         # Confidence score
+                    "box": box.xyxy.tolist()                 # Bounding box coordinates
+                })
+        return {"detections": detections}
+
+    def __call__(self, inputs):
+        """
+        Full pipeline: preprocess, predict, postprocess.
+        :param inputs: JSON input data
+        :return: JSON output results
+        """
+        preprocessed_data = self.preprocess(inputs)
+        predictions = self.predict(preprocessed_data)
+        return self.postprocess(predictions)