Spaces:
Sleeping
Sleeping
| # Copyright (c) Meta Platforms, Inc. and affiliates. | |
| # All rights reserved. | |
| # | |
| # This source code is licensed under the license found in the | |
| # LICENSE file in the root directory of this source tree. | |
| # | |
| # Modified from https://github.com/facebookresearch/vggt | |
| import numpy as np | |
| import pycolmap | |
| from mapanything.third_party.projection import project_3D_points_np | |
| def batch_np_matrix_to_pycolmap( | |
| points3d, | |
| extrinsics, | |
| intrinsics, | |
| tracks, | |
| image_size, | |
| masks=None, | |
| max_reproj_error=None, | |
| max_points3D_val=3000, | |
| shared_camera=False, | |
| camera_type="SIMPLE_PINHOLE", | |
| extra_params=None, | |
| min_inlier_per_frame=64, | |
| points_rgb=None, | |
| ): | |
| """ | |
| Convert Batched NumPy Arrays to PyCOLMAP | |
| Check https://github.com/colmap/pycolmap for more details about its format | |
| NOTE that colmap expects images/cameras/points3D to be 1-indexed | |
| so there is a +1 offset between colmap index and batch index | |
| NOTE: different from VGGSfM, this function: | |
| 1. Use np instead of torch | |
| 2. Frame index and camera id starts from 1 rather than 0 (to fit the format of PyCOLMAP) | |
| """ | |
| # points3d: Px3 | |
| # extrinsics: Nx3x4 | |
| # intrinsics: Nx3x3 | |
| # tracks: NxPx2 | |
| # masks: NxP | |
| # image_size: 2, assume all the frames have been padded to the same size | |
| # where N is the number of frames and P is the number of tracks | |
| N, P, _ = tracks.shape | |
| assert len(extrinsics) == N | |
| assert len(intrinsics) == N | |
| assert len(points3d) == P | |
| assert image_size.shape[0] == 2 | |
| reproj_mask = None | |
| if max_reproj_error is not None: | |
| projected_points_2d, projected_points_cam = project_3D_points_np( | |
| points3d, extrinsics, intrinsics | |
| ) | |
| projected_diff = np.linalg.norm(projected_points_2d - tracks, axis=-1) | |
| projected_points_2d[projected_points_cam[:, -1] <= 0] = 1e6 | |
| reproj_mask = projected_diff < max_reproj_error | |
| if masks is not None and reproj_mask is not None: | |
| masks = np.logical_and(masks, reproj_mask) | |
| elif masks is not None: | |
| masks = masks | |
| else: | |
| masks = reproj_mask | |
| assert masks is not None | |
| if masks.sum(1).min() < min_inlier_per_frame: | |
| print("Not enough inliers per frame, skip BA.") | |
| return None, None | |
| # Reconstruction object, following the format of PyCOLMAP/COLMAP | |
| reconstruction = pycolmap.Reconstruction() | |
| inlier_num = masks.sum(0) | |
| valid_mask = inlier_num >= 2 # a track is invalid if without two inliers | |
| valid_idx = np.nonzero(valid_mask)[0] | |
| # Only add 3D points that have sufficient 2D points | |
| for vidx in valid_idx: | |
| # Use RGB colors if provided, otherwise use zeros | |
| rgb = points_rgb[vidx] if points_rgb is not None else np.zeros(3) | |
| reconstruction.add_point3D(points3d[vidx], pycolmap.Track(), rgb) | |
| num_points3D = len(valid_idx) | |
| camera = None | |
| # frame idx | |
| for fidx in range(N): | |
| # set camera | |
| if camera is None or (not shared_camera): | |
| pycolmap_intri = _build_pycolmap_intri( | |
| fidx, intrinsics, camera_type, extra_params | |
| ) | |
| camera = pycolmap.Camera( | |
| model=camera_type, | |
| width=image_size[0], | |
| height=image_size[1], | |
| params=pycolmap_intri, | |
| camera_id=fidx + 1, | |
| ) | |
| # add camera | |
| reconstruction.add_camera(camera) | |
| # set image | |
| cam_from_world = pycolmap.Rigid3d( | |
| pycolmap.Rotation3d(extrinsics[fidx][:3, :3]), extrinsics[fidx][:3, 3] | |
| ) # Rot and Trans | |
| image = pycolmap.Image( | |
| id=fidx + 1, | |
| name=f"image_{fidx + 1}", | |
| camera_id=camera.camera_id, | |
| cam_from_world=cam_from_world, | |
| ) | |
| points2D_list = [] | |
| point2D_idx = 0 | |
| # NOTE point3D_id start by 1 | |
| for point3D_id in range(1, num_points3D + 1): | |
| original_track_idx = valid_idx[point3D_id - 1] | |
| if (reconstruction.points3D[point3D_id].xyz < max_points3D_val).all(): | |
| if masks[fidx][original_track_idx]: | |
| # It seems we don't need +0.5 for BA | |
| point2D_xy = tracks[fidx][original_track_idx] | |
| # Please note when adding the Point2D object | |
| # It not only requires the 2D xy location, but also the id to 3D point | |
| points2D_list.append(pycolmap.Point2D(point2D_xy, point3D_id)) | |
| # add element | |
| track = reconstruction.points3D[point3D_id].track | |
| track.add_element(fidx + 1, point2D_idx) | |
| point2D_idx += 1 | |
| assert point2D_idx == len(points2D_list) | |
| try: | |
| image.points2D = pycolmap.ListPoint2D(points2D_list) | |
| image.registered = True | |
| except: # noqa | |
| print(f"frame {fidx + 1} is out of BA") | |
| image.registered = False | |
| # add image | |
| reconstruction.add_image(image) | |
| return reconstruction, valid_mask | |
| def pycolmap_to_batch_np_matrix( | |
| reconstruction, device="cpu", camera_type="SIMPLE_PINHOLE" | |
| ): | |
| """ | |
| Convert a PyCOLMAP Reconstruction Object to batched NumPy arrays. | |
| Args: | |
| reconstruction (pycolmap.Reconstruction): The reconstruction object from PyCOLMAP. | |
| device (str): Ignored in NumPy version (kept for API compatibility). | |
| camera_type (str): The type of camera model used (default: "SIMPLE_PINHOLE"). | |
| Returns: | |
| tuple: A tuple containing points3D, extrinsics, intrinsics, and optionally extra_params. | |
| """ | |
| num_images = len(reconstruction.images) | |
| max_points3D_id = max(reconstruction.point3D_ids()) | |
| points3D = np.zeros((max_points3D_id, 3)) | |
| for point3D_id in reconstruction.points3D: | |
| points3D[point3D_id - 1] = reconstruction.points3D[point3D_id].xyz | |
| extrinsics = [] | |
| intrinsics = [] | |
| extra_params = [] if camera_type == "SIMPLE_RADIAL" else None | |
| for i in range(num_images): | |
| # Extract and append extrinsics | |
| pyimg = reconstruction.images[i + 1] | |
| pycam = reconstruction.cameras[pyimg.camera_id] | |
| matrix = pyimg.cam_from_world.matrix() | |
| extrinsics.append(matrix) | |
| # Extract and append intrinsics | |
| calibration_matrix = pycam.calibration_matrix() | |
| intrinsics.append(calibration_matrix) | |
| if camera_type == "SIMPLE_RADIAL": | |
| extra_params.append(pycam.params[-1]) | |
| # Convert lists to NumPy arrays instead of torch tensors | |
| extrinsics = np.stack(extrinsics) | |
| intrinsics = np.stack(intrinsics) | |
| if camera_type == "SIMPLE_RADIAL": | |
| extra_params = np.stack(extra_params) | |
| extra_params = extra_params[:, None] | |
| return points3D, extrinsics, intrinsics, extra_params | |
| ######################################################## | |
| def batch_np_matrix_to_pycolmap_wo_track( | |
| points3d, | |
| points_xyf, | |
| points_rgb, | |
| extrinsics, | |
| intrinsics, | |
| image_size, | |
| shared_camera=False, | |
| camera_type="SIMPLE_PINHOLE", | |
| ): | |
| """ | |
| Convert Batched NumPy Arrays to PyCOLMAP | |
| Different from batch_np_matrix_to_pycolmap, this function does not use tracks. | |
| It saves points3d to colmap reconstruction format only to serve as init for Gaussians or other nvs methods. | |
| Do NOT use this for BA. | |
| """ | |
| # points3d: Px3 | |
| # points_xyf: Px3, with x, y coordinates and frame indices | |
| # points_rgb: Px3, rgb colors | |
| # extrinsics: Nx3x4 | |
| # intrinsics: Nx3x3 | |
| # image_size: 2, assume all the frames have been padded to the same size | |
| # where N is the number of frames and P is the number of tracks | |
| N = len(extrinsics) | |
| P = len(points3d) | |
| # Reconstruction object, following the format of PyCOLMAP/COLMAP | |
| reconstruction = pycolmap.Reconstruction() | |
| for vidx in range(P): | |
| reconstruction.add_point3D(points3d[vidx], pycolmap.Track(), points_rgb[vidx]) | |
| camera = None | |
| # frame idx | |
| for fidx in range(N): | |
| # set camera | |
| if camera is None or (not shared_camera): | |
| pycolmap_intri = _build_pycolmap_intri(fidx, intrinsics, camera_type) | |
| camera = pycolmap.Camera( | |
| model=camera_type, | |
| width=image_size[0], | |
| height=image_size[1], | |
| params=pycolmap_intri, | |
| camera_id=fidx + 1, | |
| ) | |
| # add camera | |
| reconstruction.add_camera(camera) | |
| # set image | |
| cam_from_world = pycolmap.Rigid3d( | |
| pycolmap.Rotation3d(extrinsics[fidx][:3, :3]), extrinsics[fidx][:3, 3] | |
| ) # Rot and Trans | |
| image = pycolmap.Image( | |
| id=fidx + 1, | |
| name=f"image_{fidx + 1}", | |
| camera_id=camera.camera_id, | |
| cam_from_world=cam_from_world, | |
| ) | |
| points2D_list = [] | |
| point2D_idx = 0 | |
| points_belong_to_fidx = points_xyf[:, 2].astype(np.int32) == fidx | |
| points_belong_to_fidx = np.nonzero(points_belong_to_fidx)[0] | |
| for point3D_batch_idx in points_belong_to_fidx: | |
| point3D_id = point3D_batch_idx + 1 | |
| point2D_xyf = points_xyf[point3D_batch_idx] | |
| point2D_xy = point2D_xyf[:2] | |
| points2D_list.append(pycolmap.Point2D(point2D_xy, point3D_id)) | |
| # add element | |
| track = reconstruction.points3D[point3D_id].track | |
| track.add_element(fidx + 1, point2D_idx) | |
| point2D_idx += 1 | |
| assert point2D_idx == len(points2D_list) | |
| try: | |
| image.points2D = pycolmap.ListPoint2D(points2D_list) | |
| image.registered = True | |
| except: # noqa | |
| print(f"frame {fidx + 1} does not have any points") | |
| image.registered = False | |
| # add image | |
| reconstruction.add_image(image) | |
| return reconstruction | |
| def _build_pycolmap_intri(fidx, intrinsics, camera_type, extra_params=None): | |
| """ | |
| Helper function to get camera parameters based on camera type. | |
| Args: | |
| fidx: Frame index | |
| intrinsics: Camera intrinsic parameters | |
| camera_type: Type of camera model | |
| extra_params: Additional parameters for certain camera types | |
| Returns: | |
| pycolmap_intri: NumPy array of camera parameters | |
| """ | |
| if camera_type == "PINHOLE": | |
| pycolmap_intri = np.array( | |
| [ | |
| intrinsics[fidx][0, 0], | |
| intrinsics[fidx][1, 1], | |
| intrinsics[fidx][0, 2], | |
| intrinsics[fidx][1, 2], | |
| ] | |
| ) | |
| elif camera_type == "SIMPLE_PINHOLE": | |
| focal = (intrinsics[fidx][0, 0] + intrinsics[fidx][1, 1]) / 2 | |
| pycolmap_intri = np.array( | |
| [focal, intrinsics[fidx][0, 2], intrinsics[fidx][1, 2]] | |
| ) | |
| elif camera_type == "SIMPLE_RADIAL": | |
| raise NotImplementedError("SIMPLE_RADIAL is not supported yet") | |
| focal = (intrinsics[fidx][0, 0] + intrinsics[fidx][1, 1]) / 2 | |
| pycolmap_intri = np.array( | |
| [ | |
| focal, | |
| intrinsics[fidx][0, 2], | |
| intrinsics[fidx][1, 2], | |
| extra_params[fidx][0], | |
| ] | |
| ) | |
| else: | |
| raise ValueError(f"Camera type {camera_type} is not supported yet") | |
| return pycolmap_intri | |