Buckets:
| import json | |
| import os | |
| import sys | |
| from pathlib import Path | |
| import cv2 | |
| import numpy as np | |
| from scipy.spatial.transform import Rotation | |
| from .plane import plane_pointmap | |
| def reconstruct_object( | |
| folder, output, source_root, checkpoints, frame=10, seed=42, device="cuda:1" | |
| ): | |
| import torch | |
| from hydra.utils import instantiate | |
| from omegaconf import OmegaConf | |
| folder, output = Path(folder), Path(output) | |
| output.mkdir(parents=True, exist_ok=True) | |
| os.environ["LIDRA_SKIP_INIT"] = "true" | |
| sys.path.insert(0, str(Path(source_root).resolve())) | |
| torch.cuda.set_device(device) | |
| scene = json.loads((folder / "scene.json").read_text()) | |
| with np.load(folder / "cameras.npz") as saved: | |
| frame_indices = saved["frame_indices"] | |
| view_index = int(np.argmin(np.abs(frame_indices - frame))) | |
| frame = int(frame_indices[view_index]) | |
| camera = saved["camera/T_world_cam"][view_index] | |
| intrinsics = saved["camera/K"] | |
| capture = cv2.VideoCapture(str(folder / "real_video.mp4")) | |
| try: | |
| capture.set(cv2.CAP_PROP_POS_FRAMES, view_index) | |
| ok, image = capture.read() | |
| if not ok: | |
| raise ValueError(f"Cannot read object source frame {frame}") | |
| image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) | |
| finally: | |
| capture.release() | |
| with np.load(folder / "real_masks.npz") as saved: | |
| if str(saved["model"]) != "sam3.1_multiplex": | |
| raise ValueError("Object reconstruction requires SAM3.1 masks") | |
| actor = saved["names"].tolist().index("object") | |
| index = int(np.flatnonzero(saved["frame_indices"] == frame)[0]) | |
| mask = saved["masks"][actor, index] | |
| if mask.shape != image.shape[:2]: | |
| raise ValueError("Object mask and source video dimensions must match") | |
| if np.count_nonzero(mask) < 16: | |
| raise ValueError("Object mask is empty") | |
| pointmap = plane_pointmap( | |
| camera, | |
| intrinsics, | |
| image.shape[1::-1], | |
| scene["table_height_m"], | |
| mask, | |
| scene["object_height_prior_m"], | |
| ) | |
| if np.count_nonzero(np.isfinite(pointmap).all(axis=-1) & mask) < 16: | |
| raise ValueError( | |
| "Object mask has no valid forward intersection with the table prior" | |
| ) | |
| pointmap *= np.array([-1, -1, 1], dtype=np.float32) | |
| mask = mask.astype(np.uint8) | |
| rgba = np.concatenate((image, (mask * 255)[..., None]), axis=-1) | |
| cv2.imwrite(str(output / "input.png"), cv2.cvtColor(rgba, cv2.COLOR_RGBA2BGRA)) | |
| config = OmegaConf.load(Path(checkpoints) / "pipeline.yaml") | |
| config.workspace_dir = str(Path(checkpoints).resolve()) | |
| config.compile_model = False | |
| config.rendering_engine = "pytorch3d" | |
| config.device = device | |
| config.depth_model = None | |
| print( | |
| json.dumps({"stage": "loading_sam3d", "frame": frame, "device": device}), | |
| flush=True, | |
| ) | |
| pipeline = instantiate(config) | |
| with torch.inference_mode(): | |
| result = pipeline.run( | |
| rgba, | |
| seed=seed, | |
| with_mesh_postprocess=False, | |
| with_texture_baking=False, | |
| with_layout_postprocess=False, | |
| use_vertex_color=True, | |
| pointmap=torch.from_numpy(pointmap.astype(np.float32)), | |
| ) | |
| mesh = result["glb"] | |
| mesh.export(output / "object.glb") | |
| vertices = ( | |
| np.asarray(mesh.vertices) @ np.array([[1, 0, 0], [0, 0, -1], [0, 1, 0]]).T | |
| ) | |
| rotation = result["rotation"].detach().cpu().numpy().reshape(-1, 4)[0] | |
| scale = result["scale"].detach().cpu().numpy().reshape(-1, 3)[0] | |
| translation = result["translation"].detach().cpu().numpy().reshape(-1, 3)[0] | |
| matrix = Rotation.from_quat(rotation[[1, 2, 3, 0]]).as_matrix() | |
| camera_vertices = (vertices * scale) @ matrix + translation | |
| camera_vertices *= np.array([-1, -1, 1]) | |
| world_vertices = camera_vertices @ camera[:3, :3].T + camera[:3, 3] | |
| colors = np.asarray(mesh.visual.vertex_colors)[:, :3].astype(np.uint8) | |
| faces = np.asarray(mesh.faces, dtype=np.int32) | |
| np.savez_compressed( | |
| output / "object.npz", | |
| world_vertices=world_vertices, | |
| faces=faces, | |
| colors=colors, | |
| source_frame_index=frame, | |
| rotation=rotation, | |
| scale=scale, | |
| translation=translation, | |
| T_world_camera=camera, | |
| intrinsics=intrinsics, | |
| ) | |
| import trimesh | |
| world_mesh = trimesh.Trimesh( | |
| world_vertices, faces, vertex_colors=colors, process=False | |
| ) | |
| world_mesh.export(output / "object_world.ply") | |
| metrics = { | |
| "model": "facebook/sam-3d-objects", | |
| "frame": frame, | |
| "seed": seed, | |
| "vertices": len(vertices), | |
| "triangles": len(faces), | |
| "units": "meters", | |
| "extent_m": np.ptp(world_vertices, axis=0).tolist(), | |
| "center_world_m": world_vertices.mean(axis=0).tolist(), | |
| "watertight": bool(world_mesh.is_watertight), | |
| "layout": "sam3d_pose_with_calibrated_table_prior", | |
| "pointmap_source": "analytic_camera_ray_plane_intersection", | |
| "table_height_m": scene["table_height_m"], | |
| "object_height_prior_m": scene["object_height_prior_m"], | |
| } | |
| (output / "metrics.json").write_text(json.dumps(metrics, indent=2) + "\n") | |
| print(json.dumps(metrics), flush=True) | |
| return output / "object.npz" | |
Xet Storage Details
- Size:
- 5.31 kB
- Xet hash:
- b4184a02ca278c5c4330a6942d8e73551c606d377cabf990984aa00576f5ce4c
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.