Blame for main.py

from __future__ import annotations

import importlib.util
import json
import os
import sys
from pathlib import Path
from typing import Any

import cv2
import numpy as np


def _load_local_miner_class():
    miner_path = Path(__file__).resolve().parent / "miner.py"
    spec = importlib.util.spec_from_file_location("manako_bridge_local_miner", str(miner_path))
    if spec is None or spec.loader is None:
        raise RuntimeError(f"Could not load miner module from {miner_path}")
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    miner_class = getattr(module, "Miner", None)
    if miner_class is None:
        raise RuntimeError(f"miner.py does not export Miner in {miner_path}")
    return miner_class


Miner = _load_local_miner_class()


CLASS_NAMES = ['car', 'bus', 'truck', 'motorcycle', 'van']
MODEL_TYPE = 'onnxruntime'


def _to_dict(value: Any) -> dict[str, Any]:
    if isinstance(value, dict):
        return value
    if hasattr(value, "model_dump") and callable(value.model_dump):
        dumped = value.model_dump()
        if isinstance(dumped, dict):
            return dumped
    if hasattr(value, "__dict__"):
        return dict(value.__dict__)
    return {}


def _extract_boxes(frame_result: Any) -> list[Any]:
    frame = _to_dict(frame_result)
    boxes = frame.get("boxes", [])
    if isinstance(boxes, list):
        return boxes
    return []


def _extract_polygons(frame_result: Any) -> list[Any]:
    frame = _to_dict(frame_result)
    polygons = frame.get("polygons", [])
    if isinstance(polygons, list):
        return polygons
    polygon = frame.get("polygon", [])
    if isinstance(polygon, list):
        return [polygon]
    return []


def _normalize_polygon_item(polygon: Any) -> list[Any]:
    payload = _to_dict(polygon)
    points = payload.get("polygon")
    if isinstance(points, list) and points:
        return points
    points = payload.get("points")
    if isinstance(points, list) and points:
        return points
    if isinstance(polygon, list):
        return polygon
    return []


def _bbox_from_polygon(polygon: Any) -> tuple[float, float, float, float] | None:
    if not isinstance(polygon, list) or not polygon:
        return None
    xs: list[float] = []
    ys: list[float] = []
    for point in polygon:
        if not isinstance(point, (list, tuple)) or len(point) < 2:
            continue
        try:
            xs.append(float(point[0]))
            ys.append(float(point[1]))
        except (TypeError, ValueError):
            continue
    if not xs or not ys:
        return None
    return min(xs), min(ys), max(xs), max(ys)


def _resolve_runtime_class_names(miner: Any) -> list[str]:
    value = getattr(miner, "class_names", None)
    if isinstance(value, (list, tuple)):
        resolved = [str(item) for item in value]
        if resolved:
            return resolved
    return list(CLASS_NAMES)


def _to_detection(box: Any, class_names: list[str]) -> dict[str, Any]:
    payload = _to_dict(box)
    cls_id = int(payload.get("cls_id", 0))
    polygon = payload.get("polygon")
    if not (isinstance(polygon, list) and polygon):
        polygons = payload.get("polygons")
        if isinstance(polygons, list) and polygons:
            polygon = polygons[0]
    if not (isinstance(polygon, list) and polygon):
        points = payload.get("points")
        if isinstance(points, list) and points:
            polygon = points

    bbox = _bbox_from_polygon(polygon) if polygon is not None else None
    x1 = float(payload.get("x1", bbox[0] if bbox else 0.0))
    y1 = float(payload.get("y1", bbox[1] if bbox else 0.0))
    x2 = float(payload.get("x2", bbox[2] if bbox else 0.0))
    y2 = float(payload.get("y2", bbox[3] if bbox else 0.0))
    width = max(0.0, x2 - x1)
    height = max(0.0, y2 - y1)
    detection = {
        "bbox": [x1, y1, x2, y2],
        "x": x1 + width / 2.0,
        "y": y1 + height / 2.0,
        "width": width,
        "height": height,
        "confidence": float(payload.get("conf", 0.0)),
        "class_id": cls_id,
        "class": class_names[cls_id] if 0 <= cls_id < len(class_names) else str(cls_id),
    }
    if isinstance(polygon, list) and polygon:
        detection["polygon"] = polygon
    polygons = payload.get("polygons")
    if isinstance(polygons, list) and polygons:
        detection["polygons"] = polygons
        if "polygon" not in detection:
            detection["polygon"] = polygons[0]
    if "polygon" not in detection:
        points = payload.get("points")
        if isinstance(points, list) and points:
            detection["polygon"] = points
    return detection


