from __future__ import annotations import json from pathlib import Path from typing import Any, Iterable, Optional import cv2 import numpy as np DEFAULT_CNCAP_PATH_PREFIX_SRC = "/mnt/hfs/project-G1M3" DEFAULT_CNCAP_PATH_PREFIX_DST = "/mnt/G1M3" DEFAULT_OUTPUT_RELATIVE_ANCHORS = ("CNCAP2024数采", "gt_org_data") def rewrite_path_prefix(path_str: str, prefix_src: str, prefix_dst: str) -> str: normalized_path = str(path_str).strip() normalized_src = str(prefix_src).rstrip("/") normalized_dst = str(prefix_dst).rstrip("/") if normalized_src and normalized_path.startswith(normalized_src): return f"{normalized_dst}{normalized_path[len(normalized_src):]}" return normalized_path def load_path_list_from_json(json_file: str | Path, values_key: str = "values") -> list[str]: json_path = Path(json_file).resolve() with json_path.open("r", encoding="utf-8") as file: payload = json.load(file) if not isinstance(payload, dict): raise ValueError(f"Expected top-level dict in {json_path}, got {type(payload).__name__}") values = payload.get(values_key) if not isinstance(values, list): raise ValueError(f"Expected {values_key!r} list in {json_path}, got {type(values).__name__}") return [str(item).strip() for item in values if str(item).strip()] def _normalize_video_case_input(video_case_dir: str | Path) -> Path: input_path = Path(video_case_dir).resolve() if input_path.is_file() and input_path.name == "camera4.bin": return input_path.parent.parent if input_path.is_dir() and input_path.name == "sigmastar.1": return input_path.parent return input_path def build_case_output_rel_dir( case_dir: str | Path, preferred_anchor_names: Iterable[str] = DEFAULT_OUTPUT_RELATIVE_ANCHORS, ) -> Path: resolved_case_dir = Path(case_dir).resolve() parts = resolved_case_dir.parts for anchor_name in preferred_anchor_names: if anchor_name in parts: anchor_index = parts.index(anchor_name) suffix_parts = parts[anchor_index + 1 :] if suffix_parts: return Path(*suffix_parts) if len(parts) >= 3: return Path(*parts[-3:]) if len(parts) >= 2: return Path(*parts[-2:]) if parts: return Path(parts[-1]) return Path("case") def resolve_video_case_paths(video_case_dir: str | Path) -> tuple[Path, Path, Path]: case_dir = _normalize_video_case_input(video_case_dir) if not case_dir.is_dir(): raise FileNotFoundError(f"Video case directory not found: {case_dir}") video_path = case_dir / "sigmastar.1" / "camera4.bin" if not video_path.is_file(): raise FileNotFoundError(f"camera4.bin not found under {case_dir}") calib_candidates = [ case_dir / "test_data" / "calibs" / "camera4.json", case_dir.parent / "test_data" / "calibs" / "camera4.json", case_dir / "sigmastar.1" / "calibs" / "camera4.json", case_dir / "calibs" / "camera4.json", ] calib_path = next((path for path in calib_candidates if path.is_file()), None) if calib_path is None: checked = ", ".join(str(path) for path in calib_candidates) raise FileNotFoundError(f"camera4.json not found for {case_dir}. Checked: {checked}") return case_dir, video_path, calib_path def collect_video_case_dirs(video_root_dir: str | Path) -> list[Path]: root_dir = Path(video_root_dir).resolve() if not root_dir.is_dir(): raise FileNotFoundError(f"Video root directory not found: {root_dir}") case_dirs: list[Path] = [] for path in sorted(root_dir.iterdir()): if not path.is_dir(): continue try: resolve_video_case_paths(path) except FileNotFoundError: continue case_dirs.append(path) if not case_dirs: raise FileNotFoundError(f"No valid video case directories found under {root_dir}") return case_dirs def read_video_frame_index(video_path: str | Path) -> Optional[dict[str, Any]]: try: video_path_obj = Path(video_path).resolve() case_dir = video_path_obj.parent.parent video_folder = video_path_obj.parent.name video_file = video_path_obj.stem index_path = case_dir / "L2" / f"{video_folder}.{video_file}.index.json" if not index_path.is_file(): return None with index_path.open("r", encoding="utf-8") as file: payload = json.load(file) return payload if isinstance(payload, dict) else None except Exception: return None def get_video_frame_info(frame_index_payload: Optional[dict[str, Any]], frame_idx: int) -> Optional[dict[str, Any]]: if not frame_index_payload: return None try: fields = frame_index_payload.get("fields", {}) index_list = frame_index_payload.get("index", []) if not isinstance(fields, dict) or not isinstance(index_list, list) or frame_idx >= len(index_list): return None frame_data = index_list[frame_idx] if not isinstance(frame_data, (list, tuple)): return None frame_info = {} for field_name, field_idx in fields.items(): if isinstance(field_idx, int) and 0 <= field_idx < len(frame_data): frame_info[str(field_name)] = frame_data[field_idx] return frame_info except Exception: return None def _normalize_frame_info_token(value: Any) -> str: token = str(value or "").strip() if not token or token.lower() == "none": return "" return token.replace("/", "_").replace("\\", "_").replace(" ", "") def _safe_int(value: Any) -> Optional[int]: try: if value is None: return None return int(str(value).strip()) except (TypeError, ValueError): return None def get_video_frame_id(frame_info: Optional[dict[str, Any]]) -> Optional[int]: if not frame_info: return None for key in ("frame_id", "cve_frame_id", "frameId"): frame_id = _safe_int(frame_info.get(key)) if frame_id is not None: return frame_id return None def iter_video_case_frames( video_path: str | Path, *, frame_index_payload: Optional[dict[str, Any]] = None, frame_stride: int = 1, max_frames: int = 0, frame_index_start: Optional[int] = None, frame_index_end: Optional[int] = None, frame_id_start: Optional[int] = None, frame_id_end: Optional[int] = None, ) -> Iterable[tuple[int, np.ndarray, str, Optional[dict[str, Any]]]]: resolved_video_path = Path(video_path).resolve() cap = cv2.VideoCapture(str(resolved_video_path)) if not cap.isOpened(): raise RuntimeError(f"Failed to open video file: {resolved_video_path}") stride = max(1, int(frame_stride)) resolved_frame_index_start = None if frame_index_start is None else max(0, int(frame_index_start)) resolved_frame_index_end = None if frame_index_end is None else int(frame_index_end) resolved_frame_id_start = None if frame_id_start is None else int(frame_id_start) resolved_frame_id_end = None if frame_id_end is None else int(frame_id_end) read_frame_index = 0 emitted_count = 0 try: while True: ret, frame = cap.read() if not ret: break if resolved_frame_index_end is not None and read_frame_index > resolved_frame_index_end: break frame_info = get_video_frame_info(frame_index_payload, read_frame_index) frame_id_value = get_video_frame_id(frame_info) if resolved_frame_index_start is not None and read_frame_index < resolved_frame_index_start: read_frame_index += 1 continue if resolved_frame_id_start is not None: if frame_id_value is None or frame_id_value < resolved_frame_id_start: read_frame_index += 1 continue if resolved_frame_id_end is not None and frame_id_value is not None and frame_id_value > resolved_frame_id_end: break if read_frame_index % stride != 0: read_frame_index += 1 continue frame_id_token = _normalize_frame_info_token(frame_id_value) timestamp_token = _normalize_frame_info_token(None if frame_info is None else frame_info.get("timestamp")) if frame_id_token and timestamp_token: frame_name = f"{resolved_video_path.stem}_{frame_id_token}_{timestamp_token}.png" elif frame_id_token: frame_name = f"{resolved_video_path.stem}_{frame_id_token}.png" elif timestamp_token: frame_name = f"{resolved_video_path.stem}_{read_frame_index:06d}_{timestamp_token}.png" else: frame_name = f"{resolved_video_path.stem}_{read_frame_index:06d}.png" yield read_frame_index, frame, frame_name, frame_info emitted_count += 1 read_frame_index += 1 if max_frames > 0 and emitted_count >= max_frames: break finally: cap.release()