Add image output and update dependencies list for replicate to work (#19)

Signed-off-by: Onuralp SEZER <thunderbirdtr@gmail.com>

Author: onuralpszr

yolo11n/cog.yaml

@@ -2,14 +2,18 @@
 
 build:
   gpu: false
-  python_version: "3.11"
+  python_version: "3.12"
   system_packages:
     - "libgl1-mesa-glx"
     - "libglib2.0-0"
+    - "ffmpeg"
   python_packages:
-    - ultralytics>=8.3.0
-    - torch>=2.0.0
-    - torchvision
-    - pillow
+    - "ultralytics>=8.3.0"
+    - "torch==2.3.1"
+    - "torchvision"
+    - "numpy>=1.24.0"
+    - "opencv-python"
+    - "pillow"
+
 predict: predict.py:Predictor
 image: r8.im/ultralytics/yolo11n

yolo11n/predict.py

@@ -1,10 +1,14 @@
 # Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 
 import json
-from typing import Any, Dict
-
-from cog import BasePredictor, Input, Path
+from typing import Optional
 from ultralytics import YOLO
+from cog import BasePredictor, Input, Path, BaseModel
+
+
+class Output(BaseModel):
+    image: Optional[Path] = None
+    json_str: Optional[str] = None
 
 
 class Predictor(BasePredictor):
@@ -21,13 +25,16 @@ def predict(
         iou: float = Input(description="IoU threshold for NMS", default=0.45, ge=0.0, le=1.0),
         imgsz: int = Input(description="Image size", default=640, choices=[320, 416, 512, 640, 832, 1024, 1280]),
         return_json: bool = Input(description="Return detection results as JSON", default=False),
-    ) -> Dict[str, Any] | Path:
+    ) -> Output:
         """Run inference and return annotated image with optional JSON results."""
         result = self.model(str(image), conf=conf, iou=iou, imgsz=imgsz)[0]
         image_path = "output.png"
         result.save(image_path)
 
         if return_json:
-            return {"image": Path(image_path), "results": json.loads(result.to_json())}
+            return Output(
+                image=Path(image_path), 
+                json_str=result.to_json()
+            )
         else:
-            return Path(image_path)
+            return Output(image=Path(image_path))