Self-Improving CV Pipeline

feature_extractor_name = "yolo_detector"
version = "1.0.0"
description = "YOLOv8 object detection with bounding boxes and class embeddings"

dependencies = ["ultralytics==8.2.0", "torch>=2.0"]

features = [
    {
        "feature_type": "json",
        "feature_name": "detections",
    },
    {
        "feature_type": "embedding",
        "feature_name": "detection_embedding",
        "embedding_dim": 512,
        "distance_metric": "cosine",
    },
]

output_schema = {
    "detections": {
        "type": "array",
        "items": {
            "type": "object",
            "properties": {
                "class": {"type": "string"},
                "confidence": {"type": "number"},
                "bbox": {
                    "type": "object",
                    "properties": {
                        "x": {"type": "number"},
                        "y": {"type": "number"},
                        "w": {"type": "number"},
                        "h": {"type": "number"},
                    },
                },
            },
        },
    },
    "detection_embedding": {
        "type": "array",
        "items": {"type": "number"},
        "description": "512-dim CLIP embedding of the highest-confidence crop",
    },
}

input_mappings = {"image": "image"}
tier = 1
tier_label = "OBJECT_DETECTION"
compute_profile = {"resource_type": "gpu"}

import numpy as np
import pandas as pd
from engine.models.lazy import LazyModelMixin
from engine.inference.services import BaseBatchInferenceService
from engine.io import parallel_io


class YOLODetector(LazyModelMixin, BaseBatchInferenceService):
    model_id = "ultralytics/yolov8m"
    model_source = "huggingface"

    def _instantiate_model(self, cached_data):
        from ultralytics import YOLO
        model = YOLO("yolov8m.pt")
        model.to(self._detect_device())
        return model, None

    def _process_batch(self, batch):
        model, _ = self.get_model()

        images = parallel_io(batch["data"].tolist())

        results = model(images, conf=0.25)

        all_detections = []
        all_embeddings = []

        for result in results:
            detections = []
            for box in result.boxes:
                detections.append({
                    "class": result.names[int(box.cls)],
                    "confidence": float(box.conf),
                    "bbox": {
                        "x": float(box.xywh[0][0]),
                        "y": float(box.xywh[0][1]),
                        "w": float(box.xywh[0][2]),
                        "h": float(box.xywh[0][3]),
                    },
                })
            all_detections.append(detections)

            if detections:
                best = max(detections, key=lambda d: d["confidence"])
                crop = result.orig_img[
                    int(best["bbox"]["y"] - best["bbox"]["h"]/2):int(best["bbox"]["y"] + best["bbox"]["h"]/2),
                    int(best["bbox"]["x"] - best["bbox"]["w"]/2):int(best["bbox"]["x"] + best["bbox"]["w"]/2),
                ]
                embedding = self._embed_crop(crop)
            else:
                embedding = np.zeros(512).tolist()
            all_embeddings.append(embedding)

        batch["detections"] = all_detections
        batch["detection_embedding"] = all_embeddings
        return batch

    def _embed_crop(self, crop):
        # Replace with CLIP or similar for production
        return np.random.randn(512).astype(np.float32).tolist()


def build_steps(extractor_request=None, base_steps=None, **kwargs):
    steps = list(base_steps or [])
    steps.append(YOLODetector())
    return {"steps": steps, "prepare": lambda ds: ds}


def extract(extractor_request=None, base_steps=None, **kwargs):
    result = build_steps(
        extractor_request=extractor_request,
        base_steps=base_steps, **kwargs
    )
    class PipelineResult:
        def __init__(self, steps, prepare):
            self.steps = steps
            self.prepare = prepare
    return PipelineResult(result["steps"], result["prepare"])