| 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: |
| |
| |
| |
| 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_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() |
|
|
