248 lines
9.0 KiB
Python
Executable File
248 lines
9.0 KiB
Python
Executable File
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()
|