1574 lines
59 KiB
Python
Executable File
1574 lines
59 KiB
Python
Executable File
from __future__ import annotations
|
|
|
|
import json
|
|
import math
|
|
import sys
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Any, Iterable, Optional
|
|
|
|
import cv2
|
|
import numpy as np
|
|
import torch
|
|
import yaml
|
|
|
|
FILE = Path(__file__).resolve()
|
|
ROOT = FILE.parents[2]
|
|
if str(ROOT) not in sys.path:
|
|
sys.path.append(str(ROOT))
|
|
|
|
from tools.model_inference.adapters.video_dir_inference_utils import (
|
|
iter_video_case_frames,
|
|
read_video_frame_index,
|
|
resolve_video_case_paths,
|
|
)
|
|
from ultralytics.data.ground3d_augment import (
|
|
adjust_calib_for_roi_crop,
|
|
build_final_resized_calib,
|
|
compute_centered_roi_bounds,
|
|
)
|
|
from ultralytics.nn.tasks import load_checkpoint
|
|
from ultralytics.utils.plotting import colors
|
|
from ultralytics.utils.plotting_3d import (
|
|
EDGE_YAW_MAX_LATERAL_DIST_M,
|
|
decode_edge_yaw_selection_from_prediction,
|
|
decode_3d_prediction,
|
|
draw_3d_box,
|
|
extract_face_regressed_size_priors_from_prediction,
|
|
extract_3d_attrs_from_prediction,
|
|
project_face_bottom_edge,
|
|
project_partial_face_bottom_edge,
|
|
reconstruct_edge_based_box_from_selection,
|
|
)
|
|
from ultralytics.utils.torch_utils import select_device
|
|
|
|
|
|
DEFAULT_VISUALIZATION_ROOT = FILE.parent / "visualization"
|
|
IMAGE_SUFFIXES = (".png", ".jpg", ".jpeg", ".bmp", ".webp")
|
|
DEFAULT_EDGE_YAW_MAX_LATERAL_DIST_M = EDGE_YAW_MAX_LATERAL_DIST_M
|
|
GRID_BG_COLOR = (24, 24, 24)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ROIModelSpec:
|
|
name: str
|
|
model_path: str
|
|
roi_size: tuple[int, int]
|
|
crop_center_mode: str
|
|
virtual_fx: float
|
|
imgsz: Optional[tuple[int, int]] = None
|
|
conf: float = 0.25
|
|
max_det: int = 300
|
|
|
|
|
|
@dataclass
|
|
class LoadedROIModel:
|
|
spec: ROIModelSpec
|
|
model: torch.nn.Module
|
|
names: dict[int, str]
|
|
face_3d_classes: set[int]
|
|
complete_3d_classes: set[int]
|
|
fake_3d_classes: set[int]
|
|
imgsz: tuple[int, int]
|
|
|
|
|
|
@dataclass
|
|
class PreparedROI:
|
|
name: str
|
|
image: np.ndarray
|
|
crop_bounds: tuple[int, int, int, int]
|
|
calib: dict[str, Any]
|
|
vp_x: float
|
|
vp_y: float
|
|
crop_center_x: float
|
|
crop_center_y: float
|
|
|
|
|
|
@dataclass
|
|
class InferenceContext:
|
|
device: torch.device
|
|
use_half: bool
|
|
classes: Optional[set[int]]
|
|
roi_models: list[LoadedROIModel]
|
|
edge_yaw_max_lateral_dist_m: float
|
|
inference_batch_size: int
|
|
|
|
|
|
@dataclass
|
|
class RawROIOutputs:
|
|
detections: np.ndarray
|
|
preds_3d: np.ndarray
|
|
preds_3d_fake: Optional[np.ndarray]
|
|
preds_edge: Optional[np.ndarray]
|
|
anchors: np.ndarray
|
|
strides: np.ndarray
|
|
|
|
|
|
def _coerce_imgsz(imgsz: Any) -> Optional[tuple[int, int]]:
|
|
if imgsz is None:
|
|
return None
|
|
if isinstance(imgsz, int):
|
|
return (int(imgsz), int(imgsz))
|
|
if isinstance(imgsz, str):
|
|
parts = [part.strip() for part in imgsz.split(",") if part.strip()]
|
|
if len(parts) == 1:
|
|
value = int(parts[0])
|
|
return (value, value)
|
|
if len(parts) == 2:
|
|
return (int(parts[0]), int(parts[1]))
|
|
raise ValueError(f"Unable to parse imgsz={imgsz!r}")
|
|
if isinstance(imgsz, (list, tuple)):
|
|
if len(imgsz) == 1:
|
|
value = int(imgsz[0])
|
|
return (value, value)
|
|
if len(imgsz) == 2:
|
|
return (int(imgsz[0]), int(imgsz[1]))
|
|
raise ValueError(f"Unsupported imgsz value: {imgsz!r}")
|
|
|
|
|
|
def _infer_model_imgsz(model: torch.nn.Module, override: Optional[tuple[int, int]]) -> tuple[int, int]:
|
|
if override is not None:
|
|
return override
|
|
model_args = getattr(model, "args", {}) or {}
|
|
imgsz = _coerce_imgsz(model_args.get("imgsz"))
|
|
return imgsz or (768, 352)
|
|
|
|
|
|
def _load_yaml_if_present(path: str) -> dict[str, Any]:
|
|
if not path:
|
|
return {}
|
|
yaml_path = Path(path)
|
|
if not yaml_path.exists():
|
|
return {}
|
|
with yaml_path.open("r", encoding="utf-8") as file:
|
|
return yaml.safe_load(file) or {}
|
|
|
|
|
|
def camera4_payload_to_raw_calib(payload: dict[str, Any]) -> dict[str, Any]:
|
|
"""Normalize a flat or combined `camera4.json` payload to the raw calib format used by ROI inference."""
|
|
if "intrinsics" in payload:
|
|
intrinsics = payload.get("intrinsics", {}).get("camera4.json", {}) or {}
|
|
extrinsics = payload.get("extrinsics", {}).get("camera4.json", {}) or {}
|
|
pitch = 0.0
|
|
rpy = extrinsics.get("rpy", [0.0, 0.0, 0.0])
|
|
if len(rpy) > 1:
|
|
pitch = math.radians(float(rpy[1]))
|
|
return {
|
|
"focal_u": float(intrinsics["focal_u"]),
|
|
"focal_v": float(intrinsics["focal_v"]),
|
|
"cu": float(intrinsics["cu"]),
|
|
"cv": float(intrinsics["cv"]),
|
|
"pitch": pitch,
|
|
"distort_coeffs": list(intrinsics.get("distort_coeffs", [])),
|
|
"image_width": payload.get("image_width", payload.get("img_width")),
|
|
"image_height": payload.get("image_height", payload.get("img_height")),
|
|
"source_format": "combined_calibration",
|
|
"angle_unit": "radians",
|
|
}
|
|
|
|
required = ("focal_u", "focal_v", "cu", "cv")
|
|
missing = [key for key in required if key not in payload]
|
|
if missing:
|
|
raise KeyError(f"camera4 payload missing required keys: {missing}")
|
|
|
|
# Match ultralytics.data.ground3d_augment.read_calib_from_path() flat-file behavior:
|
|
# keep the payload values as-is instead of applying extra angle conversions here.
|
|
calib = dict(payload)
|
|
calib["focal_u"] = float(payload["focal_u"])
|
|
calib["focal_v"] = float(payload["focal_v"])
|
|
calib["cu"] = float(payload["cu"])
|
|
calib["cv"] = float(payload["cv"])
|
|
calib["distort_coeffs"] = list(payload.get("distort_coeffs", []))
|
|
calib["source_format"] = "flat_camera4"
|
|
calib["angle_unit"] = "degrees"
|
|
return calib
|
|
|
|
|
|
def load_camera4_calib(calib_path: str | Path) -> dict[str, Any]:
|
|
"""Load a flat or combined `camera4.json` file into the raw calib format used by training-time ROI code."""
|
|
calib_path = Path(calib_path)
|
|
with calib_path.open("r", encoding="utf-8") as file:
|
|
payload = json.load(file)
|
|
return camera4_payload_to_raw_calib(payload)
|
|
|
|
|
|
def _to_radians(angle: float, unit: str) -> float:
|
|
if unit == "degrees":
|
|
return math.radians(float(angle))
|
|
return float(angle)
|
|
|
|
|
|
def _compute_vanishing_point_xy(raw_calib: dict[str, Any], ori_w: int, ori_h: int) -> tuple[float, float]:
|
|
"""Compute vanishing point using the explicit cx/cy/fx/fy/pitch/yaw formula."""
|
|
if raw_calib is None:
|
|
return ori_w / 2.0, ori_h / 2.0
|
|
|
|
cx = float(raw_calib.get("cu", ori_w / 2.0))
|
|
cy = float(raw_calib.get("cv", ori_h / 2.0))
|
|
fx = float(raw_calib.get("focal_u", ori_w))
|
|
fy = float(raw_calib.get("focal_v", ori_h))
|
|
angle_unit = str(raw_calib.get("angle_unit", "radians"))
|
|
yaw = _to_radians(float(raw_calib.get("yaw", 0.0)), angle_unit)
|
|
pitch = _to_radians(float(raw_calib.get("pitch", 0.0)), angle_unit)
|
|
vp_x = cx + fx * math.tan(yaw)
|
|
vp_y = cy - fy * math.tan(pitch)
|
|
return float(vp_x), float(vp_y)
|
|
|
|
|
|
def _resize_ground3d_image_in_steps(
|
|
image_bgr: np.ndarray,
|
|
target_size: tuple[int, int],
|
|
interpolation: int = cv2.INTER_LINEAR,
|
|
) -> np.ndarray:
|
|
"""Match Ground3D training resize with repeated 0.5x downsampling before the final resize."""
|
|
target_w, target_h = int(target_size[0]), int(target_size[1])
|
|
current_h, current_w = image_bgr.shape[:2]
|
|
|
|
if (current_w, current_h) == (target_w, target_h):
|
|
return image_bgr
|
|
|
|
resized = image_bgr
|
|
if target_w < current_w and target_h < current_h:
|
|
while True:
|
|
next_w = math.ceil(current_w * 0.5)
|
|
next_h = math.ceil(current_h * 0.5)
|
|
if next_w < target_w or next_h < target_h:
|
|
break
|
|
|
|
resized = cv2.resize(resized, (next_w, next_h), interpolation=interpolation)
|
|
current_h, current_w = resized.shape[:2]
|
|
if (current_w, current_h) == (target_w, target_h):
|
|
return resized
|
|
|
|
return cv2.resize(resized, (target_w, target_h), interpolation=interpolation)
|
|
|
|
|
|
def _prepare_roi_image(
|
|
image_bgr: np.ndarray,
|
|
raw_calib: dict[str, Any],
|
|
spec: ROIModelSpec,
|
|
target_size: tuple[int, int],
|
|
) -> PreparedROI:
|
|
source_image = np.ascontiguousarray(image_bgr).copy()
|
|
ori_h, ori_w = source_image.shape[:2]
|
|
roi_w = min(int(spec.roi_size[0]), ori_w)
|
|
roi_h = min(int(spec.roi_size[1]), ori_h)
|
|
vp_x, vp_y = _compute_vanishing_point_xy(raw_calib, ori_w, ori_h)
|
|
crop_center_x = vp_x if spec.crop_center_mode == "vxvy" else ori_w / 2.0
|
|
crop_center_y = vp_y
|
|
crop_bounds = compute_centered_roi_bounds(ori_w, ori_h, roi_w, roi_h, crop_center_x, vp_y)
|
|
crop_x1, crop_y1, crop_x2, crop_y2 = crop_bounds
|
|
|
|
cropped = source_image[crop_y1:crop_y2, crop_x1:crop_x2].copy()
|
|
crop_size = (crop_x2 - crop_x1, crop_y2 - crop_y1)
|
|
if crop_size == target_size:
|
|
resized = cropped
|
|
else:
|
|
resized = _resize_ground3d_image_in_steps(cropped, target_size, interpolation=cv2.INTER_LINEAR)
|
|
|
|
calib_for_resize = adjust_calib_for_roi_crop(raw_calib, ori_w, ori_h, crop_bounds)
|
|
final_calib = build_final_resized_calib(
|
|
calib_for_resize["focal_u"],
|
|
calib_for_resize["focal_v"],
|
|
calib_for_resize["cu"],
|
|
calib_for_resize["cv"],
|
|
calib_for_resize["src_w"],
|
|
calib_for_resize["src_h"],
|
|
target_size[0],
|
|
target_size[1],
|
|
spec.virtual_fx,
|
|
distort_coeffs=calib_for_resize["distort_coeffs"],
|
|
)
|
|
return PreparedROI(
|
|
name=spec.name,
|
|
image=resized,
|
|
crop_bounds=crop_bounds,
|
|
calib=final_calib,
|
|
vp_x=float(vp_x),
|
|
vp_y=float(vp_y),
|
|
crop_center_x=float(crop_center_x),
|
|
crop_center_y=float(crop_center_y),
|
|
)
|
|
|
|
|
|
def prepare_roi_image(
|
|
image_bgr: np.ndarray,
|
|
raw_calib: dict[str, Any],
|
|
spec: ROIModelSpec,
|
|
target_size: tuple[int, int],
|
|
) -> PreparedROI:
|
|
return _prepare_roi_image(image_bgr, raw_calib, spec, target_size)
|
|
|
|
|
|
def prepare_roi_batch(
|
|
images_bgr: list[np.ndarray],
|
|
raw_calibs: list[dict[str, Any]],
|
|
spec: ROIModelSpec,
|
|
target_size: tuple[int, int],
|
|
) -> list[PreparedROI]:
|
|
if len(images_bgr) != len(raw_calibs):
|
|
raise ValueError(
|
|
f"Expected the same number of images and calibrations, got {len(images_bgr)} images and {len(raw_calibs)} calibrations."
|
|
)
|
|
return [prepare_roi_image(image_bgr, raw_calib, spec, target_size) for image_bgr, raw_calib in zip(images_bgr, raw_calibs)]
|
|
|
|
|
|
def _image_to_tensor(image_bgr: np.ndarray, device: torch.device, use_half: bool) -> torch.Tensor:
|
|
image_rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)
|
|
tensor = torch.from_numpy(np.ascontiguousarray(image_rgb.transpose(2, 0, 1))).to(device)
|
|
tensor = tensor.half() if use_half else tensor.float()
|
|
tensor = tensor / 256.0
|
|
return tensor.unsqueeze(0)
|
|
|
|
|
|
def _images_to_batch_tensor(images_bgr: list[np.ndarray], device: torch.device, use_half: bool) -> torch.Tensor:
|
|
if not images_bgr:
|
|
raise ValueError("Expected at least one ROI image to build a batch tensor.")
|
|
image_batch_rgb = [cv2.cvtColor(np.ascontiguousarray(image_bgr), cv2.COLOR_BGR2RGB).transpose(2, 0, 1) for image_bgr in images_bgr]
|
|
tensor = torch.from_numpy(np.ascontiguousarray(np.stack(image_batch_rgb, axis=0))).to(device)
|
|
tensor = tensor.half() if use_half else tensor.float()
|
|
tensor = tensor / 256.0
|
|
return tensor
|
|
|
|
|
|
def _restore_3d_depth_scale(preds_3d: Optional[np.ndarray], depth_scale: float) -> Optional[np.ndarray]:
|
|
if preds_3d is None:
|
|
return None
|
|
preds_3d = preds_3d.copy()
|
|
for channel in (0, 6, 12, 18, 24):
|
|
preds_3d[:, channel] *= depth_scale
|
|
return preds_3d
|
|
|
|
|
|
def _restore_depth_scale(
|
|
preds_3d: np.ndarray,
|
|
preds_3d_fake: Optional[np.ndarray],
|
|
preds_edge: Optional[np.ndarray],
|
|
depth_scale: float,
|
|
) -> tuple[np.ndarray, Optional[np.ndarray], Optional[np.ndarray]]:
|
|
"""Restore canonical ROI-space depths back to metric camera depth for both prediction branches."""
|
|
preds_3d = _restore_3d_depth_scale(preds_3d, depth_scale)
|
|
preds_3d_fake = _restore_3d_depth_scale(preds_3d_fake, depth_scale)
|
|
|
|
if preds_edge is None:
|
|
return preds_3d, preds_3d_fake, None
|
|
|
|
preds_edge = preds_edge.copy()
|
|
preds_edge[:, 2::3] *= depth_scale
|
|
return preds_3d, preds_3d_fake, preds_edge
|
|
|
|
|
|
def _unpack_model_outputs(outputs: Any) -> tuple[np.ndarray, np.ndarray, Optional[np.ndarray], Optional[np.ndarray], np.ndarray, np.ndarray]:
|
|
batch_outputs = _unpack_model_outputs_batch(outputs)
|
|
if not batch_outputs:
|
|
raise RuntimeError("Expected model forward to return at least one inference sample.")
|
|
first = batch_outputs[0]
|
|
return first.detections, first.preds_3d, first.preds_3d_fake, first.preds_edge, first.anchors, first.strides
|
|
|
|
|
|
def _to_numpy_batch(value: Any) -> Optional[np.ndarray]:
|
|
if value is None:
|
|
return None
|
|
if torch.is_tensor(value):
|
|
return value.detach().float().cpu().numpy()
|
|
return np.asarray(value)
|
|
|
|
|
|
def _split_batch_array(values: Optional[np.ndarray], batch_size: int, value_name: str) -> list[Optional[np.ndarray]]:
|
|
if values is None:
|
|
return [None] * batch_size
|
|
if values.ndim == 0:
|
|
raise RuntimeError(f"Expected batched `{value_name}` values, but got a scalar.")
|
|
if values.ndim == 1:
|
|
if batch_size != 1:
|
|
raise RuntimeError(f"Expected batched `{value_name}` values for {batch_size} samples, but got shape {values.shape}.")
|
|
values = values[None, ...]
|
|
if values.shape[0] != batch_size:
|
|
raise RuntimeError(
|
|
f"Batch size mismatch for `{value_name}`: detections batch={batch_size}, {value_name} batch={values.shape[0]}."
|
|
)
|
|
return [values[index] for index in range(batch_size)]
|
|
|
|
|
|
def _unpack_model_outputs_batch(outputs: Any) -> list[RawROIOutputs]:
|
|
if not isinstance(outputs, (list, tuple)) or len(outputs) < 2:
|
|
raise RuntimeError("Expected model forward to return `(detections, raw_preds)` for Detect3D inference.")
|
|
|
|
detections = outputs[0]
|
|
raw_preds = outputs[1]
|
|
if not isinstance(raw_preds, dict):
|
|
raise RuntimeError("Detect3D raw predictions are missing from model output.")
|
|
|
|
one2one = raw_preds.get("one2one", {})
|
|
preds_3d = one2one.get("preds_3d_selected")
|
|
preds_3d_fake = one2one.get("preds_3d_fake_selected")
|
|
anchors = one2one.get("anchors_selected")
|
|
strides = one2one.get("strides_selected")
|
|
preds_edge = one2one.get("preds_edge_selected")
|
|
if preds_3d is None or anchors is None or strides is None:
|
|
raise RuntimeError("Detect3D metadata (`preds_3d_selected` / `anchors_selected` / `strides_selected`) is missing.")
|
|
|
|
if isinstance(detections, (list, tuple)):
|
|
det_batches = [_to_numpy_batch(det) for det in detections]
|
|
else:
|
|
det_np = _to_numpy_batch(detections)
|
|
if det_np is None:
|
|
raise RuntimeError("Detect3D detections are missing from model output.")
|
|
if det_np.ndim == 2:
|
|
det_np = det_np[None, ...]
|
|
det_batches = [det_np[index] for index in range(det_np.shape[0])]
|
|
|
|
batch_size = len(det_batches)
|
|
preds_3d_batches = _split_batch_array(_to_numpy_batch(preds_3d), batch_size, "preds_3d_selected")
|
|
preds_3d_fake_batches = _split_batch_array(_to_numpy_batch(preds_3d_fake), batch_size, "preds_3d_fake_selected")
|
|
preds_edge_batches = _split_batch_array(_to_numpy_batch(preds_edge), batch_size, "preds_edge_selected")
|
|
anchors_batches = _split_batch_array(_to_numpy_batch(anchors), batch_size, "anchors_selected")
|
|
strides_batches = _split_batch_array(_to_numpy_batch(strides), batch_size, "strides_selected")
|
|
|
|
return [
|
|
RawROIOutputs(
|
|
detections=np.asarray(det_batches[index], dtype=np.float32),
|
|
preds_3d=np.asarray(preds_3d_batches[index], dtype=np.float32),
|
|
preds_3d_fake=None
|
|
if preds_3d_fake_batches[index] is None
|
|
else np.asarray(preds_3d_fake_batches[index], dtype=np.float32),
|
|
preds_edge=None if preds_edge_batches[index] is None else np.asarray(preds_edge_batches[index], dtype=np.float32),
|
|
anchors=np.asarray(anchors_batches[index], dtype=np.float32),
|
|
strides=np.asarray(strides_batches[index], dtype=np.float32),
|
|
)
|
|
for index in range(batch_size)
|
|
]
|
|
|
|
|
|
def _filter_prediction_rows(
|
|
detections: np.ndarray,
|
|
preds_3d: np.ndarray,
|
|
preds_3d_fake: Optional[np.ndarray],
|
|
preds_edge: Optional[np.ndarray],
|
|
anchors: np.ndarray,
|
|
strides: np.ndarray,
|
|
conf_thres: float,
|
|
max_det: int,
|
|
classes: Optional[set[int]],
|
|
) -> tuple[np.ndarray, np.ndarray, Optional[np.ndarray], np.ndarray, np.ndarray]:
|
|
keep_idx = _compute_prediction_keep_idx(detections, conf_thres=conf_thres, max_det=max_det, classes=classes)
|
|
if detections.size == 0:
|
|
empty = detections[:0]
|
|
return empty, preds_3d[:0], None if preds_3d_fake is None else preds_3d_fake[:0], None if preds_edge is None else preds_edge[:0], anchors[:, :0], strides[:0]
|
|
return (
|
|
detections[keep_idx],
|
|
preds_3d[keep_idx],
|
|
None if preds_3d_fake is None else preds_3d_fake[keep_idx],
|
|
None if preds_edge is None else preds_edge[keep_idx],
|
|
anchors[:, keep_idx],
|
|
strides[keep_idx],
|
|
)
|
|
|
|
|
|
def _compute_prediction_keep_idx(
|
|
detections: np.ndarray,
|
|
conf_thres: float,
|
|
max_det: int,
|
|
classes: Optional[set[int]],
|
|
) -> np.ndarray:
|
|
if detections.size == 0:
|
|
return np.zeros((0,), dtype=np.int64)
|
|
|
|
keep = detections[:, 4] >= conf_thres
|
|
if classes is not None:
|
|
keep &= np.isin(detections[:, 5].astype(np.int32), np.asarray(sorted(classes), dtype=np.int32))
|
|
return np.flatnonzero(keep)[:max_det]
|
|
|
|
|
|
def iter_batches(items: Iterable[Any], batch_size: int) -> Iterable[list[Any]]:
|
|
resolved_batch_size = max(1, int(batch_size))
|
|
batch: list[Any] = []
|
|
for item in items:
|
|
batch.append(item)
|
|
if len(batch) >= resolved_batch_size:
|
|
yield batch
|
|
batch = []
|
|
if batch:
|
|
yield batch
|
|
|
|
|
|
def run_model_for_prepared_roi_batch(
|
|
bundle: LoadedROIModel,
|
|
prepared_batch: list[PreparedROI],
|
|
device: Optional[torch.device] = None,
|
|
use_half: Optional[bool] = None,
|
|
) -> list[RawROIOutputs]:
|
|
if not prepared_batch:
|
|
return []
|
|
resolved_device = device if device is not None else next(bundle.model.parameters()).device
|
|
resolved_use_half = bool(use_half if use_half is not None else next(bundle.model.parameters()).dtype == torch.float16)
|
|
image_tensor = _images_to_batch_tensor([prepared.image for prepared in prepared_batch], device=resolved_device, use_half=resolved_use_half)
|
|
with torch.inference_mode():
|
|
outputs = bundle.model(image_tensor)
|
|
|
|
raw_outputs_batch = _unpack_model_outputs_batch(outputs)
|
|
if len(raw_outputs_batch) != len(prepared_batch):
|
|
raise RuntimeError(
|
|
f"Model output batch size mismatch for {bundle.spec.name}: "
|
|
f"prepared={len(prepared_batch)} raw_outputs={len(raw_outputs_batch)}."
|
|
)
|
|
|
|
restored_outputs = []
|
|
for prepared, raw_outputs in zip(prepared_batch, raw_outputs_batch):
|
|
preds_3d, preds_3d_fake, preds_edge = _restore_depth_scale(
|
|
raw_outputs.preds_3d,
|
|
raw_outputs.preds_3d_fake,
|
|
raw_outputs.preds_edge,
|
|
float(prepared.calib.get("depth_scale", 1.0)),
|
|
)
|
|
restored_outputs.append(
|
|
RawROIOutputs(
|
|
detections=raw_outputs.detections,
|
|
preds_3d=preds_3d,
|
|
preds_3d_fake=preds_3d_fake,
|
|
preds_edge=preds_edge,
|
|
anchors=raw_outputs.anchors,
|
|
strides=raw_outputs.strides,
|
|
)
|
|
)
|
|
return restored_outputs
|
|
|
|
|
|
def run_model_for_prepared_roi(
|
|
bundle: LoadedROIModel,
|
|
prepared: PreparedROI,
|
|
device: Optional[torch.device] = None,
|
|
use_half: Optional[bool] = None,
|
|
) -> RawROIOutputs:
|
|
outputs_batch = run_model_for_prepared_roi_batch(bundle, [prepared], device=device, use_half=use_half)
|
|
if not outputs_batch:
|
|
raise RuntimeError(f"No inference outputs were produced for {bundle.spec.name}.")
|
|
return outputs_batch[0]
|
|
|
|
|
|
def filter_prediction_outputs(
|
|
raw_outputs: RawROIOutputs,
|
|
conf_thres: float,
|
|
max_det: int,
|
|
classes: Optional[set[int]],
|
|
) -> RawROIOutputs:
|
|
detections, preds_3d, preds_3d_fake, preds_edge, anchors, strides = _filter_prediction_rows(
|
|
raw_outputs.detections,
|
|
raw_outputs.preds_3d,
|
|
raw_outputs.preds_3d_fake,
|
|
raw_outputs.preds_edge,
|
|
raw_outputs.anchors,
|
|
raw_outputs.strides,
|
|
conf_thres=conf_thres,
|
|
max_det=max_det,
|
|
classes=classes,
|
|
)
|
|
return RawROIOutputs(
|
|
detections=detections,
|
|
preds_3d=preds_3d,
|
|
preds_3d_fake=preds_3d_fake,
|
|
preds_edge=preds_edge,
|
|
anchors=anchors,
|
|
strides=strides,
|
|
)
|
|
|
|
|
|
def filter_prediction_outputs_batch(
|
|
raw_outputs_batch: list[RawROIOutputs],
|
|
conf_thres: float,
|
|
max_det: int,
|
|
classes: Optional[set[int]],
|
|
) -> list[RawROIOutputs]:
|
|
return [
|
|
filter_prediction_outputs(
|
|
raw_outputs=raw_outputs,
|
|
conf_thres=conf_thres,
|
|
max_det=max_det,
|
|
classes=classes,
|
|
)
|
|
for raw_outputs in raw_outputs_batch
|
|
]
|
|
|
|
|
|
def _class_name(names: dict[int, str], cls_id: int) -> str:
|
|
if isinstance(names, dict):
|
|
return str(names.get(cls_id, cls_id))
|
|
return str(cls_id)
|
|
|
|
|
|
def _infer_fake_3d_classes(model: torch.nn.Module, names: dict[int, str]) -> set[int]:
|
|
explicit = set(getattr(model, "fake_3d_classes", set()) or set())
|
|
if explicit:
|
|
return explicit
|
|
if not isinstance(names, dict):
|
|
return set()
|
|
return {int(cls_id) for cls_id, name in names.items() if str(name).endswith("_fake")}
|
|
|
|
|
|
def _draw_2d_boxes(
|
|
image: np.ndarray,
|
|
detections: np.ndarray,
|
|
names: dict[int, str],
|
|
) -> np.ndarray:
|
|
drawn = image.copy()
|
|
for det in detections:
|
|
x1, y1, x2, y2 = np.round(det[:4]).astype(np.int32)
|
|
conf = float(det[4])
|
|
cls_id = int(det[5])
|
|
color = colors(cls_id, bgr=True)
|
|
cv2.rectangle(drawn, (x1, y1), (x2, y2), color, 1, cv2.LINE_AA)
|
|
label = f"{_class_name(names, cls_id)} {conf:.2f}"
|
|
(tw, th), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.45, 1)
|
|
cv2.rectangle(drawn, (x1, max(0, y1 - th - 6)), (x1 + tw + 2, y1), color, -1)
|
|
cv2.putText(drawn, label, (x1 + 1, max(th + 1, y1 - 3)), cv2.FONT_HERSHEY_SIMPLEX, 0.45, (255, 255, 255), 1, cv2.LINE_AA)
|
|
return drawn
|
|
|
|
|
|
def _annotate_panel_title(image: np.ndarray, title: str) -> np.ndarray:
|
|
drawn = image.copy()
|
|
cv2.putText(drawn, title, (10, 24), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 255), 2, cv2.LINE_AA)
|
|
return drawn
|
|
|
|
|
|
def _serialize_scalar(value: Any) -> Any:
|
|
if isinstance(value, (np.floating, np.integer)):
|
|
value = value.item()
|
|
if isinstance(value, float) and not math.isfinite(value):
|
|
return None
|
|
return value
|
|
|
|
|
|
def _serialize_array(values: Any) -> Any:
|
|
if values is None:
|
|
return None
|
|
arr = np.asarray(values)
|
|
if arr.ndim == 0:
|
|
return _serialize_scalar(arr.item())
|
|
return [_serialize_array(item) for item in arr.tolist()]
|
|
|
|
|
|
def _wrapped_angle_diff_rad(lhs: Optional[float], rhs: Optional[float]) -> Optional[float]:
|
|
if lhs is None or rhs is None:
|
|
return None
|
|
lhs_val = float(lhs)
|
|
rhs_val = float(rhs)
|
|
if not math.isfinite(lhs_val) or not math.isfinite(rhs_val):
|
|
return None
|
|
return float((lhs_val - rhs_val + math.pi) % (2 * math.pi) - math.pi)
|
|
|
|
|
|
def _decoded_visible_face_types(decoded: Optional[dict[str, Any]]) -> tuple[int, ...]:
|
|
if not decoded:
|
|
return ()
|
|
raw_types = decoded.get("visible_face_types", ())
|
|
if raw_types is None:
|
|
return ()
|
|
face_types = []
|
|
for face_type in raw_types:
|
|
face_type_scalar = _serialize_scalar(face_type)
|
|
if face_type_scalar is None:
|
|
continue
|
|
face_types.append(int(face_type_scalar))
|
|
return tuple(face_types)
|
|
|
|
|
|
def _build_edge_heading_decoded(
|
|
base_decoded: Optional[dict[str, Any]],
|
|
pred_41: np.ndarray,
|
|
pred_edge_60: Optional[np.ndarray],
|
|
anchor_xy: np.ndarray,
|
|
stride: float,
|
|
bbox_xyxy: np.ndarray,
|
|
calib: dict[str, Any],
|
|
img_w: int,
|
|
whole_attrs: Optional[dict[str, Any]],
|
|
edge_yaw_max_lateral_dist_m: float,
|
|
img_h: Optional[int] = None,
|
|
edge_selection: Optional[dict[str, Any]] = None,
|
|
) -> tuple[Optional[dict[str, Any]], float, bool]:
|
|
edge_selection = edge_selection or decode_edge_yaw_selection_from_prediction(
|
|
pred_41,
|
|
pred_edge_60,
|
|
anchor_xy,
|
|
stride,
|
|
calib,
|
|
bbox_xyxy=bbox_xyxy,
|
|
img_w=img_w,
|
|
img_h=img_h,
|
|
max_lateral_dist_m=edge_yaw_max_lateral_dist_m,
|
|
)
|
|
selected_face_types = tuple(int(face_type) for face_type in edge_selection.get("face_types", ()))
|
|
edge_yaw = float(edge_selection["yaw"])
|
|
if not bool(edge_selection.get("lateral_ok")):
|
|
return None, float(edge_yaw), False
|
|
|
|
regressed_dims = pred_41[27:30] if whole_attrs is None else whole_attrs.get("dims", pred_41[27:30])
|
|
edge_box = reconstruct_edge_based_box_from_selection(
|
|
edge_selection,
|
|
box_center_y_m=None,
|
|
regressed_dims=regressed_dims,
|
|
face_regressed_dims_by_type=extract_face_regressed_size_priors_from_prediction(pred_41),
|
|
)
|
|
if edge_box is None:
|
|
return None, float(edge_yaw), False
|
|
|
|
decoded = dict(base_decoded or {})
|
|
resolved_face_types = tuple(int(face_type) for face_type in edge_box.get("face_types", ()) or selected_face_types)
|
|
decoded["corners_3d"] = edge_box["corners_3d"]
|
|
decoded["edge_points_2d"] = edge_selection.get("edge_points_2d")
|
|
decoded["edge_points_3d"] = edge_selection.get("edge_points_3d")
|
|
decoded["visible_face_types"] = resolved_face_types
|
|
decoded["face_center_2d"] = None
|
|
decoded["face_color"] = None
|
|
decoded["edge_box_center_3d"] = edge_box["center"]
|
|
decoded["edge_box_dims"] = edge_box["dims"]
|
|
decoded["edge_box_mode"] = edge_box.get("mode")
|
|
return decoded, float(edge_yaw), True
|
|
|
|
|
|
def _build_edge_prediction_artifacts(
|
|
base_decoded: Optional[dict[str, Any]],
|
|
pred_41: np.ndarray,
|
|
pred_edge_60: Optional[np.ndarray],
|
|
anchor_xy: np.ndarray,
|
|
stride: float,
|
|
bbox_xyxy: np.ndarray,
|
|
calib: dict[str, Any],
|
|
img_w: int,
|
|
whole_attrs: Optional[dict[str, Any]],
|
|
edge_yaw_max_lateral_dist_m: float,
|
|
img_h: Optional[int] = None,
|
|
) -> dict[str, Any]:
|
|
edge_selection = decode_edge_yaw_selection_from_prediction(
|
|
pred_41,
|
|
pred_edge_60,
|
|
anchor_xy,
|
|
stride,
|
|
calib,
|
|
bbox_xyxy=bbox_xyxy,
|
|
img_w=img_w,
|
|
img_h=img_h,
|
|
max_lateral_dist_m=edge_yaw_max_lateral_dist_m,
|
|
)
|
|
regressed_dims = pred_41[27:30] if whole_attrs is None else whole_attrs.get("dims", pred_41[27:30])
|
|
edge_box = reconstruct_edge_based_box_from_selection(
|
|
edge_selection,
|
|
box_center_y_m=None,
|
|
regressed_dims=regressed_dims,
|
|
face_regressed_dims_by_type=extract_face_regressed_size_priors_from_prediction(pred_41),
|
|
)
|
|
heading_decoded, edge_yaw, edge_confident = _build_edge_heading_decoded(
|
|
base_decoded=base_decoded,
|
|
pred_41=pred_41,
|
|
pred_edge_60=pred_edge_60,
|
|
anchor_xy=anchor_xy,
|
|
stride=stride,
|
|
bbox_xyxy=bbox_xyxy,
|
|
calib=calib,
|
|
img_w=img_w,
|
|
img_h=img_h,
|
|
whole_attrs=whole_attrs,
|
|
edge_yaw_max_lateral_dist_m=edge_yaw_max_lateral_dist_m,
|
|
edge_selection=edge_selection,
|
|
)
|
|
return {
|
|
"edge_selection": edge_selection,
|
|
"edge_box": edge_box,
|
|
"heading_decoded": heading_decoded,
|
|
"edge_yaw": float(edge_yaw),
|
|
"edge_confident": bool(edge_confident),
|
|
}
|
|
|
|
|
|
def _draw_heading_lines(
|
|
img: np.ndarray,
|
|
center_uv: Any,
|
|
lines: list[tuple[str, tuple[int, int, int]]],
|
|
font_scale: float = 0.45,
|
|
thickness: int = 1,
|
|
line_gap: int = 4,
|
|
) -> np.ndarray:
|
|
if center_uv is None:
|
|
return img
|
|
uv = np.asarray(center_uv, dtype=np.float32).reshape(-1)
|
|
if uv.size < 2 or not np.isfinite(uv[:2]).all():
|
|
return img
|
|
|
|
cx, cy = int(round(float(uv[0]))), int(round(float(uv[1])))
|
|
font = cv2.FONT_HERSHEY_SIMPLEX
|
|
text_sizes = [cv2.getTextSize(text, font, font_scale, thickness)[0] for text, _ in lines]
|
|
total_h = sum(size[1] for size in text_sizes) + line_gap * max(0, len(lines) - 1)
|
|
y = cy - total_h // 2
|
|
|
|
for (text, color), (tw, th) in zip(lines, text_sizes):
|
|
baseline_y = y + th
|
|
cv2.putText(img, text, (cx - tw // 2, baseline_y), font, font_scale, color, thickness, cv2.LINE_AA)
|
|
y = baseline_y + line_gap
|
|
return img
|
|
|
|
|
|
def _extract_center_3d(center_3d: Any) -> Optional[np.ndarray]:
|
|
if center_3d is None:
|
|
return None
|
|
center = np.asarray(center_3d, dtype=np.float32).reshape(-1)
|
|
if center.size < 3 or not np.isfinite(center[:3]).all():
|
|
return None
|
|
return center
|
|
|
|
|
|
def _format_pose_lines(center_3d: Any, yaw_rad: Optional[float], color: tuple[int, int, int]) -> list[tuple[str, tuple[int, int, int]]]:
|
|
if yaw_rad is None:
|
|
return []
|
|
yaw_value = float(yaw_rad)
|
|
if not math.isfinite(yaw_value):
|
|
return []
|
|
center = _extract_center_3d(center_3d)
|
|
if center is None:
|
|
return []
|
|
lateral = abs(float(center[0]))
|
|
depth_z = float(center[2])
|
|
return [
|
|
(f"{lateral:.1f}", color),
|
|
(f"{depth_z:.1f}", color),
|
|
(f"{math.degrees(yaw_value):.1f}", color),
|
|
]
|
|
|
|
|
|
def _edge_batch_list(edge_points_2d: Any) -> list[np.ndarray]:
|
|
if edge_points_2d is None:
|
|
return []
|
|
arr = np.asarray(edge_points_2d, dtype=np.float32)
|
|
if arr.ndim == 2:
|
|
return [arr]
|
|
if arr.ndim == 3:
|
|
return [arr[index] for index in range(arr.shape[0])]
|
|
return []
|
|
|
|
|
|
def _project_selected_face_edges_from_box(
|
|
corners_3d: Any,
|
|
face_types: tuple[int, ...],
|
|
face_is_partial: tuple[bool, ...],
|
|
calib: dict[str, Any],
|
|
img_w: int,
|
|
img_h: int,
|
|
) -> Optional[np.ndarray]:
|
|
if corners_3d is None:
|
|
return None
|
|
projected_edges = []
|
|
for face_type, is_partial in zip(face_types, face_is_partial):
|
|
if bool(is_partial):
|
|
_, points_2d = project_partial_face_bottom_edge(corners_3d, int(face_type), calib, img_w, img_h, num_samples=5)
|
|
else:
|
|
_, points_2d = project_face_bottom_edge(corners_3d, int(face_type), calib, num_samples=5)
|
|
if points_2d is None:
|
|
return None
|
|
projected_edges.append(np.asarray(points_2d, dtype=np.float32))
|
|
if not projected_edges:
|
|
return None
|
|
if len(projected_edges) == 1:
|
|
return projected_edges[0]
|
|
return np.stack(projected_edges, axis=0)
|
|
|
|
|
|
def _selected_edge_residual_stats(reference_edge_points_2d: Any, projected_edge_points_2d: Any) -> dict[str, Any]:
|
|
reference_batches = _edge_batch_list(reference_edge_points_2d)
|
|
projected_batches = _edge_batch_list(projected_edge_points_2d)
|
|
if not reference_batches or len(reference_batches) != len(projected_batches):
|
|
return {"available": False, "mean_px": None, "max_px": None, "per_face_mean_px": None}
|
|
|
|
per_face_mean = []
|
|
all_residuals = []
|
|
for ref_points, proj_points in zip(reference_batches, projected_batches):
|
|
if ref_points.shape != proj_points.shape or ref_points.ndim != 2 or ref_points.shape[1] != 2:
|
|
return {"available": False, "mean_px": None, "max_px": None, "per_face_mean_px": None}
|
|
residuals = np.linalg.norm(np.asarray(ref_points, dtype=np.float32) - np.asarray(proj_points, dtype=np.float32), axis=1)
|
|
per_face_mean.append(float(np.mean(residuals)))
|
|
all_residuals.append(residuals)
|
|
|
|
flat = np.concatenate(all_residuals, axis=0) if all_residuals else np.zeros((0,), dtype=np.float32)
|
|
if flat.size == 0:
|
|
return {"available": False, "mean_px": None, "max_px": None, "per_face_mean_px": None}
|
|
return {
|
|
"available": True,
|
|
"mean_px": float(np.mean(flat)),
|
|
"max_px": float(np.max(flat)),
|
|
"per_face_mean_px": [float(value) for value in per_face_mean],
|
|
}
|
|
|
|
|
|
def decode_prepared_roi_predictions(
|
|
bundle: LoadedROIModel,
|
|
prepared: PreparedROI,
|
|
filtered_outputs: RawROIOutputs,
|
|
edge_yaw_max_lateral_dist_m: float,
|
|
) -> list[dict[str, Any]]:
|
|
detections = filtered_outputs.detections
|
|
preds_3d = filtered_outputs.preds_3d
|
|
preds_3d_fake = filtered_outputs.preds_3d_fake
|
|
preds_edge = filtered_outputs.preds_edge
|
|
anchors = filtered_outputs.anchors
|
|
strides = filtered_outputs.strides
|
|
|
|
predictions = []
|
|
for index, det in enumerate(detections):
|
|
pred_edge = None if preds_edge is None else preds_edge[index]
|
|
bbox_xyxy = det[:4].astype(np.float32)
|
|
cls_id = int(det[5])
|
|
is_face_3d_class = cls_id in bundle.face_3d_classes
|
|
pred_3d = preds_3d_fake[index] if preds_3d_fake is not None and cls_id in bundle.fake_3d_classes else preds_3d[index]
|
|
decoded = decode_3d_prediction(
|
|
pred_3d,
|
|
anchors[:, index],
|
|
float(strides[index]),
|
|
prepared.calib,
|
|
prepared.image.shape[1],
|
|
prepared.image.shape[0],
|
|
bundle.face_3d_classes,
|
|
bundle.complete_3d_classes,
|
|
cls_id,
|
|
pred_edge_60=pred_edge,
|
|
bbox_xyxy=bbox_xyxy,
|
|
)
|
|
face_anchor_type = None if decoded is None else decoded.get("visible_face_type")
|
|
attr_face_type = int(face_anchor_type) if is_face_3d_class and face_anchor_type is not None else None
|
|
attrs = extract_3d_attrs_from_prediction(
|
|
pred_3d,
|
|
anchors[:, index],
|
|
float(strides[index]),
|
|
prepared.calib,
|
|
face_type=attr_face_type,
|
|
pred_edge_60=pred_edge,
|
|
)
|
|
if is_face_3d_class:
|
|
edge_artifacts = _build_edge_prediction_artifacts(
|
|
base_decoded=decoded,
|
|
pred_41=pred_3d,
|
|
pred_edge_60=pred_edge,
|
|
anchor_xy=anchors[:, index],
|
|
stride=float(strides[index]),
|
|
bbox_xyxy=bbox_xyxy,
|
|
calib=prepared.calib,
|
|
img_w=prepared.image.shape[1],
|
|
img_h=prepared.image.shape[0],
|
|
whole_attrs=attrs,
|
|
edge_yaw_max_lateral_dist_m=float(edge_yaw_max_lateral_dist_m),
|
|
)
|
|
edge_selection = edge_artifacts["edge_selection"]
|
|
edge_box = edge_artifacts["edge_box"]
|
|
edge_heading_decoded = edge_artifacts["heading_decoded"]
|
|
edge_yaw = edge_artifacts["edge_yaw"]
|
|
edge_confident = edge_artifacts["edge_confident"]
|
|
else:
|
|
edge_selection = None
|
|
edge_box = None
|
|
edge_heading_decoded = None
|
|
edge_yaw = float("nan")
|
|
edge_confident = False
|
|
predictions.append(
|
|
{
|
|
"bbox_xyxy": bbox_xyxy,
|
|
"confidence": float(det[4]),
|
|
"cls_id": cls_id,
|
|
"pred_41": pred_3d,
|
|
"used_fake_3d_head": bool(preds_3d_fake is not None and cls_id in bundle.fake_3d_classes),
|
|
"pred_edge_60": pred_edge,
|
|
"anchor_xy": anchors[:, index],
|
|
"stride": float(strides[index]),
|
|
"attrs": attrs,
|
|
"decoded": decoded,
|
|
"edge_selection": edge_selection,
|
|
"edge_box": edge_box,
|
|
"edge_heading_decoded": edge_heading_decoded,
|
|
"edge_yaw": float(edge_yaw),
|
|
"edge_confident": bool(edge_confident),
|
|
}
|
|
)
|
|
return predictions
|
|
|
|
|
|
def infer_prepared_roi_batch(
|
|
bundle: LoadedROIModel,
|
|
prepared_batch: list[PreparedROI],
|
|
classes: Optional[set[int]],
|
|
edge_yaw_max_lateral_dist_m: float,
|
|
conf_thres: Optional[float] = None,
|
|
raw_outputs_batch: Optional[list[RawROIOutputs]] = None,
|
|
device: Optional[torch.device] = None,
|
|
use_half: Optional[bool] = None,
|
|
) -> list[dict[str, Any]]:
|
|
resolved_raw_outputs_batch = (
|
|
raw_outputs_batch
|
|
if raw_outputs_batch is not None
|
|
else run_model_for_prepared_roi_batch(bundle, prepared_batch, device=device, use_half=use_half)
|
|
)
|
|
if len(resolved_raw_outputs_batch) != len(prepared_batch):
|
|
raise RuntimeError(
|
|
f"Prepared ROI batch size mismatch for {bundle.spec.name}: "
|
|
f"prepared={len(prepared_batch)} raw_outputs={len(resolved_raw_outputs_batch)}."
|
|
)
|
|
|
|
resolved_conf_thres = float(bundle.spec.conf if conf_thres is None else conf_thres)
|
|
results = []
|
|
for prepared, raw_outputs in zip(prepared_batch, resolved_raw_outputs_batch):
|
|
filtered_outputs = filter_prediction_outputs(
|
|
raw_outputs=raw_outputs,
|
|
conf_thres=resolved_conf_thres,
|
|
max_det=bundle.spec.max_det,
|
|
classes=classes,
|
|
)
|
|
results.append(
|
|
{
|
|
"prepared": prepared,
|
|
"raw_outputs": raw_outputs,
|
|
"filtered_outputs": filtered_outputs,
|
|
"predictions": decode_prepared_roi_predictions(
|
|
bundle=bundle,
|
|
prepared=prepared,
|
|
filtered_outputs=filtered_outputs,
|
|
edge_yaw_max_lateral_dist_m=edge_yaw_max_lateral_dist_m,
|
|
),
|
|
}
|
|
)
|
|
return results
|
|
|
|
|
|
def _build_visualized_roi_result(
|
|
bundle: LoadedROIModel,
|
|
prepared: PreparedROI,
|
|
filtered_outputs: RawROIOutputs,
|
|
predictions: list[dict[str, Any]],
|
|
) -> dict[str, Any]:
|
|
panel_2d = _draw_2d_boxes(prepared.image, filtered_outputs.detections, bundle.names)
|
|
panel_3d = prepared.image.copy()
|
|
panel_3d_edge = prepared.image.copy()
|
|
pred_records = []
|
|
|
|
for prediction in predictions:
|
|
bbox_xyxy = np.asarray(prediction["bbox_xyxy"], dtype=np.float32)
|
|
cls_id = int(prediction["cls_id"])
|
|
decoded = prediction.get("decoded")
|
|
edge_selection = prediction.get("edge_selection") or {
|
|
"face_types": (),
|
|
"face_is_partial": (),
|
|
"edge_points_2d": None,
|
|
"lateral_distance_m": None,
|
|
"lateral_ok": False,
|
|
"two_face_eligible": False,
|
|
}
|
|
edge_box = prediction.get("edge_box")
|
|
heading_decoded = prediction.get("edge_heading_decoded")
|
|
pred_attrs = prediction.get("attrs")
|
|
visible_face_types = _decoded_visible_face_types(decoded)
|
|
|
|
selected_face_types = tuple(int(face_type) for face_type in edge_selection.get("face_types", ()))
|
|
selected_face_is_partial = tuple(bool(flag) for flag in edge_selection.get("face_is_partial", ()))
|
|
selected_edge_points_2d = edge_selection.get("edge_points_2d")
|
|
direct_box_selected_edges_2d = _project_selected_face_edges_from_box(
|
|
None if decoded is None else decoded.get("corners_3d"),
|
|
selected_face_types,
|
|
selected_face_is_partial,
|
|
prepared.calib,
|
|
prepared.image.shape[1],
|
|
prepared.image.shape[0],
|
|
)
|
|
edge_box_face_types = tuple(int(face_type) for face_type in (edge_box or {}).get("face_types", ()) or selected_face_types)
|
|
edge_box_selected_edges_2d = _project_selected_face_edges_from_box(
|
|
None if heading_decoded is None else heading_decoded.get("corners_3d"),
|
|
edge_box_face_types,
|
|
selected_face_is_partial,
|
|
prepared.calib,
|
|
prepared.image.shape[1],
|
|
prepared.image.shape[0],
|
|
)
|
|
direct_edge_fit = _selected_edge_residual_stats(selected_edge_points_2d, direct_box_selected_edges_2d)
|
|
edge_edge_fit = _selected_edge_residual_stats(selected_edge_points_2d, edge_box_selected_edges_2d)
|
|
|
|
reg_yaw = None if pred_attrs is None else _serialize_scalar(pred_attrs["yaw"])
|
|
center_3d = None if pred_attrs is None else pred_attrs.get("center")
|
|
edge_yaw = prediction.get("edge_yaw")
|
|
edge_vs_reg_yaw_rad = _wrapped_angle_diff_rad(edge_yaw, reg_yaw)
|
|
record = {
|
|
"bbox_xyxy": _serialize_array(bbox_xyxy),
|
|
"confidence": float(prediction["confidence"]),
|
|
"cls_id": cls_id,
|
|
"cls_name": _class_name(bundle.names, cls_id),
|
|
"yaw_rad": reg_yaw,
|
|
"edge_yaw_rad": _serialize_scalar(edge_yaw),
|
|
"edge_yaw_confident": bool(prediction.get("edge_confident")),
|
|
"edge_yaw_lateral_distance_m": _serialize_scalar(edge_selection.get("lateral_distance_m")),
|
|
"edge_yaw_lateral_ok": bool(edge_selection.get("lateral_ok")),
|
|
"edge_yaw_two_face_eligible": bool(edge_selection.get("two_face_eligible")),
|
|
"edge_yaw_selected_face_types": _serialize_array(selected_face_types),
|
|
"edge_yaw_selected_face_is_partial": _serialize_array(selected_face_is_partial),
|
|
"edge_vs_reg_yaw_rad": _serialize_scalar(edge_vs_reg_yaw_rad),
|
|
"selected_edge_direct_box_fit_available": bool(direct_edge_fit["available"]),
|
|
"selected_edge_direct_box_fit_mean_px": _serialize_scalar(direct_edge_fit["mean_px"]),
|
|
"selected_edge_direct_box_fit_max_px": _serialize_scalar(direct_edge_fit["max_px"]),
|
|
"selected_edge_direct_box_fit_per_face_mean_px": _serialize_array(direct_edge_fit["per_face_mean_px"]),
|
|
"selected_edge_edgeyaw_box_fit_available": bool(edge_edge_fit["available"]),
|
|
"selected_edge_edgeyaw_box_fit_mean_px": _serialize_scalar(edge_edge_fit["mean_px"]),
|
|
"selected_edge_edgeyaw_box_fit_max_px": _serialize_scalar(edge_edge_fit["max_px"]),
|
|
"selected_edge_edgeyaw_box_fit_per_face_mean_px": _serialize_array(edge_edge_fit["per_face_mean_px"]),
|
|
"selected_edge_fit_gain_px": _serialize_scalar(
|
|
None
|
|
if direct_edge_fit["mean_px"] is None or edge_edge_fit["mean_px"] is None
|
|
else float(direct_edge_fit["mean_px"]) - float(edge_edge_fit["mean_px"])
|
|
),
|
|
"edge_box_center_3d": None if edge_box is None else _serialize_array(edge_box["center"]),
|
|
"edge_box_dims": None if edge_box is None else _serialize_array(edge_box["dims"]),
|
|
"edge_box_length_m": _serialize_scalar(None if edge_box is None else edge_box.get("side_length_m")),
|
|
"edge_box_width_m": _serialize_scalar(None if edge_box is None else edge_box.get("width_m")),
|
|
"edge_box_mode": None if edge_box is None else str(edge_box.get("mode")),
|
|
"edge_box_length_source": None if edge_box is None else edge_box.get("length_source"),
|
|
"edge_box_width_source": None if edge_box is None else edge_box.get("width_source"),
|
|
"center_uv": None if pred_attrs is None else _serialize_array(pred_attrs["uv"]),
|
|
"center_3d": _serialize_array(center_3d),
|
|
"visible_face_type": None if decoded is None else _serialize_scalar(decoded.get("visible_face_type")),
|
|
"visible_face_count": len(visible_face_types),
|
|
"visible_face_types": None if decoded is None else _serialize_array(visible_face_types),
|
|
"crop_bounds": list(prepared.crop_bounds),
|
|
}
|
|
pred_records.append(record)
|
|
|
|
center_uv = record.get("center_uv")
|
|
if decoded is not None and decoded.get("corners_3d") is not None:
|
|
draw_3d_box(
|
|
panel_3d,
|
|
decoded["corners_3d"],
|
|
prepared.calib,
|
|
decoded.get("face_center_2d"),
|
|
decoded.get("face_color"),
|
|
edge_points_2d=decoded.get("edge_points_2d"),
|
|
edge_color=(0, 255, 0),
|
|
thickness=1,
|
|
)
|
|
reg_lines = _format_pose_lines(record.get("center_3d"), record.get("yaw_rad"), (0, 0, 0))
|
|
if reg_lines:
|
|
_draw_heading_lines(panel_3d, center_uv, reg_lines, font_scale=0.35, thickness=1, line_gap=2)
|
|
|
|
if bool(record.get("edge_yaw_confident")) and heading_decoded is not None and heading_decoded.get("corners_3d") is not None:
|
|
draw_3d_box(
|
|
panel_3d_edge,
|
|
heading_decoded["corners_3d"],
|
|
prepared.calib,
|
|
heading_decoded.get("face_center_2d"),
|
|
heading_decoded.get("face_color"),
|
|
edge_points_2d=heading_decoded.get("edge_points_2d"),
|
|
edge_color=(0, 255, 0),
|
|
thickness=1,
|
|
)
|
|
edge_lines = _format_pose_lines(record.get("edge_box_center_3d") or record.get("center_3d"), record.get("edge_yaw_rad"), (0, 0, 0))
|
|
if edge_lines:
|
|
_draw_heading_lines(panel_3d_edge, center_uv, edge_lines, font_scale=0.35, thickness=1, line_gap=2)
|
|
|
|
return {
|
|
"prepared": prepared,
|
|
"detections": filtered_outputs.detections,
|
|
"predictions": pred_records,
|
|
"panel_2d": _annotate_panel_title(panel_2d, f"{bundle.spec.name} 2D"),
|
|
"panel_3d": _annotate_panel_title(panel_3d, f"{bundle.spec.name} 3D"),
|
|
"panel_3d_edge": _annotate_panel_title(panel_3d_edge, f"{bundle.spec.name} 3D EdgeRecon (1+ face)"),
|
|
}
|
|
|
|
|
|
def _run_single_roi(
|
|
bundle: LoadedROIModel,
|
|
prepared: PreparedROI,
|
|
device: torch.device,
|
|
use_half: bool,
|
|
classes: Optional[set[int]],
|
|
edge_yaw_max_lateral_dist_m: float,
|
|
) -> dict[str, Any]:
|
|
inference_result = infer_prepared_roi_batch(
|
|
bundle=bundle,
|
|
prepared_batch=[prepared],
|
|
classes=classes,
|
|
edge_yaw_max_lateral_dist_m=edge_yaw_max_lateral_dist_m,
|
|
raw_outputs_batch=None,
|
|
device=device,
|
|
use_half=use_half,
|
|
)[0]
|
|
return _build_visualized_roi_result(
|
|
bundle=bundle,
|
|
prepared=inference_result["prepared"],
|
|
filtered_outputs=inference_result["filtered_outputs"],
|
|
predictions=inference_result["predictions"],
|
|
)
|
|
|
|
|
|
def _assemble_grid(roi_results: list[dict[str, Any]]) -> np.ndarray:
|
|
row_images = []
|
|
for row_idx, roi_result in enumerate(roi_results):
|
|
panels = [roi_result["panel_2d"], roi_result["panel_3d"], roi_result["panel_3d_edge"]]
|
|
row_h = max(panel.shape[0] for panel in panels)
|
|
row_w = sum(panel.shape[1] for panel in panels)
|
|
row_canvas = np.full((row_h, row_w, 3), GRID_BG_COLOR, dtype=np.uint8)
|
|
x0 = 0
|
|
for panel in panels:
|
|
panel_h, panel_w = panel.shape[:2]
|
|
y0 = max(0, (row_h - panel_h) // 2)
|
|
row_canvas[y0 : y0 + panel_h, x0 : x0 + panel_w] = panel
|
|
x0 += panel_w
|
|
row_images.append(row_canvas)
|
|
|
|
grid_h = sum(row.shape[0] for row in row_images)
|
|
grid_w = max(row.shape[1] for row in row_images)
|
|
grid = np.full((grid_h, grid_w, 3), GRID_BG_COLOR, dtype=np.uint8)
|
|
y0 = 0
|
|
for row in row_images:
|
|
row_h, row_w = row.shape[:2]
|
|
x0 = max(0, (grid_w - row_w) // 2)
|
|
grid[y0 : y0 + row_h, x0 : x0 + row_w] = row
|
|
y0 += row_h
|
|
return grid
|
|
|
|
|
|
def iter_case_images(images_dir: str | Path, glob_pattern: str, max_images: int = 0) -> Iterable[Path]:
|
|
images_dir = Path(images_dir)
|
|
image_paths = [path for path in sorted(images_dir.glob(glob_pattern)) if path.is_file() and path.suffix.lower() in IMAGE_SUFFIXES]
|
|
if max_images > 0:
|
|
image_paths = image_paths[:max_images]
|
|
return image_paths
|
|
|
|
|
|
def iter_loaded_case_images(
|
|
images_dir: str | Path,
|
|
glob_pattern: str,
|
|
max_images: int = 0,
|
|
) -> Iterable[tuple[str, np.ndarray]]:
|
|
for image_path in iter_case_images(images_dir, glob_pattern=glob_pattern, max_images=max_images):
|
|
image_bgr = cv2.imread(str(image_path), cv2.IMREAD_COLOR)
|
|
if image_bgr is None:
|
|
raise FileNotFoundError(f"Failed to read image: {image_path}")
|
|
yield image_path.name, image_bgr
|
|
|
|
|
|
def build_inference_context(
|
|
roi_specs: list[ROIModelSpec],
|
|
device: str = "0",
|
|
half: bool = False,
|
|
classes: Optional[list[int]] = None,
|
|
edge_yaw_max_lateral_dist_m: float = DEFAULT_EDGE_YAW_MAX_LATERAL_DIST_M,
|
|
inference_batch_size: int = 1,
|
|
) -> InferenceContext:
|
|
requested_device = device or ""
|
|
try:
|
|
selected_device = select_device(requested_device)
|
|
except ValueError as exc:
|
|
# Allow batch jobs to keep running on CPU-only environments even if the
|
|
# default CLI device is set to `0`.
|
|
if requested_device and requested_device.lower() != "cpu" and not torch.cuda.is_available():
|
|
print(
|
|
f"[WARN] Requested device={requested_device!r} but CUDA is unavailable; "
|
|
"falling back to CPU."
|
|
)
|
|
selected_device = select_device("cpu")
|
|
else:
|
|
raise exc
|
|
use_half = bool(half and selected_device.type != "cpu")
|
|
roi_models = []
|
|
for spec in roi_specs:
|
|
model, _ = load_checkpoint(spec.model_path, device=selected_device, fuse=False)
|
|
if use_half:
|
|
model = model.half()
|
|
names = getattr(model, "names", {}) or {}
|
|
fake_3d_classes = _infer_fake_3d_classes(model, names)
|
|
roi_models.append(
|
|
LoadedROIModel(
|
|
spec=spec,
|
|
model=model.eval(),
|
|
names=names,
|
|
face_3d_classes=set(getattr(model, "face_3d_classes", set())),
|
|
complete_3d_classes=set(getattr(model, "complete_3d_classes", set())),
|
|
fake_3d_classes=fake_3d_classes,
|
|
imgsz=_infer_model_imgsz(model, spec.imgsz),
|
|
)
|
|
)
|
|
return InferenceContext(
|
|
device=selected_device,
|
|
use_half=use_half,
|
|
classes=None if classes is None else set(int(cls_id) for cls_id in classes),
|
|
roi_models=roi_models,
|
|
edge_yaw_max_lateral_dist_m=float(edge_yaw_max_lateral_dist_m),
|
|
inference_batch_size=max(1, int(inference_batch_size)),
|
|
)
|
|
|
|
|
|
def _resolve_case_inputs(case_dir: str = "", images_dir: str = "", calib_file: str = "") -> tuple[Path, Path, Optional[Path]]:
|
|
if case_dir:
|
|
case_path = Path(case_dir).resolve()
|
|
return case_path / "images", case_path / "calib" / "L2_calib" / "camera4.json", case_path
|
|
if not images_dir or not calib_file:
|
|
raise ValueError("Either --case-dir or both --images-dir and --calib-file are required.")
|
|
return Path(images_dir).resolve(), Path(calib_file).resolve(), None
|
|
|
|
|
|
def _resolve_visualization_output_dir(output_dir: str | Path, case_name: str) -> Path:
|
|
output_dir = Path(output_dir)
|
|
if output_dir == DEFAULT_VISUALIZATION_ROOT or output_dir == Path(str(DEFAULT_VISUALIZATION_ROOT)):
|
|
output_dir = output_dir / case_name
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
return output_dir
|
|
|
|
|
|
def create_predictions_payload(
|
|
context: InferenceContext,
|
|
case_name: str,
|
|
source_info: Optional[dict[str, Any]] = None,
|
|
) -> dict[str, Any]:
|
|
payload = {
|
|
"case_name": case_name,
|
|
"edge_yaw_max_lateral_dist_m": context.edge_yaw_max_lateral_dist_m,
|
|
"frames": [],
|
|
}
|
|
if source_info:
|
|
payload.update(source_info)
|
|
return payload
|
|
|
|
|
|
def _run_inference_on_frame_batch(
|
|
context: InferenceContext,
|
|
raw_calib: dict[str, Any],
|
|
frame_batch: list[tuple[int, str, np.ndarray]],
|
|
output_dir: Path,
|
|
) -> list[dict[str, Any]]:
|
|
if not frame_batch:
|
|
return []
|
|
|
|
frame_payloads = [
|
|
{
|
|
"frame_index": int(frame_index),
|
|
"frame_name": frame_name,
|
|
"rois": {},
|
|
}
|
|
for frame_index, frame_name, _image_bgr in frame_batch
|
|
]
|
|
roi_results_by_frame: list[list[dict[str, Any]]] = [[] for _ in frame_batch]
|
|
batch_images = [image_bgr for _frame_index, _frame_name, image_bgr in frame_batch]
|
|
raw_calibs = [raw_calib for _ in frame_batch]
|
|
|
|
for bundle in context.roi_models:
|
|
prepared_batch = prepare_roi_batch(batch_images, raw_calibs, bundle.spec, bundle.imgsz)
|
|
inference_results = infer_prepared_roi_batch(
|
|
bundle=bundle,
|
|
prepared_batch=prepared_batch,
|
|
classes=context.classes,
|
|
edge_yaw_max_lateral_dist_m=context.edge_yaw_max_lateral_dist_m,
|
|
device=context.device,
|
|
use_half=context.use_half,
|
|
)
|
|
for frame_offset, inference_result in enumerate(inference_results):
|
|
prepared = inference_result["prepared"]
|
|
roi_result = _build_visualized_roi_result(
|
|
bundle=bundle,
|
|
prepared=prepared,
|
|
filtered_outputs=inference_result["filtered_outputs"],
|
|
predictions=inference_result["predictions"],
|
|
)
|
|
roi_results_by_frame[frame_offset].append(roi_result)
|
|
frame_payloads[frame_offset]["rois"][bundle.spec.name.lower()] = {
|
|
"crop_bounds": list(prepared.crop_bounds),
|
|
"vp_x": prepared.vp_x,
|
|
"vp_y": prepared.vp_y,
|
|
"crop_center_x": prepared.crop_center_x,
|
|
"crop_center_y": prepared.crop_center_y,
|
|
"edge_yaw_max_lateral_dist_m": context.edge_yaw_max_lateral_dist_m,
|
|
"calib": {key: _serialize_scalar(value) for key, value in prepared.calib.items()},
|
|
"predictions": roi_result["predictions"],
|
|
}
|
|
|
|
for frame_offset, ((frame_index, frame_name, _image_bgr), roi_results) in enumerate(zip(frame_batch, roi_results_by_frame)):
|
|
grid = _assemble_grid(roi_results)
|
|
frame_stem = Path(frame_name).stem or f"frame_{frame_index:06d}"
|
|
grid_path = output_dir / f"{frame_stem}.jpg"
|
|
if not cv2.imwrite(str(grid_path), grid):
|
|
raise IOError(f"Failed to write visualization image: {grid_path}")
|
|
frame_payloads[frame_offset]["visualization"] = str(grid_path)
|
|
|
|
return frame_payloads
|
|
|
|
|
|
def _run_inference_on_frame(
|
|
context: InferenceContext,
|
|
raw_calib: dict[str, Any],
|
|
image_bgr: np.ndarray,
|
|
frame_name: str,
|
|
frame_index: int,
|
|
output_dir: Path,
|
|
) -> dict[str, Any]:
|
|
frame_payloads = _run_inference_on_frame_batch(
|
|
context=context,
|
|
raw_calib=raw_calib,
|
|
frame_batch=[(frame_index, frame_name, image_bgr)],
|
|
output_dir=output_dir,
|
|
)
|
|
if not frame_payloads:
|
|
raise RuntimeError(f"Failed to run inference for frame {frame_name}.")
|
|
return frame_payloads[0]
|
|
|
|
|
|
def append_image_stream_inference(
|
|
context: InferenceContext,
|
|
frames: Iterable[tuple[str, np.ndarray]],
|
|
raw_calib: dict[str, Any],
|
|
output_dir: str | Path,
|
|
predictions_payload: dict[str, Any],
|
|
frame_index_offset: int = 0,
|
|
frame_name_prefix: str = "",
|
|
) -> int:
|
|
output_dir = Path(output_dir)
|
|
num_frames = 0
|
|
batch_size = max(1, int(getattr(context, "inference_batch_size", 1)))
|
|
pending_frames: list[tuple[int, str, np.ndarray]] = []
|
|
for local_frame_index, (frame_name, image_bgr) in enumerate(frames):
|
|
effective_frame_name = f"{frame_name_prefix}_{frame_name}" if frame_name_prefix else frame_name
|
|
pending_frames.append((frame_index_offset + local_frame_index, effective_frame_name, image_bgr))
|
|
if len(pending_frames) < batch_size:
|
|
continue
|
|
frame_payloads = _run_inference_on_frame_batch(
|
|
context=context,
|
|
raw_calib=raw_calib,
|
|
frame_batch=pending_frames,
|
|
output_dir=output_dir,
|
|
)
|
|
predictions_payload["frames"].extend(frame_payloads)
|
|
num_frames += len(frame_payloads)
|
|
pending_frames = []
|
|
|
|
if pending_frames:
|
|
frame_payloads = _run_inference_on_frame_batch(
|
|
context=context,
|
|
raw_calib=raw_calib,
|
|
frame_batch=pending_frames,
|
|
output_dir=output_dir,
|
|
)
|
|
predictions_payload["frames"].extend(frame_payloads)
|
|
num_frames += len(frame_payloads)
|
|
|
|
return num_frames
|
|
|
|
|
|
def run_image_stream_inference(
|
|
context: InferenceContext,
|
|
frames: Iterable[tuple[str, np.ndarray]],
|
|
raw_calib: dict[str, Any],
|
|
case_name: str,
|
|
output_dir: str | Path = DEFAULT_VISUALIZATION_ROOT,
|
|
source_info: Optional[dict[str, Any]] = None,
|
|
) -> dict[str, Any]:
|
|
output_dir = _resolve_visualization_output_dir(output_dir, case_name)
|
|
predictions_payload = create_predictions_payload(
|
|
context=context,
|
|
case_name=case_name,
|
|
source_info=source_info,
|
|
)
|
|
num_frames = append_image_stream_inference(
|
|
context=context,
|
|
frames=frames,
|
|
raw_calib=raw_calib,
|
|
output_dir=output_dir,
|
|
predictions_payload=predictions_payload,
|
|
frame_index_offset=0,
|
|
frame_name_prefix="",
|
|
)
|
|
|
|
if num_frames == 0:
|
|
raise FileNotFoundError(f"No frames available for case {case_name}")
|
|
|
|
predictions_path = output_dir / "predictions.json"
|
|
with predictions_path.open("w", encoding="utf-8") as file:
|
|
json.dump(predictions_payload, file, indent=2, ensure_ascii=False)
|
|
|
|
return {
|
|
"case_name": case_name,
|
|
"output_dir": str(output_dir),
|
|
"predictions_path": str(predictions_path),
|
|
"num_frames": num_frames,
|
|
}
|
|
|
|
|
|
def run_case_inference(
|
|
context: InferenceContext,
|
|
case_dir: str = "",
|
|
images_dir: str = "",
|
|
calib_file: str = "",
|
|
output_dir: str | Path = DEFAULT_VISUALIZATION_ROOT,
|
|
glob_pattern: str = "*.png",
|
|
max_images: int = 0,
|
|
) -> dict[str, Any]:
|
|
images_dir_path, calib_path, resolved_case_dir = _resolve_case_inputs(case_dir=case_dir, images_dir=images_dir, calib_file=calib_file)
|
|
raw_calib = load_camera4_calib(calib_path)
|
|
|
|
case_name = resolved_case_dir.name if resolved_case_dir is not None else images_dir_path.name
|
|
return run_image_stream_inference(
|
|
context=context,
|
|
frames=iter_loaded_case_images(images_dir_path, glob_pattern=glob_pattern, max_images=max_images),
|
|
raw_calib=raw_calib,
|
|
case_name=case_name,
|
|
output_dir=output_dir,
|
|
source_info={
|
|
"images_dir": str(images_dir_path),
|
|
"calib_file": str(calib_path),
|
|
},
|
|
)
|
|
|
|
|
|
def iter_video_case_images(
|
|
video_path: str | Path,
|
|
frame_index_payload: Optional[dict[str, Any]] = None,
|
|
video_stride: int = 1,
|
|
max_images: int = 0,
|
|
) -> Iterable[tuple[str, np.ndarray]]:
|
|
for frame_index, image_bgr, frame_name, _frame_info in iter_video_case_frames(
|
|
video_path,
|
|
frame_index_payload=frame_index_payload,
|
|
frame_stride=video_stride,
|
|
max_frames=max_images,
|
|
):
|
|
yield f"{int(frame_index):06d}_{Path(frame_name).name}", image_bgr
|
|
|
|
|
|
def run_video_case_inference(
|
|
context: InferenceContext,
|
|
video_case_dir: str | Path,
|
|
output_dir: str | Path = DEFAULT_VISUALIZATION_ROOT,
|
|
max_images: int = 0,
|
|
video_stride: int = 1,
|
|
) -> dict[str, Any]:
|
|
case_dir, video_path, calib_path = resolve_video_case_paths(video_case_dir)
|
|
raw_calib = load_camera4_calib(calib_path)
|
|
frame_index_payload = read_video_frame_index(video_path)
|
|
|
|
return run_image_stream_inference(
|
|
context=context,
|
|
frames=iter_video_case_images(
|
|
video_path=video_path,
|
|
frame_index_payload=frame_index_payload,
|
|
video_stride=video_stride,
|
|
max_images=max_images,
|
|
),
|
|
raw_calib=raw_calib,
|
|
case_name=case_dir.name,
|
|
output_dir=output_dir,
|
|
source_info={
|
|
"images_dir": str(video_path.parent),
|
|
"calib_file": str(calib_path),
|
|
"input_mode": "video_case",
|
|
"video_case_dir": str(case_dir),
|
|
"video_path": str(video_path),
|
|
},
|
|
)
|
|
|
|
|
|
def build_roi_specs_from_args(args: Any) -> list[ROIModelSpec]:
|
|
from tools.pdcl_inference.run_batch_two_roi_infer import build_roi_specs_from_args as _build_roi_specs_from_args
|
|
|
|
return _build_roi_specs_from_args(args)
|
|
|
|
|
|
def add_two_roi_inference_args(parser: Any, include_output_dir: bool = True) -> None:
|
|
from tools.pdcl_inference.run_batch_two_roi_infer import add_two_roi_inference_args as _add_two_roi_inference_args
|
|
|
|
_add_two_roi_inference_args(parser, include_output_dir=include_output_dir)
|