单目3D初始代码
This commit is contained in:
zhao.zhu
265
tools/pdcl_inference/read_calib_from_clips.py
Executable file
265
tools/pdcl_inference/read_calib_from_clips.py
Executable file
@@ -0,0 +1,265 @@
|
||||
"""
|
||||
Read calibration parameters for each clip in clip_list.txt and save to a CSV table.
|
||||
|
||||
Usage:
|
||||
python tools/pdcl_inference/read_calib_from_clips.py \
|
||||
--clip-list-file tools/pdcl_inference/clip_list.txt \
|
||||
--output calib_table.csv
|
||||
|
||||
Calibration source priority (same as run_batch_merged_infer.py):
|
||||
1. Nearby camera4.json found by resolve_calib_file_for_clip()
|
||||
2. camera4.json embedded as MCAP attachment
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
import sys
|
||||
import traceback
|
||||
from pathlib import Path
|
||||
from typing import List, Optional
|
||||
|
||||
import pandas as pd
|
||||
|
||||
FILE = Path(__file__).resolve()
|
||||
ROOT = FILE.parents[2]
|
||||
if str(ROOT) not in sys.path:
|
||||
sys.path.append(str(ROOT))
|
||||
|
||||
os.environ["STS_UID"] = "dis-uploader"
|
||||
os.environ["STS_SECRET_KEY"] = "277310cc09724d315514a79701fecb0f"
|
||||
|
||||
try:
|
||||
from dotenv import load_dotenv
|
||||
load_dotenv()
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
from tools.pdcl_inference.pdcl_clip_service import PDCLClipService
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Calibration helpers (mirrors run_batch_merged_infer.py logic)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
CALIB_ATTACHMENT_NAMES = (
|
||||
"sigmastar.1/calibs/camera4.json",
|
||||
"test_data/calibs/camera4.json",
|
||||
"calibs/camera4.json",
|
||||
)
|
||||
|
||||
|
||||
def resolve_calib_file_for_clip(clip_path: str) -> Optional[Path]:
|
||||
"""Search for camera4.json in directories near the clip file (same as run_batch_merged_infer.py)."""
|
||||
base = Path(clip_path).resolve()
|
||||
search_dir = base.parent
|
||||
candidates = []
|
||||
for _ in range(4):
|
||||
candidates.extend([
|
||||
search_dir / "calibs" / "camera4.json",
|
||||
search_dir / "test_data" / "calibs" / "camera4.json",
|
||||
search_dir / "L2_calib" / "camera4.json",
|
||||
search_dir / "calib" / "L2_calib" / "camera4.json",
|
||||
])
|
||||
parent = search_dir.parent
|
||||
if parent == search_dir:
|
||||
break
|
||||
search_dir = parent
|
||||
for p in candidates:
|
||||
if p.exists():
|
||||
return p
|
||||
return None
|
||||
|
||||
|
||||
def read_calib_from_mcap_attachment(clip_path: str) -> Optional[dict]:
|
||||
"""Extract camera4.json from MCAP attachments."""
|
||||
try:
|
||||
from pdcl_pyclip.reader import ClipReader
|
||||
except ImportError:
|
||||
return None
|
||||
|
||||
try:
|
||||
reader = ClipReader(clip_path)
|
||||
for attachment in reader.iter_attachments():
|
||||
if attachment.name in CALIB_ATTACHMENT_NAMES:
|
||||
return json.loads(attachment.data.decode("utf-8"))
|
||||
except Exception:
|
||||
pass
|
||||
return None
|
||||
|
||||
|
||||
def get_calib_for_clip(clip_path: str) -> tuple[Optional[dict], str]:
|
||||
"""Return (calib_dict, source_description). Returns (None, error_msg) on failure."""
|
||||
# Priority 1: nearby file
|
||||
nearby = resolve_calib_file_for_clip(clip_path)
|
||||
if nearby is not None:
|
||||
try:
|
||||
with open(nearby) as f:
|
||||
return json.load(f), f"file:{nearby}"
|
||||
except Exception as e:
|
||||
pass # fall through to MCAP attachment
|
||||
|
||||
# Priority 2: MCAP attachment
|
||||
if clip_path.lower().endswith(".mcap"):
|
||||
calib = read_calib_from_mcap_attachment(clip_path)
|
||||
if calib is not None:
|
||||
return calib, "mcap_attachment"
|
||||
|
||||
return None, "not_found"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Clip list parser (same as run_batch_merged_infer.py)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def parse_clip_list(clip_list_file: str) -> List[dict]:
|
||||
tasks = []
|
||||
with open(clip_list_file) as f:
|
||||
for line in f:
|
||||
line = line.strip()
|
||||
if not line:
|
||||
continue
|
||||
parts = line.split()
|
||||
clip_uuid = parts[0]
|
||||
date_name = parts[1] if len(parts) >= 2 else ""
|
||||
tasks.append({"clip_uuid": clip_uuid, "date_name": date_name})
|
||||
return tasks
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Calibration dict → flat record
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def flatten_calib(calib: dict) -> dict:
|
||||
coeffs = calib.get("distort_coeffs", [])
|
||||
pos = calib.get("pos", [])
|
||||
|
||||
fx = calib.get("focal_u")
|
||||
fy = calib.get("focal_v")
|
||||
cx = calib.get("cu")
|
||||
cy = calib.get("cv")
|
||||
pitch = calib.get("pitch") # degrees
|
||||
yaw = calib.get("yaw", 0.0) # degrees
|
||||
|
||||
# Vanishing point (same formula as dataloaders3d.py __getitem__):
|
||||
# vanish_x = cx + fx * tan(yaw_deg * pi/180)
|
||||
# vanish_y = cy - fy * tan(pitch_deg * pi/180)
|
||||
if None not in (fx, fy, cx, cy, pitch):
|
||||
vanish_x = cx + fx * math.tan(yaw * math.pi / 180)
|
||||
vanish_y = cy - fy * math.tan(pitch * math.pi / 180)
|
||||
else:
|
||||
vanish_x = vanish_y = None
|
||||
|
||||
return {
|
||||
"focal_u": fx,
|
||||
"focal_v": fy,
|
||||
"cu": cx,
|
||||
"cv": cy,
|
||||
"pitch": pitch,
|
||||
"yaw": yaw,
|
||||
"roll": calib.get("roll"),
|
||||
"fov": calib.get("fov"),
|
||||
"pos_x": pos[0] if len(pos) > 0 else None,
|
||||
"pos_y": pos[1] if len(pos) > 1 else None,
|
||||
"pos_z": pos[2] if len(pos) > 2 else None,
|
||||
"vanish_x": vanish_x,
|
||||
"vanish_y": vanish_y,
|
||||
"k1": coeffs[0] if len(coeffs) > 0 else None,
|
||||
"k2": coeffs[1] if len(coeffs) > 1 else None,
|
||||
"k3": coeffs[2] if len(coeffs) > 2 else None,
|
||||
"k4": coeffs[3] if len(coeffs) > 3 else None,
|
||||
"img_width": calib.get("image_width", calib.get("img_width")),
|
||||
"img_height": calib.get("image_height", calib.get("img_height")),
|
||||
}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Main
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Extract calibration params from PDCL clips to CSV")
|
||||
parser.add_argument("--clip-list-file", type=str,
|
||||
default=str(FILE.parent / "clip_list.txt"),
|
||||
help="Path to clip list text file")
|
||||
parser.add_argument("--output", type=str, default="calib_table.csv",
|
||||
help="Output CSV path")
|
||||
args = parser.parse_args()
|
||||
|
||||
tasks = parse_clip_list(args.clip_list_file)
|
||||
print(f"Found {len(tasks)} clips in {args.clip_list_file}")
|
||||
|
||||
rows = []
|
||||
for idx, task in enumerate(tasks, start=1):
|
||||
clip_uuid = task["clip_uuid"]
|
||||
date_name = task["date_name"]
|
||||
print(f"[{idx}/{len(tasks)}] clip_uuid={clip_uuid} ...", end=" ", flush=True)
|
||||
|
||||
row = {
|
||||
"clip_uuid": clip_uuid,
|
||||
"date_name": date_name,
|
||||
"clip_path": "",
|
||||
"calib_source": "",
|
||||
"error": "",
|
||||
}
|
||||
|
||||
try:
|
||||
clip_path = PDCLClipService.get_clip_path(clip_uuid)
|
||||
if clip_path is None:
|
||||
row["error"] = "clip_not_found"
|
||||
print("SKIP (clip not found)")
|
||||
rows.append(row)
|
||||
continue
|
||||
|
||||
row["clip_path"] = clip_path
|
||||
calib, source = get_calib_for_clip(clip_path)
|
||||
row["calib_source"] = source
|
||||
|
||||
if calib is None:
|
||||
row["error"] = "calib_not_found"
|
||||
print("SKIP (calib not found)")
|
||||
else:
|
||||
row.update(flatten_calib(calib))
|
||||
print(f"OK ({source})")
|
||||
|
||||
except Exception as exc:
|
||||
row["error"] = f"{type(exc).__name__}: {exc}"
|
||||
print(f"ERROR: {exc}")
|
||||
traceback.print_exc()
|
||||
|
||||
rows.append(row)
|
||||
|
||||
df = pd.DataFrame(rows)
|
||||
|
||||
# Reorder columns for readability
|
||||
lead_cols = ["clip_uuid", "date_name",
|
||||
"focal_u", "focal_v", "cu", "cv",
|
||||
"pitch", "yaw", "roll", "fov",
|
||||
"pos_x", "pos_y", "pos_z",
|
||||
"vanish_x", "vanish_y",
|
||||
"k1", "k2", "k3", "k4",
|
||||
"img_width", "img_height", "calib_source", "clip_path", "error"]
|
||||
present = [c for c in lead_cols if c in df.columns]
|
||||
extra = [c for c in df.columns if c not in present]
|
||||
df = df[present + extra]
|
||||
|
||||
out_path = Path(args.output)
|
||||
out_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
df.to_csv(out_path, index=False)
|
||||
print(f"\nSaved {len(df)} rows → {out_path}")
|
||||
|
||||
# Also save Excel if openpyxl is available
|
||||
try:
|
||||
xlsx_path = out_path.with_suffix(".xlsx")
|
||||
df.to_excel(xlsx_path, index=False)
|
||||
print(f"Saved Excel → {xlsx_path}")
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Print summary table to console
|
||||
print("\n" + df.to_string(index=False))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user