Buckets:
| import json | |
| import os | |
| import sys | |
| from pathlib import Path | |
| import cv2 | |
| import numpy as np | |
| from scipy.spatial.transform import Rotation | |
| 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) | |
| 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]) | |
| with np.load(folder / f"view_{view_index:04d}.npz") as saved: | |
| image = saved["img_no_norm"] | |
| points = saved["pts3d"] | |
| camera = saved["camera_poses"] | |
| intrinsics = saved["intrinsics"] | |
| valid = saved["mask"] | |
| 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] | |
| mask = cv2.resize( | |
| mask.astype(np.uint8), image.shape[1::-1], interpolation=cv2.INTER_NEAREST | |
| ) | |
| if np.count_nonzero(mask) < 16: | |
| raise ValueError("Object mask is empty") | |
| pointmap = (points - camera[:3, 3]) @ camera[:3, :3] | |
| pointmap *= np.array([-1, -1, 1], dtype=np.float32) | |
| pointmap[~valid] = np.nan | |
| 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]) | |
| reconstruction_vertices = camera_vertices @ camera[:3, :3].T + camera[:3, 3] | |
| with np.load(folder / "surface.npz") as saved: | |
| anchor = saved["T_world_reconstruction"] | |
| world_vertices = reconstruction_vertices @ anchor[:3, :3].T + anchor[: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=anchor @ 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_mapanything_metric_pointmap", | |
| } | |
| (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:
- 4.69 kB
- Xet hash:
- 837a038f0f4dfcc85eacd1103ab7a1725b5bf56bc5de31101a5170ff0e54ddb3
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.