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 = ['fire', 'smoke', 'fire extinguisher']
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 _normalize_image_for_miner(image: Any) -> Any:
    if image is None or hasattr(image, "shape"):
        return image
    if isinstance(image, (bytes, bytearray, memoryview)):
        try:
            buffer = np.frombuffer(bytes(image), dtype=np.uint8)
            decoded = cv2.imdecode(buffer, cv2.IMREAD_COLOR)
            if decoded is not None:
                return decoded
        except Exception:
            return image
    if hasattr(image, "convert") and callable(image.convert):
        try:
            rgb = image.convert("RGB")
            array = np.array(rgb)
            if getattr(array, "ndim", 0) == 3 and array.shape[-1] == 3:
                return cv2.cvtColor(array, cv2.COLOR_RGB2BGR)
            return array
        except Exception:
            return image
    try:
        array = np.asarray(image)
        if getattr(array, "shape", None):
            return array
    except Exception:
        return image
    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,
    }


def _candidate_keypoint_counts(miner: Any) -> list[int]:
    counts: list[int] = [0]
    for attr in ("n_keypoints", "num_keypoints", "keypoint_count", "num_joints"):
        value = getattr(miner, attr, None)
        if isinstance(value, int) and value > 0:
            counts.append(value)
    counts.append(32)

    seen: set[int] = set()
    ordered: list[int] = []
    for count in counts:
        if count in seen:
            continue
        seen.add(count)
        ordered.append(count)
    return ordered


def _predict_batch_with_fallbacks(miner: Any, image: Any) -> list[Any]:
    normalized_image = _normalize_image_for_miner(image)
    errors: list[str] = []
    for n_keypoints in _candidate_keypoint_counts(miner):
        try:
            return miner.predict_batch([normalized_image], offset=0, n_keypoints=n_keypoints)
        except Exception as exc:
            errors.append(f"n_keypoints={n_keypoints} -> {exc}")
            continue
    raise RuntimeError("predict_batch failed for all keypoint candidates: " + " | ".join(errors))


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_batch_with_fallbacks(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()