def _require_bgr_frame(image: Any) -> Any:
    # Same frame the turbovision validator hands predict_batch: cv2.imdecode(...,
    # IMREAD_COLOR) -> BGR uint8 HWC with 3 channels. Anything else is a harness
    # bug and must fail loudly rather than be coerced.
    if not isinstance(image, np.ndarray):
        raise TypeError(f"frame must be a numpy.ndarray (BGR uint8 HWC), got {type(image).__name__}")
    if image.dtype != np.uint8:
        raise TypeError(f"frame dtype must be uint8, got {image.dtype}")
    if image.ndim != 3 or image.shape[2] != 3:
        raise TypeError(f"frame must be HWC with 3 channels (BGR), got shape {image.shape}")
    return image


def load_model(onnx_path: str | None = None, data_dir: str | None = None):
    del onnx_path
    repo_dir = Path(data_dir) if data_dir else Path(__file__).resolve().parent
    miner = Miner(repo_dir)
    class_names = _resolve_runtime_class_names(miner)
    return {
        "miner": miner,
        "model_type": MODEL_TYPE,
        "class_names": class_names,
    }


# Validator chute: predict_batch(batch_images=..., offset=0, n_keypoints=32)
# (turbovision chute_template :376-380; scorevision/utils/challenges.py:412).
VALIDATOR_N_KEYPOINTS = 32


def _predict_frame(miner: Any, image: Any) -> list[Any]:
    frame = _require_bgr_frame(image)
    return miner.predict_batch(batch_images=[frame], offset=0, n_keypoints=VALIDATOR_N_KEYPOINTS)


def run_model(model: Any, image: Any = None, onnx_path: str | None = None, data_dir: str | None = None):
    del onnx_path
    if image is None:
        image = model
        model = load_model(data_dir=data_dir)
    miner = model["miner"]
    class_names = model.get("class_names")
    if not isinstance(class_names, list):
        class_names = list(CLASS_NAMES)
    results = _predict_frame(miner, image)
    if not results:
        return [[]]
    frame_boxes = _extract_boxes(results[0])
    frame_polygons = _extract_polygons(results[0])
    detections = []
    for idx, box in enumerate(frame_boxes):
        detection = _to_detection(box, class_names)
        box_polygon = _normalize_polygon_item(box)
        if box_polygon:
            detection["polygons"] = [box_polygon]
        if idx < len(frame_polygons):
            polygon = _normalize_polygon_item(frame_polygons[idx])
            if polygon:
                detection["polygon"] = polygon
                detection["polygons"] = [polygon]
        detections.append(detection)
    if not detections and frame_polygons:
        for idx, polygon in enumerate(frame_polygons):
            normalized_polygon = _normalize_polygon_item(polygon)
            if not normalized_polygon:
                continue
            detection = _to_detection({"polygon": normalized_polygon}, class_names)
            detection["polygons"] = [normalized_polygon]
            detections.append(detection)
    return [detections]


def main() -> None:
    if len(sys.argv) < 2:
        print("Usage: main.py <image_path>", file=sys.stderr)
        raise SystemExit(1)
    image_path = sys.argv[1]
    image = cv2.imread(image_path, cv2.IMREAD_COLOR)
    if image is None:
        print(f"Could not read image: {image_path}", file=sys.stderr)
        raise SystemExit(1)
    data_dir = os.path.dirname(os.path.abspath(__file__))
    model = load_model(data_dir=data_dir)
    output = run_model(model, image)
    print(json.dumps(output, indent=2))


if __name__ == "__main__":
    main()