Initial commit

Co-authored-by: Jensen Zhou <jensen.zhou@stability.ai> Co-authored-by: Aaryaman Vasishta <aaryaman.vasishta@stability.ai>
2025-03-18 07:49:36 -07:00
commit 8a640b7316
53 changed files with 8131 additions and 0 deletions
@@ -0,0 +1,116 @@
+import contextlib
+import os
+import os.path as osp
+import sys
+from typing import cast
+
+import imageio.v3 as iio
+import numpy as np
+import torch
+
+
+class Dust3rPipeline(object):
+    def __init__(self, device: str | torch.device = "cuda"):
+        submodule_path = osp.realpath(
+            osp.join(osp.dirname(__file__), "../../third_party/dust3r/")
+        )
+        if submodule_path not in sys.path:
+            sys.path.insert(0, submodule_path)
+        try:
+            with open(os.devnull, "w") as f, contextlib.redirect_stdout(f):
+                from dust3r.cloud_opt import (  # type: ignore[import]
+                    GlobalAlignerMode,
+                    global_aligner,
+                )
+                from dust3r.image_pairs import make_pairs  # type: ignore[import]
+                from dust3r.inference import inference  # type: ignore[import]
+                from dust3r.model import AsymmetricCroCo3DStereo  # type: ignore[import]
+                from dust3r.utils.image import load_images  # type: ignore[import]
+        except ImportError:
+            raise ImportError(
+                "Missing required submodule: 'dust3r'. Please ensure that all submodules are properly set up.\n\n"
+                "To initialize them, run the following command in the project root:\n"
+                "  git submodule update --init --recursive"
+            )
+
+        self.device = torch.device(device)
+        self.model = AsymmetricCroCo3DStereo.from_pretrained(
+            "naver/DUSt3R_ViTLarge_BaseDecoder_512_dpt"
+        ).to(self.device)
+
+        self._GlobalAlignerMode = GlobalAlignerMode
+        self._global_aligner = global_aligner
+        self._make_pairs = make_pairs
+        self._inference = inference
+        self._load_images = load_images
+
+    def infer_cameras_and_points(
+        self,
+        img_paths: list[str],
+        Ks: list[list] = None,
+        c2ws: list[list] = None,
+        batch_size: int = 16,
+        schedule: str = "cosine",
+        lr: float = 0.01,
+        niter: int = 500,
+        min_conf_thr: int = 3,
+    ) -> tuple[
+        list[np.ndarray], np.ndarray, np.ndarray, list[np.ndarray], list[np.ndarray]
+    ]:
+        num_img = len(img_paths)
+        if num_img == 1:
+            print("Only one image found, duplicating it to create a stereo pair.")
+            img_paths = img_paths * 2
+
+        images = self._load_images(img_paths, size=512)
+        pairs = self._make_pairs(
+            images,
+            scene_graph="complete",
+            prefilter=None,
+            symmetrize=True,
+        )
+        output = self._inference(pairs, self.model, self.device, batch_size=batch_size)
+
+        ori_imgs = [iio.imread(p) for p in img_paths]
+        ori_img_whs = np.array([img.shape[1::-1] for img in ori_imgs])
+        img_whs = np.concatenate([image["true_shape"][:, ::-1] for image in images], 0)
+
+        scene = self._global_aligner(
+            output,
+            device=self.device,
+            mode=self._GlobalAlignerMode.PointCloudOptimizer,
+            same_focals=True,
+            optimize_pp=False,  # True,
+            min_conf_thr=min_conf_thr,
+        )
+
+        # if Ks is not None:
+        #     scene.preset_focal(
+        #         torch.tensor([[K[0, 0], K[1, 1]] for K in Ks])
+        #     )
+
+        if c2ws is not None:
+            scene.preset_pose(c2ws)
+
+        _ = scene.compute_global_alignment(
+            init="msp", niter=niter, schedule=schedule, lr=lr
+        )
+
+        imgs = cast(list, scene.imgs)
+        Ks = scene.get_intrinsics().detach().cpu().numpy().copy()
+        c2ws = scene.get_im_poses().detach().cpu().numpy()  # type: ignore
+        pts3d = [x.detach().cpu().numpy() for x in scene.get_pts3d()]  # type: ignore
+        if num_img > 1:
+            masks = [x.detach().cpu().numpy() for x in scene.get_masks()]
+            points = [p[m] for p, m in zip(pts3d, masks)]
+            point_colors = [img[m] for img, m in zip(imgs, masks)]
+        else:
+            points = [p.reshape(-1, 3) for p in pts3d]
+            point_colors = [img.reshape(-1, 3) for img in imgs]
+
+        # Convert back to the original image size.
+        imgs = ori_imgs
+        Ks[:, :2, -1] *= ori_img_whs / img_whs
+        Ks[:, :2, :2] *= (ori_img_whs / img_whs).mean(axis=1, keepdims=True)[..., None]
+
+        return imgs, Ks, c2ws, points, point_colors