Spaces:
Running on Zero
Running on Zero
| """BA-T demo: two-view bundle adjustment from an image pair. | |
| Upload two images (A, B). BA-T — an iterative transformer layer inspired by | |
| bundle adjustment — refines the pair and predicts: | |
| * a dense, coloured 3D point cloud (both views, in camera A's frame), | |
| * the relative camera pose between the two images. | |
| The model runs entirely on the GPU via ZeroGPU. All outputs share one | |
| arbitrary (non-metric) scale within a pair. | |
| """ | |
| import os | |
| import tempfile | |
| import spaces # MUST come before any torch / CUDA-touching import | |
| import numpy as np | |
| import torch | |
| import gradio as gr | |
| from ba_t.data import load_views, denormalize | |
| from ba_t.inference import load_model, run_pairs | |
| MODEL_ID = "zhangganlin/BA-T" | |
| CKPT_FILENAME = "bat.pth" | |
| DEFAULT_K = 8 | |
| DEFAULT_CONF_THRESH = 3.0 | |
| # ---------------------------------------------------------------- model load | |
| ckpt_path = os.path.join(tempfile.gettempdir(), CKPT_FILENAME) | |
| if not os.path.exists(ckpt_path): | |
| from huggingface_hub import hf_hub_download | |
| ckpt_path = hf_hub_download(MODEL_ID, CKPT_FILENAME) | |
| model = load_model(ckpt_path, device="cuda", verbose=False) | |
| print(f"BA-T loaded from {MODEL_ID}/{CKPT_FILENAME}", flush=True) | |
| # ---------------------------------------------------------------- helpers | |
| def rotation_angle_deg(R): | |
| return float(np.degrees(np.arccos(np.clip((np.trace(R) - 1) / 2, -1.0, 1.0)))) | |
| def transform_points(T, pts): | |
| return pts @ T[:3, :3].T + T[:3, 3] | |
| def save_ply(path, pts, colors): | |
| """Binary little-endian PLY with xyz float32 + rgb uint8.""" | |
| pts = np.asarray(pts, dtype=np.float32).reshape(-1, 3) | |
| colors = np.asarray(colors, dtype=np.uint8).reshape(-1, 3) | |
| vertex = np.empty( | |
| len(pts), | |
| dtype=[ | |
| ("x", "<f4"), ("y", "<f4"), ("z", "<f4"), | |
| ("red", "u1"), ("green", "u1"), ("blue", "u1"), | |
| ], | |
| ) | |
| vertex["x"], vertex["y"], vertex["z"] = pts[:, 0], pts[:, 1], pts[:, 2] | |
| vertex["red"], vertex["green"], vertex["blue"] = colors[:, 0], colors[:, 1], colors[:, 2] | |
| header = ( | |
| "ply\nformat binary_little_endian 1.0\n" | |
| f"element vertex {len(pts)}\n" | |
| "property float x\nproperty float y\nproperty float z\n" | |
| "property uchar red\nproperty uchar green\nproperty uchar blue\nend_header\n" | |
| ) | |
| with open(path, "wb") as f: | |
| f.write(header.encode("ascii")) | |
| f.write(vertex.tobytes()) | |
| def pose_matrix_str(T): | |
| return "\n".join(" ".join(f"{v:+9.4f}" for v in row) for row in T) | |
| # ---------------------------------------------------------------- inference | |
| def bundle_adjust( | |
| image_a, | |
| image_b, | |
| num_refine_iter: int = DEFAULT_K, | |
| conf_thresh: float = DEFAULT_CONF_THRESH, | |
| ): | |
| """Refine an image pair with BA-T and reconstruct the 3D geometry. | |
| Args: | |
| image_a: first image (camera A). | |
| image_b: second image (camera B, taken from a different viewpoint). | |
| num_refine_iter: number of BA-layer refinement iterations K (4-24). | |
| conf_thresh: keep 3D points whose confidence is above this value. | |
| Returns: | |
| (ply file of the coloured point cloud in camera A's frame, | |
| npz file with point maps / confidences / relative poses, | |
| pose text report) | |
| """ | |
| if image_a is None or image_b is None: | |
| raise gr.Error("Please provide both images A and B.") | |
| view_a, view_b = load_views([image_a, image_b]) | |
| out = run_pairs(model, view_a, [view_b], num_refine_iter=int(num_refine_iter)) | |
| a, b = out["a"][0], out["b"][0] | |
| T_ba = a["relative_pose"][0].numpy() # camera A -> camera B | |
| T_ab = b["relative_pose"][0].numpy() # camera B -> camera A | |
| pts_a = a["pts3d"][0].numpy() | |
| pts_b = b["pts3d"][0].numpy() | |
| conf_a = a["conf"][0].numpy() | |
| conf_b = b["conf"][0].numpy() | |
| img_a = denormalize(view_a["img"][0]) | |
| img_b = denormalize(view_b["img"][0]) | |
| # fused point cloud, both views in the camera frame of A | |
| keep_a = conf_a > conf_thresh | |
| keep_b = conf_b > conf_thresh | |
| if not keep_a.any() and not keep_b.any(): | |
| raise gr.Error( | |
| "No 3D point passed the confidence threshold. " | |
| "Try lowering it, or use two images that see the same scene." | |
| ) | |
| pts = np.concatenate([pts_a[keep_a], transform_points(T_ab, pts_b[keep_b])]) | |
| cols = np.concatenate([img_a[keep_a], img_b[keep_b]]) | |
| with tempfile.NamedTemporaryFile(suffix=".ply", delete=False) as f: | |
| ply_path = f.name | |
| save_ply(ply_path, pts, cols) | |
| with tempfile.NamedTemporaryFile(suffix=".npz", delete=False) as f: | |
| npz_path = f.name | |
| np.savez_compressed( | |
| npz_path, | |
| pts3d_a=pts_a, pts3d_b=pts_b, conf_a=conf_a, conf_b=conf_b, | |
| T_b_from_a=T_ba, T_a_from_b=T_ab, img_a=img_a, img_b=img_b, | |
| ) | |
| # camera centre of B, as a unit direction in camera A's frame (scale is arbitrary) | |
| t = T_ab[:3, 3] / max(np.linalg.norm(T_ab[:3, 3]), 1e-12) | |
| report = ( | |
| f"Relative rotation (A -> B): {rotation_angle_deg(T_ba[:3, :3]):.2f} deg\n" | |
| f"Pose confidence (A->B / B->A): {float(a['relative_pose_conf'][0]):.3f} / " | |
| f"{float(b['relative_pose_conf'][0]):.3f}\n" | |
| f"Camera B centre direction (in A's frame): " | |
| f"[{t[0]:+.3f}, {t[1]:+.3f}, {t[2]:+.3f}] (unit; scale is arbitrary)\n" | |
| f"Points kept: {len(pts)} of {pts_a.size // 3 + pts_b.size // 3}\n\n" | |
| f"T_B<-A (camera A -> camera B):\n{pose_matrix_str(T_ba)}\n\n" | |
| f"T_A<-B (camera B -> camera A):\n{pose_matrix_str(T_ab)}" | |
| ) | |
| return ply_path, npz_path, report | |
| # ---------------------------------------------------------------- ui | |
| CSS = """ | |
| #col-container { max-width: 1100px; margin: 0 auto; } | |
| .dark .gradio-container { color: var(--body-text-color); } | |
| """ | |
| with gr.Blocks() as demo: | |
| gr.Markdown( | |
| """ | |
| # BA-T: An Iterative Transformer for Two-View Bundle Adjustment | |
| Upload two images of a scene from different viewpoints. BA-T refines the | |
| pair with a recursive bundle-adjustment transformer layer and predicts a | |
| dense 3D point cloud and the relative camera pose between the views. | |
| *Point maps and translation share one arbitrary (non-metric) scale.* | |
| [Paper](https://arxiv.org/abs/2606.03287) · | |
| [Project page](https://ganlinzhang.xyz/BA-T/) · | |
| [Code](https://github.com/zhangganlin/BA-T) · | |
| [Model](https://huggingface.co/zhangganlin/BA-T) | |
| """ | |
| ) | |
| with gr.Column(elem_id="col-container"): | |
| with gr.Row(): | |
| image_a = gr.Image(type="filepath", label="Image A", height=280) | |
| image_b = gr.Image(type="filepath", label="Image B", height=280) | |
| with gr.Accordion("Advanced settings", open=False): | |
| num_refine_iter = gr.Slider( | |
| minimum=4, maximum=24, step=1, value=DEFAULT_K, | |
| label="Refinement iterations K", | |
| info="The BA layer runs K times; trained with 4-24. More = more refined.", | |
| ) | |
| conf_thresh = gr.Slider( | |
| minimum=1.0, maximum=10.0, step=0.5, value=DEFAULT_CONF_THRESH, | |
| label="Point confidence threshold", | |
| info="Only 3D points above this confidence enter the point cloud.", | |
| ) | |
| run_btn = gr.Button("Reconstruct", variant="primary") | |
| point_cloud = gr.Model3D( | |
| label="3D point cloud (camera A's frame)", | |
| camera_position=(0, 0, -3.5), | |
| ) | |
| with gr.Row(): | |
| pose_report = gr.Textbox(label="Relative camera pose", lines=13) | |
| npz_file = gr.File(label="Raw predictions (.npz)") | |
| gr.Examples( | |
| examples=[ | |
| ["examples/teaser_0.jpg", "examples/teaser_1.jpg"], | |
| ["examples/twoview_desk_0.jpg", "examples/twoview_desk_1.jpg"], | |
| ["examples/twoview_shelf_0.jpg", "examples/twoview_shelf_1.jpg"], | |
| ["examples/twoview_corner_0.jpg", "examples/twoview_corner_1.jpg"], | |
| ], | |
| inputs=[image_a, image_b], | |
| outputs=[point_cloud, npz_file, pose_report], | |
| fn=bundle_adjust, | |
| cache_examples=True, | |
| cache_mode="lazy", | |
| ) | |
| run_btn.click( | |
| bundle_adjust, | |
| inputs=[image_a, image_b, num_refine_iter, conf_thresh], | |
| outputs=[point_cloud, npz_file, pose_report], | |
| api_name="reconstruct", | |
| ) | |
| demo.launch(mcp_server=True, theme=gr.themes.Citrus(), css=CSS) | |