backend/scripts/seed.py

from pathlib import Path
import sys

from qdrant_client import models

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))

from app.config import get_settings  # noqa: E402
from app.embedder import ImageEmbedder  # noqa: E402
from app.manifest import Manifest  # noqa: E402
from app.qdrant_store import QdrantStore, point_id  # noqa: E402


def payload_base(manifest: Manifest, scenario: dict, item: dict, crop: str, memory_type: str) -> dict:
    camera = next(cam for cam in manifest.cameras(scenario["id"]) if cam["id"] == item["camera_id"])
    return {
        "scenario": scenario["id"],
        "camera_id": item["camera_id"],
        "zone": item.get("zone", camera["zone"]),
        "timestamp": int(item.get("timestamp", 1716900000)),
        "frame_url": manifest.asset_url(item.get("frame", camera.get("incident_frame", camera["baseline_frame"]))),
        "crop_url": manifest.asset_url(crop),
        "bbox": item.get("bbox", [0, 0, 0, 0]),
        "asset_id": item["asset_id"],
        "memory_type": memory_type,
        "is_baseline": bool(item.get("is_baseline", False)),
        "is_incident": bool(item.get("is_incident", False)),
        "object_label": item.get("object_label", "region"),
        "track_id": item.get("track_id"),
        "region_id": item.get("region_id"),
        "region_label": item.get("region_label"),
        "alarm_enabled": item.get("alarm_enabled"),
    }


def region_points(manifest: Manifest, embedder: ImageEmbedder, scenario: dict):
    points = []
    for region in manifest.regions(scenario["id"]):
        camera = next(cam for cam in manifest.cameras(scenario["id"]) if cam["id"] == region["camera_id"])
        for state, crop, baseline in [
            ("baseline", region["baseline_crop"], True),
            ("incident", region["incident_crop"], False),
        ]:
            asset_id = f"{region['region_id']}_{state}"
            payload = payload_base(
                manifest,
                scenario,
                {
                    **region,
                    "asset_id": asset_id,
                    "frame": camera[f"{state}_frame"],
                    "is_baseline": baseline,
                    "is_incident": not baseline,
                    "object_label": region.get(f"{state}_label", region["region_id"]),
                },
                crop,
                f"region_{state}",
            )
            vector = embedder.embed_path(manifest.asset_path(crop))
            points.append(models.PointStruct(id=point_id(f"{scenario['id']}:{asset_id}"), vector=vector, payload=payload))
            if baseline:
                for index, variant_crop in enumerate(baseline_variants(manifest, crop, region.get("variant_prefix")), start=1):
                    variant_id = f"{asset_id}_v{index}"
                    variant_payload = {
                        **payload,
                        "asset_id": variant_id,
                        "crop_url": manifest.asset_url(variant_crop),
                        "memory_variant": index,
                    }
                    vector = embedder.embed_path(manifest.asset_path(variant_crop))
                    points.append(models.PointStruct(id=point_id(f"{scenario['id']}:{variant_id}"), vector=vector, payload=variant_payload))
    return points


def baseline_variants(manifest: Manifest, crop: str, variant_prefix: str | None = None) -> list[str]:
    path = Path(crop)
    stem = variant_prefix or path.stem
    variant_dir = manifest.settings.resolved_asset_root / path.parent / "variants"
    if not variant_dir.exists():
        return []
    return [
        f"{path.parent.as_posix()}/variants/{variant.name}"
        for variant in sorted(variant_dir.glob(f"{stem}_v*.jpg"))
    ]


def object_points(manifest: Manifest, embedder: ImageEmbedder, scenario: dict):
    points = []
    for item in manifest.objects(scenario["id"]):
        payload = payload_base(manifest, scenario, item, item["crop"], item["memory_type"])
        vector = embedder.embed_path(manifest.asset_path(item["crop"]))
        points.append(models.PointStruct(id=point_id(f"{scenario['id']}:{item['asset_id']}"), vector=vector, payload=payload))
    return points


def main() -> int:
    settings = get_settings()
    manifest = Manifest(settings)
    embedder = ImageEmbedder(settings)
    store = QdrantStore(settings)
    store.ensure_collection(recreate=True)
    points = []
    for scenario in manifest.scenarios:
        points.extend(region_points(manifest, embedder, scenario))
        points.extend(object_points(manifest, embedder, scenario))
    store.upsert(points)
    print(f"Seeded {len(points)} points into {settings.collection_name} using {embedder.mode}.")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())