cog.yaml

build:
  gpu: true
  python_version: "3.12"
  system_packages:
    - "libglib2.0-0"
    - "ffmpeg"
  python_packages:
    - "ultralytics>=8.3.0"

predict: predict.py:Predictor
image: r8.im/alexgodin/yolov11-fasionpedia

predict.py

import os
from typing import Optional

from cog import BaseModel, BasePredictor, Input, Path
from ultralytics import YOLO


class Output(BaseModel):
    """Output model for predictions."""

    image: Optional[Path] = None
    detections: Optional[list] = None


class Predictor(BasePredictor):
    """Fashion YOLO11 model predictor for Replicate deployment."""

    def setup(self) -> None:
        """Load YOLO model into memory."""
        self.model = YOLO("best.pt")

    def format_detections(self, result):
        """Format YOLO detections to match requested output format."""
        detections = []
        if result.boxes is not None:
            for box in result.boxes:
                x1, y1, x2, y2 = box.xyxy[0].cpu().numpy()
                conf_val = box.conf[0].cpu().numpy().item()
                class_id = int(box.cls[0].cpu().numpy())
                label = result.names[class_id]

                detections.append({
                    "bbox": [int(x1), int(y1), int(x2), int(y2)],
                    "label": label,
                    "confidence": float(conf_val)
                })
        return detections

    def predict(
        self,
        image: Path = Input(description="Input image"),
        conf: float = Input(description="Confidence threshold", default=0.25, ge=0.0, le=1.0),
        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_image: bool = Input(description="Return annotated image", default=False),
    ) -> Output:
        """Run inference and return JSON results with optional annotated image."""
        print(f"[DEBUG] raw image input: {image} ({type(image)})")
        if hasattr(image, "path"):
            print(f"[DEBUG] image.path: {getattr(image, 'path', None)}")
            if getattr(image, "path", None):
                print(f"[DEBUG] image.path is file: {os.path.isfile(image.path)}, is dir: {os.path.isdir(image.path)}")
                if os.path.isdir(image.path):
                    try:
                        print(f"[DEBUG] image.path contents: {os.listdir(image.path)}")
                    except Exception as exc:
                        print(f"[DEBUG] failed to list image.path contents: {exc}")

        str_image = str(image)
        print(f"[DEBUG] str(image): {str_image}")

        source_path = image.path if hasattr(image, "path") and image.path else str_image
        print(f"[DEBUG] chosen source_path: {source_path}, exists: {os.path.exists(source_path)}")

        # YOLO's data loader validates files by extension before loading them.
        # When Cog downloads URLs, it saves without an extension, causing YOLO to reject the file.
        # We add .jpg extension to pass validation - PIL/OpenCV will detect the actual format
        # from file content, so the extension doesn't need to match the real format.
        if os.path.exists(source_path) and not os.path.splitext(source_path)[1]:
            new_path = source_path + ".jpg"
            os.rename(source_path, new_path)
            source_path = new_path
            print(f"[DEBUG] renamed to: {source_path}")

        result = self.model(source_path, conf=conf, iou=iou, imgsz=imgsz)[0]
        detections = self.format_detections(result)

        if return_image:
            image_path = "output.png"
            result.save(image_path)
            return Output(image=Path(image_path), detections=detections)
        else:
            return Output(detections=detections)