Integrate ICP refinement into refine_ground_plane.py CLI

py_workspace/refine_ground_plane.py

@@ -14,6 +14,7 @@ from aruco.ground_plane import (
     Mat44,
 )
 from aruco.depth_save import load_depth_data
+from aruco.icp_registration import refine_with_icp, ICPConfig
 
 
 @click.command()
@@ -85,6 +86,23 @@ from aruco.depth_save import load_depth_data
     type=int,
     help="Random seed for RANSAC determinism.",
 )
+@click.option(
+    "--icp/--no-icp",
+    default=False,
+    help="Enable ICP refinement after ground plane alignment.",
+)
+@click.option(
+    "--icp-method",
+    type=click.Choice(["point_to_plane", "gicp"]),
+    default="point_to_plane",
+    help="ICP registration method.",
+)
+@click.option(
+    "--icp-voxel-size",
+    type=float,
+    default=0.02,
+    help="Voxel size for ICP downsampling (meters).",
+)
 @click.option(
     "--debug/--no-debug",
     default=False,
@@ -103,6 +121,9 @@ def main(
     height_range: tuple[float, float],
     stride: int,
     seed: Optional[int],
+    icp: bool,
+    icp_method: str,
+    icp_voxel_size: float,
     debug: bool,
 ):
     """
@@ -187,6 +208,30 @@ def main(
             logger.info(f"Max rotation: {metrics.rotation_deg:.2f} deg")
             logger.info(f"Max translation: {metrics.translation_m:.3f} m")
 
+        # 4.5 Optional ICP Refinement
+        icp_metrics = None
+        if icp:
+            logger.info(f"Running ICP refinement ({icp_method})...")
+            icp_config = ICPConfig(
+                method=icp_method,
+                voxel_size=icp_voxel_size,
+            )
+
+            icp_extrinsics, icp_metrics = refine_with_icp(
+                camera_data_for_refine,
+                new_extrinsics,
+                metrics.camera_planes,
+                icp_config,
+            )
+
+            if icp_metrics.success:
+                logger.info(f"ICP refinement successful: {icp_metrics.message}")
+                new_extrinsics = icp_extrinsics
+            else:
+                logger.warning(
+                    f"ICP refinement failed or skipped: {icp_metrics.message}"
+                )
+
         # 5. Save Output Extrinsics
         output_data = extrinsics_data.copy()
 
@@ -237,6 +282,23 @@ def main(
             "per_camera": per_camera_diagnostics,
         }
 
+        if icp_metrics:
+            output_data["_meta"]["icp_refined"] = {
+                "timestamp": str(np.datetime64("now")),
+                "config": {
+                    "method": icp_method,
+                    "voxel_size": icp_voxel_size,
+                },
+                "metrics": {
+                    "success": icp_metrics.success,
+                    "num_pairs_attempted": icp_metrics.num_pairs_attempted,
+                    "num_pairs_converged": icp_metrics.num_pairs_converged,
+                    "num_cameras_optimized": icp_metrics.num_cameras_optimized,
+                    "num_disconnected": icp_metrics.num_disconnected,
+                    "message": icp_metrics.message,
+                },
+            }
+
         logger.info(f"Saving refined extrinsics to {output_extrinsics}")
         with open(output_extrinsics, "w") as f:
             json.dump(output_data, f, indent=4, sort_keys=True)