Spaces:
Running on Zero
Running on Zero
| """Camera rotations of a moving camera: SimpleVO (default) or DPVO (optional). | |
| Both return world-to-camera rotations (L, 3, 3), the model's camera-motion input | |
| (see exmr.geometry.camera.cam_angvel_from_rotations). | |
| """ | |
| import contextlib | |
| import ctypes | |
| import ctypes.util | |
| import queue | |
| import threading | |
| from pathlib import Path | |
| import cv2 | |
| import imageio.v3 as iio | |
| import numpy as np | |
| import pycolmap | |
| import torch | |
| from tqdm import tqdm | |
| from exmr.geometry.rotation_conversions import quaternion_to_matrix | |
| from exmr.preprocessing.video import video_lwh | |
| DPVO_ROOT = Path(__file__).resolve().parents[2] / "third_party" / "DPVO" | |
| class SimpleVO: | |
| """SimpleVO (from GVHMR): relative camera poses from SIFT matches between every step-th | |
| frame (pycolmap two-view geometry), interpolated to all frames. | |
| Args: | |
| scale: frames are downscaled by this factor for matching | |
| step: frame stride between matched frames | |
| f_mm: assumed focal length in mm of a full-frame (36x24 mm) camera | |
| progress: show a progress bar | |
| """ | |
| def __init__(self, scale=0.5, step=8, f_mm=24, progress=True): | |
| self.scale, self.step, self.f_mm, self.progress = scale, step, f_mm, progress | |
| self.sift = cv2.SIFT_create() | |
| def __call__(self, video_path): | |
| """(L, 3, 3) world-to-camera rotations (float64), the first one identity.""" | |
| with _opencv_multithreading(): | |
| return torch.from_numpy(self.poses(video_path)[:, :3, :3]) | |
| def poses(self, video_path): | |
| """(L, 4, 4) world-to-camera transforms (float64), the first one identity.""" | |
| _reset_random_state() | |
| # keep only the frames that are matched, every step-th and the last one: all frames of a | |
| # long video would not fit into memory | |
| frames = {} # frame index -> (H, W, 3) RGB | |
| reader = iio.imiter( | |
| video_path, | |
| plugin="pyav", | |
| thread_type="FRAME", | |
| filter_sequence=[("scale", f"iw*{self.scale}:ih*{self.scale}")], | |
| ) | |
| for i, frame in enumerate(reader): | |
| if i % self.step == 0: | |
| frames[i] = frame | |
| frames[i] = frame # the last frame | |
| sample_idxs, frames = list(frames), list(frames.values()) | |
| H, W = frames[0].shape[:2] | |
| # pinhole camera with the principal point at the center | |
| focal_length = (W**2 + H**2) ** 0.5 / (24**2 + 36**2) ** 0.5 * self.f_mm | |
| camera = pycolmap.Camera( | |
| camera_id=0, model="SIMPLE_PINHOLE", width=W, height=H, | |
| params=[focal_length, W / 2, H / 2], | |
| ) # fmt: skip | |
| options = pycolmap.TwoViewGeometryOptions( | |
| min_num_inliers=10, | |
| min_E_F_inlier_ratio=0.8, | |
| max_H_inlier_ratio=0.9, | |
| compute_relative_pose=True, | |
| ) | |
| T_w2c = [np.eye(4)] # T_w2c @ p_w = p_c | |
| features = self._features(frames[0]) | |
| for frame in tqdm(frames[1:], desc="Visual odometry", disable=not self.progress): | |
| prev, features = features, self._features(frame) | |
| pts0, pts1 = _match(prev, features) | |
| matches = np.stack([np.arange(len(pts0)), np.arange(len(pts0))], axis=-1) | |
| answer = pycolmap.estimate_calibrated_two_view_geometry( | |
| camera, pts0.astype(np.float64), camera, pts1.astype(np.float64), | |
| matches=matches, options=options, | |
| ) # fmt: skip | |
| T_delta = np.eye(4) # previous -> current camera | |
| T_delta[:3] = answer.cam2_from_cam1.matrix().astype(np.float32) | |
| T_w2c.append(T_delta @ T_w2c[-1]) | |
| return _interpolate_poses(T_w2c, sample_idxs) | |
| def _features(self, img): | |
| """SIFT keypoints and descriptors of an RGB image.""" | |
| # converting with the BGR weights (as in the evaluated setup) only changes the gray mix | |
| return self.sift.detectAndCompute(cv2.cvtColor(img, cv2.COLOR_BGR2GRAY), None) | |
| def _match(features0, features1): | |
| """FLANN matches of two images' SIFT features with Lowe's ratio test -> (N, 2), (N, 2) | |
| float32 points.""" | |
| (kp0, des0), (kp1, des1) = features0, features1 | |
| flann = cv2.FlannBasedMatcher(dict(algorithm=1, trees=5), dict(checks=50)) # KD-tree | |
| good = [m for m, n in flann.knnMatch(des0, des1, k=2) if m.distance < 0.7 * n.distance] | |
| pts0 = np.float32([kp0[m.queryIdx].pt for m in good]).reshape(-1, 2) | |
| pts1 = np.float32([kp1[m.trainIdx].pt for m in good]).reshape(-1, 2) | |
| return pts0, pts1 | |
| def _opencv_multithreading(): | |
| """OpenCV with its default number of threads. Ultralytics (the person tracker) turns | |
| OpenCV's multithreading off for the whole process; SimpleVO is up to ~2x faster with it, with | |
| identical results.""" | |
| num_threads = cv2.getNumThreads() | |
| cv2.setNumThreads(-1) # the default | |
| try: | |
| yield | |
| finally: | |
| cv2.setNumThreads(num_threads) | |
| def _reset_random_state(): | |
| """Reset the process-global random state SimpleVO depends on (C rand(), used by OpenCV's FLANN | |
| kd-trees; OpenCV's RNG; COLMAP's RANSAC) to that of a fresh process, so that results do not | |
| depend on what ran before in the same process. The order matters: COLMAP's seeding also | |
| calls srand().""" | |
| pycolmap.set_random_seed(0) | |
| ctypes.CDLL(ctypes.util.find_library("c")).srand(1) # glibc's initial state | |
| cv2.setRNGSeed(0) | |
| class DPVOOdometry: | |
| """DPVO (Deep Patch Visual Odometry, https://github.com/princeton-vl/DPVO), the optional | |
| alternative to SimpleVO; needs DPVO built from third_party/DPVO (see docs/INSTALL.md). | |
| Unlike GVHMR's wrapper, frames are timestamped with their index (as in DPVO's own demo) | |
| instead of the wall-clock time, which DPVO's damped-linear motion model would otherwise turn | |
| into a dependence on processing speed. DPVO is not deterministic (random patch selection, | |
| atomic float accumulation on the GPU): repeated runs differ by ~0.1-0.4 degrees on average. | |
| Args: | |
| checkpoint: DPVO weights (dpvo.pth) | |
| resize: frames are downscaled by this factor | |
| buffer: DPVO frame buffer size | |
| """ | |
| def __init__(self, checkpoint, resize=0.5, buffer=4000): | |
| self.checkpoint, self.resize, self.buffer = checkpoint, resize, buffer | |
| def __call__(self, video_path, K): | |
| """video and its (3, 3) intrinsics -> (L, 3, 3) world-to-camera rotations (float32)""" | |
| from dpvo.config import cfg | |
| from dpvo.dpvo import DPVO | |
| intrinsics = torch.stack([K[0, 0], K[1, 1], K[0, 2], K[1, 2]]).float() # fx, fy, cx, cy | |
| # frames are decoded in a background thread (OpenCV releases the GIL), overlapping with | |
| # DPVO on the GPU. A thread rather than a process: importing dpvo switches the | |
| # multiprocessing start method to "spawn" for the whole interpreter. | |
| frames = queue.Queue(maxsize=8) | |
| stop = threading.Event() | |
| reader = threading.Thread( | |
| target=_dpvo_video_stream, | |
| args=(frames, stop, str(video_path), intrinsics, self.resize), | |
| daemon=True, | |
| ) | |
| reader.start() | |
| slam = None | |
| pbar = tqdm(total=video_lwh(video_path)[0], desc="DPVO") | |
| # no_grad is essential: DPVO writes network outputs into persistent buffers, which | |
| # would otherwise keep every frame's autograd graph alive (memory grows ~0.25 GB/frame) | |
| try: | |
| with torch.no_grad(): | |
| while (item := frames.get()) is not None: # until the end marker | |
| if isinstance(item, Exception): | |
| raise item | |
| t, image, frame_intrinsics = item | |
| pbar.update() | |
| image = torch.from_numpy(image).permute(2, 0, 1).cuda() | |
| if slam is None: | |
| cfg.merge_from_file(str(DPVO_ROOT / "config" / "default.yaml")) | |
| cfg.BUFFER_SIZE = self.buffer | |
| slam = DPVO( | |
| cfg, str(self.checkpoint), ht=image.shape[1], wd=image.shape[2], | |
| viz=False, | |
| ) # fmt: skip | |
| slam(t, image, frame_intrinsics.cuda()) | |
| for _ in range(12): | |
| slam.update() | |
| traj = _dpvo_terminate(slam) # (L, 7) camera-to-world [t, q (x, y, z, w)] | |
| finally: | |
| pbar.close() | |
| stop.set() # lets the reader exit even if it is blocked on the full queue | |
| reader.join() | |
| R_c2w = quaternion_to_matrix(torch.from_numpy(traj[:, [6, 3, 4, 5]])) | |
| return R_c2w.mT | |
| def _dpvo_terminate(slam): | |
| """slam.terminate(), without its recursion: the (L, 7) camera-to-world poses. | |
| DPVO drops frames with little camera motion and gets their poses by chaining relative poses | |
| back to a kept frame, recursively with one call per frame. When the camera barely moves for | |
| a while, the chains exceed Python's recursion limit (~1000 frames) and cost quadratic time. | |
| Here each pose is computed once, from the pose of the earlier frame it refers to, with the | |
| same operations. | |
| """ | |
| from dpvo import lietorch | |
| from dpvo.lietorch import SE3 | |
| assert not (slam.cfg.LOOP_CLOSURE or slam.cfg.CLASSIC_LOOP_CLOSURE) # not handled here | |
| for _ in range(12): | |
| slam.ran_global_ba[slam.n] = False | |
| slam.update() | |
| kept = {slam.pg.tstamps_[i]: slam.pg.poses_[i] for i in range(slam.n)} | |
| poses = [] | |
| for t in range(slam.counter): | |
| if t in kept: | |
| poses.append(SE3(kept[t])) | |
| else: | |
| t0, dP = slam.pg.delta[t] # t0 < t | |
| poses.append(dP * poses[t0]) | |
| return lietorch.stack(poses, dim=0).inv().data.cpu().numpy() | |
| def _dpvo_video_stream(frames, stop, video_path, intrinsics, resize): | |
| """Puts (frame index, downscaled BGR frame cropped to multiples of 16, intrinsics) into the | |
| frames queue, then None at the end (or the exception, if reading fails).""" | |
| def put(item): | |
| while not stop.is_set(): | |
| try: | |
| frames.put(item, timeout=0.1) | |
| return | |
| except queue.Full: | |
| pass | |
| cap = cv2.VideoCapture(video_path) | |
| try: | |
| t = 0 | |
| while not stop.is_set(): | |
| ret, frame = cap.read() | |
| if not ret: | |
| break | |
| image = cv2.resize(frame, None, fx=resize, fy=resize, interpolation=cv2.INTER_AREA) | |
| h, w, _ = image.shape | |
| put((t, image[: h - h % 16, : w - w % 16], intrinsics * resize)) | |
| t += 1 | |
| put(None) | |
| except Exception as error: | |
| put(error) | |
| finally: | |
| cap.release() | |
| def _interpolate_poses(T_w2c, sample_idxs): | |
| """Poses at sample_idxs -> all frames: linear translation, SLERP rotation.""" | |
| sample_idxs = np.array(sample_idxs) | |
| translations = np.array([T[:3, 3] for T in T_w2c]) | |
| quaternions = np.array([_rotmat_to_quat(T[:3, :3]) for T in T_w2c]) | |
| out = [] | |
| for i in range(sample_idxs[-1] + 1): | |
| if i in sample_idxs: | |
| out.append(T_w2c[np.where(sample_idxs == i)[0][0]]) | |
| continue | |
| nxt = np.searchsorted(sample_idxs, i) | |
| prv = nxt - 1 | |
| t = (i - sample_idxs[prv]) / (sample_idxs[nxt] - sample_idxs[prv]) | |
| T = np.eye(4) | |
| T[:3, :3] = _quat_to_rotmat(_slerp(quaternions[prv], quaternions[nxt], t)) | |
| T[:3, 3] = (1 - t) * translations[prv] + t * translations[nxt] | |
| out.append(T) | |
| return np.array(out) | |
| def _rotmat_to_quat(R): | |
| """3x3 rotation matrix -> quaternion [w, x, y, z].""" | |
| m00, m01, m02 = R[0, 0], R[0, 1], R[0, 2] | |
| m10, m11, m12 = R[1, 0], R[1, 1], R[1, 2] | |
| m20, m21, m22 = R[2, 0], R[2, 1], R[2, 2] | |
| tr = m00 + m11 + m22 | |
| if tr > 0: | |
| S = np.sqrt(tr + 1.0) * 2 # 4 * qw | |
| return np.array([0.25 * S, (m21 - m12) / S, (m02 - m20) / S, (m10 - m01) / S]) | |
| if (m00 > m11) and (m00 > m22): | |
| S = np.sqrt(1.0 + m00 - m11 - m22) * 2 # 4 * qx | |
| return np.array([(m21 - m12) / S, 0.25 * S, (m01 + m10) / S, (m02 + m20) / S]) | |
| if m11 > m22: | |
| S = np.sqrt(1.0 + m11 - m00 - m22) * 2 # 4 * qy | |
| return np.array([(m02 - m20) / S, (m01 + m10) / S, 0.25 * S, (m12 + m21) / S]) | |
| S = np.sqrt(1.0 + m22 - m00 - m11) * 2 # 4 * qz | |
| return np.array([(m10 - m01) / S, (m02 + m20) / S, (m12 + m21) / S, 0.25 * S]) | |
| def _quat_to_rotmat(q): | |
| """Quaternion [w, x, y, z] -> 3x3 rotation matrix.""" | |
| qw, qx, qy, qz = q | |
| return np.array( | |
| [ | |
| [1 - 2 * qy**2 - 2 * qz**2, 2 * qx * qy - 2 * qz * qw, 2 * qx * qz + 2 * qy * qw], | |
| [2 * qx * qy + 2 * qz * qw, 1 - 2 * qx**2 - 2 * qz**2, 2 * qy * qz - 2 * qx * qw], | |
| [2 * qx * qz - 2 * qy * qw, 2 * qy * qz + 2 * qx * qw, 1 - 2 * qx**2 - 2 * qy**2], | |
| ] | |
| ) | |
| def _slerp(q0, q1, t): | |
| """Spherical linear interpolation between quaternions q0 and q1.""" | |
| dot = np.dot(q0, q1) | |
| if dot < 0.0: # shorter path | |
| q1 = -q1 | |
| dot = -dot | |
| if dot > 0.9995: # nearly identical: normalized linear interpolation | |
| result = q0 + t * (q1 - q0) | |
| return result / np.linalg.norm(result) | |
| theta_0 = np.arccos(dot) | |
| theta = theta_0 * t | |
| s0 = np.cos(theta) - dot * np.sin(theta) / np.sin(theta_0) | |
| s1 = np.sin(theta) / np.sin(theta_0) | |
| return (s0 * q0) + (s1 * q1) | |