Buckets:
| import json | |
| from dataclasses import dataclass | |
| from pathlib import Path | |
| import cv2 | |
| import numpy as np | |
| def sample_video(path, stride=15, max_frames=16, start=0): | |
| if stride < 1 or max_frames < 1 or start < 0: | |
| raise ValueError("stride/max_frames must be positive and start nonnegative") | |
| capture = cv2.VideoCapture(str(path)) | |
| try: | |
| if not capture.isOpened(): | |
| raise ValueError(f"Cannot open video: {path}") | |
| fps = capture.get(cv2.CAP_PROP_FPS) | |
| if not np.isfinite(fps) or fps <= 0: | |
| raise ValueError("Video must have a valid FPS") | |
| frames, indices = [], [] | |
| index = 0 | |
| while len(frames) < max_frames: | |
| ok = capture.grab() | |
| if not ok: | |
| break | |
| if index >= start and (index - start) % stride == 0: | |
| ok, frame = capture.retrieve() | |
| if not ok: | |
| raise ValueError(f"Cannot decode frame {index}") | |
| frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)) | |
| indices.append(index) | |
| index += 1 | |
| if not frames: | |
| raise ValueError("No frames selected from video") | |
| return frames, np.asarray(indices), np.asarray(indices) / fps | |
| finally: | |
| capture.release() | |
| def _numpy(value): | |
| if hasattr(value, "detach"): | |
| value = value.detach().cpu().numpy() | |
| return np.asarray(value) | |
| class Reconstruction: | |
| views: list | |
| frame_indices: np.ndarray | |
| timestamps: np.ndarray | |
| model_id: str | |
| def from_predictions(cls, predictions, frame_indices, timestamps, model_id): | |
| if len(predictions) != len(frame_indices) or len(timestamps) != len( | |
| frame_indices | |
| ): | |
| raise ValueError("Predictions, frame indices and timestamps must match") | |
| views = [] | |
| for pred in predictions: | |
| view = { | |
| key: _numpy(pred[key])[0] | |
| for key in ( | |
| "pts3d", | |
| "depth_z", | |
| "intrinsics", | |
| "camera_poses", | |
| "conf", | |
| "mask", | |
| "img_no_norm", | |
| ) | |
| } | |
| points = view["pts3d"] | |
| height, width, channels = points.shape | |
| if ( | |
| channels != 3 | |
| or view["camera_poses"].shape != (4, 4) | |
| or view["intrinsics"].shape != (3, 3) | |
| ): | |
| raise ValueError("Unexpected MapAnything geometry/camera shape") | |
| valid = view["mask"].reshape(height, width).astype(bool) | |
| valid &= np.isfinite(points).all(axis=-1) | |
| valid &= np.isfinite(view["conf"].reshape(height, width)) | |
| valid &= np.isfinite(view["depth_z"].reshape(height, width)) | |
| valid &= view["depth_z"].reshape(height, width) > 0 | |
| view["mask"] = valid | |
| view["img_no_norm"] = np.clip(view["img_no_norm"] * 255, 0, 255).astype( | |
| np.uint8 | |
| ) | |
| views.append(view) | |
| return cls(views, np.asarray(frame_indices), np.asarray(timestamps), model_id) | |
| def save(self, folder, point_stride=4): | |
| if point_stride < 1: | |
| raise ValueError("point_stride must be positive") | |
| folder = Path(folder) | |
| folder.mkdir(parents=True, exist_ok=True) | |
| if not self.views: | |
| raise ValueError("Cannot save an empty reconstruction") | |
| cloud, colors = [], [] | |
| for i, view in enumerate(self.views): | |
| np.savez_compressed(folder / f"view_{i:04d}.npz", **view) | |
| cv2.imwrite( | |
| str(folder / f"view_{i:04d}.png"), | |
| cv2.cvtColor(view["img_no_norm"], cv2.COLOR_RGB2BGR), | |
| ) | |
| valid = view["mask"][::point_stride, ::point_stride] | |
| cloud.append(view["pts3d"][::point_stride, ::point_stride][valid]) | |
| colors.append(view["img_no_norm"][::point_stride, ::point_stride][valid]) | |
| np.savez_compressed( | |
| folder / "cameras.npz", | |
| frame_indices=self.frame_indices, | |
| timestamps=self.timestamps, | |
| **{ | |
| "camera/K": np.stack([v["intrinsics"] for v in self.views]), | |
| "camera/T_world_cam": np.stack([v["camera_poses"] for v in self.views]), | |
| "image_size": np.asarray([v["pts3d"].shape[1::-1] for v in self.views]), | |
| }, | |
| ) | |
| points, rgb = np.concatenate(cloud), np.concatenate(colors) | |
| with (folder / "scene.ply").open("w") as f: | |
| f.write(f"ply\nformat ascii 1.0\nelement vertex {len(points)}\n") | |
| f.write("property float x\nproperty float y\nproperty float z\n") | |
| f.write( | |
| "property uchar red\nproperty uchar green\nproperty uchar blue\nend_header\n" | |
| ) | |
| np.savetxt(f, np.column_stack((points, rgb)), fmt="%.6f %.6f %.6f %d %d %d") | |
| (folder / "manifest.json").write_text( | |
| json.dumps( | |
| { | |
| "model_id": self.model_id, | |
| "units": "meters", | |
| "world_frame": "mapanything_reference_view", | |
| "camera_convention": "opencv_cam2world", | |
| "pixel_coordinates": "preprocessed_images", | |
| "views": len(self.views), | |
| "point_count": len(points), | |
| "frame_indices": self.frame_indices.tolist(), | |
| "timestamps": self.timestamps.tolist(), | |
| }, | |
| indent=2, | |
| ) | |
| + "\n" | |
| ) | |
| class MapAnythingReconstructor: | |
| def __init__(self, model_id="facebook/map-anything-apache", device="cuda"): | |
| try: | |
| import torch | |
| from mapanything.models import MapAnything | |
| except ImportError as exc: | |
| raise ImportError( | |
| "Install the MapAnything optional dependency; see README.md" | |
| ) from exc | |
| if device.startswith("cuda") and not torch.cuda.is_available(): | |
| raise RuntimeError("CUDA unavailable; choose device='cpu' explicitly") | |
| self.model_id, self.device = model_id, device | |
| self.model = MapAnything.from_pretrained(model_id).to(device).eval() | |
| def reconstruct( | |
| self, frames, frame_indices, timestamps, intrinsics=None, minibatch_size=1 | |
| ): | |
| import torch | |
| from mapanything.utils.image import preprocess_inputs | |
| if ( | |
| not len(frames) | |
| or len(frames) != len(frame_indices) | |
| or len(frames) != len(timestamps) | |
| ): | |
| raise ValueError("Nonempty frames, indices and timestamps must match") | |
| if minibatch_size < 1: | |
| raise ValueError("minibatch_size must be positive") | |
| inputs = [{"img": frame} for frame in frames] | |
| if intrinsics is not None: | |
| k = np.asarray(intrinsics, dtype=np.float32) | |
| if ( | |
| k.shape != (3, 3) | |
| or not np.isfinite(k).all() | |
| or k[0, 0] <= 0 | |
| or k[1, 1] <= 0 | |
| ): | |
| raise ValueError( | |
| "intrinsics must be a finite 3x3 calibration with positive focal lengths" | |
| ) | |
| for view in inputs: | |
| view["intrinsics"] = k.copy() | |
| views = preprocess_inputs(inputs) | |
| with torch.inference_mode(): | |
| predictions = self.model.infer( | |
| views, | |
| memory_efficient_inference=True, | |
| minibatch_size=minibatch_size, | |
| use_amp=self.device.startswith("cuda"), | |
| amp_dtype="bf16", | |
| apply_mask=True, | |
| mask_edges=True, | |
| apply_confidence_mask=True, | |
| confidence_percentile=10, | |
| ) | |
| return Reconstruction.from_predictions( | |
| predictions, frame_indices, timestamps, self.model_id | |
| ) | |
Xet Storage Details
- Size:
- 7.99 kB
- Xet hash:
- 63af467712b1e5fece6b899e92f23de68e127a0be0f4b35ec3039036f49b3e28
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.