Update handler.py

handler.py

@@ -1,28 +1,20 @@
-import torch
-from transformers import AutoModelForObjectDetection, AutoProcessor
+from ultralytics import YOLO
+from huggingface_inference_toolkit.handler import BaseHandler
 
-class CustomHandler:
-    def __init__(self):
-        self.model = None
-        self.processor = None
+class CustomYOLOHandler(BaseHandler):
+    def __init__(self, model_dir, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.model = YOLO(f"{model_dir}/model.pt")
 
-    def initialize(self, model_dir):
-        # 載入模型和處理器
-        self.model = AutoModelForObjectDetection.from_pretrained(model_dir)
-        self.processor = AutoProcessor.from_pretrained(model_dir)
-
-    def preprocess(self, request):
-        # 解析輸入資料
-        inputs = request.get("inputs")
-        return self.processor(images=inputs, return_tensors="pt")
+    def preprocess(self, inputs):
+        # Preprocess inputs for YOLO
+        return inputs
 
     def inference(self, inputs):
-        # 執行推理
-        with torch.no_grad():
-            outputs = self.model(**inputs)
-        return outputs
+        # Perform inference
+        results = self.model(inputs)
+        return results
 
     def postprocess(self, outputs):
-        # 從模型輸出轉換成人類可讀格式
-        results = outputs.logits.softmax(-1).tolist()
-        return {"predictions": results}
+        # Convert YOLO results to HuggingFace pipeline outputs
+        return outputs