refactor(cli): extract depth postprocess and add tests
- Extract apply_depth_verify_refine_postprocess() from main() for testability - Add test_depth_cli_postprocess.py using mocks to validate JSON and CSV behavior - Keeps CLI behavior unchanged
This commit is contained in:
@@ -1,6 +1,7 @@
|
||||
import click
|
||||
import cv2
|
||||
import json
|
||||
import csv
|
||||
import numpy as np
|
||||
import pyzed.sl as sl
|
||||
from pathlib import Path
|
||||
@@ -21,6 +22,132 @@ from aruco.depth_verify import verify_extrinsics_with_depth
|
||||
from aruco.depth_refine import refine_extrinsics_with_depth
|
||||
|
||||
|
||||
def apply_depth_verify_refine_postprocess(
|
||||
results: Dict[str, Any],
|
||||
verification_frames: Dict[str, Any],
|
||||
marker_geometry: Dict[int, Any],
|
||||
camera_matrices: Dict[str, Any],
|
||||
verify_depth: bool,
|
||||
refine_depth: bool,
|
||||
depth_confidence_threshold: int,
|
||||
report_csv_path: Optional[str] = None,
|
||||
) -> Tuple[Dict[str, Any], List[List[Any]]]:
|
||||
"""
|
||||
Apply depth verification and refinement to computed extrinsics.
|
||||
Returns updated results and list of CSV rows.
|
||||
"""
|
||||
csv_rows: List[List[Any]] = []
|
||||
|
||||
if not (verify_depth or refine_depth):
|
||||
return results, csv_rows
|
||||
|
||||
click.echo("\nRunning depth verification/refinement on computed extrinsics...")
|
||||
|
||||
for serial, vf in verification_frames.items():
|
||||
if str(serial) not in results:
|
||||
continue
|
||||
|
||||
frame = vf["frame"]
|
||||
ids = vf["ids"]
|
||||
|
||||
# Use the FINAL COMPUTED POSE for verification
|
||||
pose_str = results[str(serial)]["pose"]
|
||||
T_mean = np.fromstring(pose_str, sep=" ").reshape(4, 4)
|
||||
cam_matrix = camera_matrices[serial]
|
||||
|
||||
marker_corners_world = {
|
||||
int(mid): marker_geometry[int(mid)]
|
||||
for mid in ids.flatten()
|
||||
if int(mid) in marker_geometry
|
||||
}
|
||||
|
||||
if marker_corners_world and frame.depth_map is not None:
|
||||
verify_res = verify_extrinsics_with_depth(
|
||||
T_mean,
|
||||
marker_corners_world,
|
||||
frame.depth_map,
|
||||
cam_matrix,
|
||||
confidence_map=frame.confidence_map,
|
||||
confidence_thresh=depth_confidence_threshold,
|
||||
)
|
||||
|
||||
results[str(serial)]["depth_verify"] = {
|
||||
"rmse": verify_res.rmse,
|
||||
"mean_abs": verify_res.mean_abs,
|
||||
"median": verify_res.median,
|
||||
"depth_normalized_rmse": verify_res.depth_normalized_rmse,
|
||||
"n_valid": verify_res.n_valid,
|
||||
"n_total": verify_res.n_total,
|
||||
}
|
||||
|
||||
click.echo(
|
||||
f"Camera {serial} verification: RMSE={verify_res.rmse:.3f}m, "
|
||||
f"Valid={verify_res.n_valid}/{verify_res.n_total}"
|
||||
)
|
||||
|
||||
if refine_depth:
|
||||
if verify_res.n_valid < 4:
|
||||
click.echo(
|
||||
f"Camera {serial}: Not enough valid depth points for refinement ({verify_res.n_valid}). Skipping."
|
||||
)
|
||||
else:
|
||||
click.echo(f"Camera {serial}: Refining extrinsics with depth...")
|
||||
T_refined, refine_stats = refine_extrinsics_with_depth(
|
||||
T_mean,
|
||||
marker_corners_world,
|
||||
frame.depth_map,
|
||||
cam_matrix,
|
||||
)
|
||||
|
||||
verify_res_post = verify_extrinsics_with_depth(
|
||||
T_refined,
|
||||
marker_corners_world,
|
||||
frame.depth_map,
|
||||
cam_matrix,
|
||||
confidence_map=frame.confidence_map,
|
||||
confidence_thresh=depth_confidence_threshold,
|
||||
)
|
||||
|
||||
pose_str_refined = " ".join(f"{x:.6f}" for x in T_refined.flatten())
|
||||
results[str(serial)]["pose"] = pose_str_refined
|
||||
results[str(serial)]["refine_depth"] = refine_stats
|
||||
results[str(serial)]["depth_verify_post"] = {
|
||||
"rmse": verify_res_post.rmse,
|
||||
"mean_abs": verify_res_post.mean_abs,
|
||||
"median": verify_res_post.median,
|
||||
"depth_normalized_rmse": verify_res_post.depth_normalized_rmse,
|
||||
"n_valid": verify_res_post.n_valid,
|
||||
"n_total": verify_res_post.n_total,
|
||||
}
|
||||
|
||||
improvement = verify_res.rmse - verify_res_post.rmse
|
||||
results[str(serial)]["refine_depth"]["improvement_rmse"] = (
|
||||
improvement
|
||||
)
|
||||
|
||||
click.echo(
|
||||
f"Camera {serial} refined: RMSE={verify_res_post.rmse:.3f}m "
|
||||
f"(Improved by {improvement:.3f}m). "
|
||||
f"Delta Rot={refine_stats['delta_rotation_deg']:.2f}deg, "
|
||||
f"Trans={refine_stats['delta_translation_norm_m']:.3f}m"
|
||||
)
|
||||
|
||||
verify_res = verify_res_post
|
||||
|
||||
if report_csv_path:
|
||||
for mid, cidx, resid in verify_res.residuals:
|
||||
csv_rows.append([serial, mid, cidx, resid])
|
||||
|
||||
if report_csv_path and csv_rows:
|
||||
with open(report_csv_path, "w", newline="") as f:
|
||||
writer = csv.writer(f)
|
||||
writer.writerow(["serial", "marker_id", "corner_idx", "residual"])
|
||||
writer.writerows(csv_rows)
|
||||
click.echo(f"Saved depth verification report to {report_csv_path}")
|
||||
|
||||
return results, csv_rows
|
||||
|
||||
|
||||
@click.command()
|
||||
@click.option("--svo", "-s", multiple=True, required=False, help="Path to SVO files.")
|
||||
@click.option("--markers", "-m", required=True, help="Path to markers parquet file.")
|
||||
@@ -260,119 +387,16 @@ def main(
|
||||
return
|
||||
|
||||
# 4. Run Depth Verification if requested
|
||||
csv_rows: List[List[Any]] = []
|
||||
if verify_depth or refine_depth:
|
||||
click.echo("\nRunning depth verification/refinement on computed extrinsics...")
|
||||
for serial, acc in accumulators.items():
|
||||
if serial not in verification_frames or str(serial) not in results:
|
||||
continue
|
||||
|
||||
# Retrieve stored frame data
|
||||
vf = verification_frames[serial]
|
||||
frame = vf["frame"]
|
||||
ids = vf["ids"]
|
||||
|
||||
# Use the FINAL COMPUTED POSE for verification
|
||||
pose_str = results[str(serial)]["pose"]
|
||||
T_mean = np.fromstring(pose_str, sep=" ").reshape(4, 4)
|
||||
cam_matrix = camera_matrices[serial]
|
||||
|
||||
marker_corners_world = {
|
||||
int(mid): marker_geometry[int(mid)]
|
||||
for mid in ids.flatten()
|
||||
if int(mid) in marker_geometry
|
||||
}
|
||||
|
||||
if marker_corners_world and frame.depth_map is not None:
|
||||
verify_res = verify_extrinsics_with_depth(
|
||||
T_mean,
|
||||
marker_corners_world,
|
||||
frame.depth_map,
|
||||
cam_matrix,
|
||||
confidence_map=frame.confidence_map,
|
||||
confidence_thresh=depth_confidence_threshold,
|
||||
)
|
||||
|
||||
results[str(serial)]["depth_verify"] = {
|
||||
"rmse": verify_res.rmse,
|
||||
"mean_abs": verify_res.mean_abs,
|
||||
"median": verify_res.median,
|
||||
"depth_normalized_rmse": verify_res.depth_normalized_rmse,
|
||||
"n_valid": verify_res.n_valid,
|
||||
"n_total": verify_res.n_total,
|
||||
}
|
||||
|
||||
click.echo(
|
||||
f"Camera {serial} verification: RMSE={verify_res.rmse:.3f}m, "
|
||||
f"Valid={verify_res.n_valid}/{verify_res.n_total}"
|
||||
)
|
||||
|
||||
if refine_depth:
|
||||
if verify_res.n_valid < 4:
|
||||
click.echo(
|
||||
f"Camera {serial}: Not enough valid depth points for refinement ({verify_res.n_valid}). Skipping."
|
||||
)
|
||||
else:
|
||||
click.echo(
|
||||
f"Camera {serial}: Refining extrinsics with depth..."
|
||||
)
|
||||
T_refined, refine_stats = refine_extrinsics_with_depth(
|
||||
T_mean,
|
||||
marker_corners_world,
|
||||
frame.depth_map,
|
||||
cam_matrix,
|
||||
)
|
||||
|
||||
verify_res_post = verify_extrinsics_with_depth(
|
||||
T_refined,
|
||||
marker_corners_world,
|
||||
frame.depth_map,
|
||||
cam_matrix,
|
||||
confidence_map=frame.confidence_map,
|
||||
confidence_thresh=depth_confidence_threshold,
|
||||
)
|
||||
|
||||
pose_str_refined = " ".join(
|
||||
f"{x:.6f}" for x in T_refined.flatten()
|
||||
)
|
||||
results[str(serial)]["pose"] = pose_str_refined
|
||||
results[str(serial)]["refine_depth"] = refine_stats
|
||||
results[str(serial)]["depth_verify_post"] = {
|
||||
"rmse": verify_res_post.rmse,
|
||||
"mean_abs": verify_res_post.mean_abs,
|
||||
"median": verify_res_post.median,
|
||||
"depth_normalized_rmse": verify_res_post.depth_normalized_rmse,
|
||||
"n_valid": verify_res_post.n_valid,
|
||||
"n_total": verify_res_post.n_total,
|
||||
}
|
||||
|
||||
improvement = verify_res.rmse - verify_res_post.rmse
|
||||
results[str(serial)]["refine_depth"]["improvement_rmse"] = (
|
||||
improvement
|
||||
)
|
||||
|
||||
click.echo(
|
||||
f"Camera {serial} refined: RMSE={verify_res_post.rmse:.3f}m "
|
||||
f"(Improved by {improvement:.3f}m). "
|
||||
f"Delta Rot={refine_stats['delta_rotation_deg']:.2f}deg, "
|
||||
f"Trans={refine_stats['delta_translation_norm_m']:.3f}m"
|
||||
)
|
||||
|
||||
verify_res = verify_res_post
|
||||
|
||||
if report_csv:
|
||||
for mid, cidx, resid in verify_res.residuals:
|
||||
csv_rows.append([serial, mid, cidx, resid])
|
||||
|
||||
# 5. Save CSV Report
|
||||
if report_csv and csv_rows:
|
||||
import csv
|
||||
|
||||
with open(report_csv, "w", newline="") as f:
|
||||
writer = csv.writer(f)
|
||||
writer.writerow(["serial", "marker_id", "corner_idx", "residual"])
|
||||
writer.writerows(csv_rows)
|
||||
click.echo(f"Saved depth verification report to {report_csv}")
|
||||
apply_depth_verify_refine_postprocess(
|
||||
results,
|
||||
verification_frames,
|
||||
marker_geometry,
|
||||
camera_matrices,
|
||||
verify_depth,
|
||||
refine_depth,
|
||||
depth_confidence_threshold,
|
||||
report_csv,
|
||||
)
|
||||
|
||||
# 6. Save to JSON
|
||||
with open(output, "w") as f:
|
||||
|
||||
Reference in New Issue
Block a user