update with rgbd fusion

depth.py

@@ -5,21 +5,27 @@
 # thrid party
 import argparse
 from tqdm import tqdm
+import numpy as np
 import cv2
 # metric 3d
-from module import Metric
+from modules import Metric
 
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description='Running Metric3D')
     parser.add_argument("--images", help='dir for image files', type=str, required=True)
     parser.add_argument("--focal", help='focal length', type=float, default=None)
+    parser.add_argument("--calib", help='calibration file, overwrite focal', type=str, default=None)
     parser.add_argument("--out", help='dir for output depth', type=str, default='')
     parser.add_argument("--depth-scale", help='depth scale factor', type=float, default=1000.0)
     parser.add_argument("--ckpt", type=str, default='./weights/metric_depth_vit_large_800k.pth', help='checkpoint file')
     parser.add_argument("--model-name", type=str, default='v2-L', choices=['v2-L', 'v2-S'], help='model type')
     args = parser.parse_args()
 
+    if args.calib:
+        calib = np.loadtxt(args.calib)
+        args.focal = (calib[0] + calib[1]) / 2
+
     metric = Metric(
         checkpoint=args.ckpt,
         model_name=args.model_name

mesh.py

@@ -0,0 +1,46 @@
+# standard library
+from typing import *
+# third party
+import argparse
+import numpy as np
+# RGBD Fusion
+from modules import RGBDFusion
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--rgb", type=str, required=True, help="directory to rgb images")
+    parser.add_argument("--depth", type=str, required=True, help="depth directory")
+    parser.add_argument("--traj", type=str, required=True, help="trajectory file")
+    parser.add_argument("--depth-scale", type=float, default=1000.0, help="depth scale factor")
+    parser.add_argument("--viz", type=bool, default=False, help="visualize")
+    parser.add_argument("--focal", type=float, default=None, help="focal length")
+    parser.add_argument("--calib", type=str, default=None, help="calib file, overwrite focal")
+    parser.add_argument("--mesh", type=str, default=None, help="save mesh")
+
+    args = parser.parse_args()
+
+    intr, distort = None, None
+    if args.calib:
+        calib = np.loadtxt(args.calib)
+        intr = calib[:4]
+        if len(calib) > 4:
+            distort = calib[4:]
+    elif args.focal:
+        intr = args.focal
+
+    mesh = RGBDFusion.pipeline(
+        image_dir=args.rgb,
+        depth_dir=args.depth,
+        traj_dir=args.traj,
+        intrinsic=intr,
+        distort=distort,
+        depth_scale=args.depth_scale,
+        viz=args.viz,
+        mesh_save=args.mesh
+    )
+
+    RGBDFusion.simplify_mesh(
+        mesh=mesh,
+        voxel_size=0.05,
+        save=args.mesh
+    )

modules/__init__.py

@@ -0,0 +1,13 @@
+from . import droid as Droid
+from . import utils
+from .data import PosedImageStream
+from .metric import Metric3D as Metric
+from .tsdf import RGBDFusion as RGBDFusion
+
+__ALL__ = [
+    "Droid"
+    "utils",
+    "Metric",
+    "RGBDFusion",
+    "PosedImageStream"
+]

modules/data.py

@@ -0,0 +1,83 @@
+# standard library
+from pathlib import Path
+from typing import *
+# thirdparty
+import cv2
+import numpy as np
+from PIL import Image
+from torch.utils.data import Dataset
+# utils
+from .utils import K_from_intr
+
+class PosedImageStream(Dataset):
+    def __init__(
+        self,
+        image_dir: Path,
+        depth_dir: Optional[Path] = None,
+        traje_dir: Optional[Path] = None,
+        depth_scale: float = 1000.0,
+        stride: Optional[int] = 1,
+        intrinsic: Optional[Union[float, np.ndarray]] = None,
+        distort: Optional[np.ndarray] = None,
+        resize: Optional[Tuple[int, int]] = None,
+    ):
+        self.depth_scale = depth_scale
+        self.distort = distort
+        self.stride = stride
+
+        self.rgb_list = \
+            sorted(list(Path(image_dir).glob('*.[p|j][n|p]g')))[::stride]
+        self.depth_list = None if not depth_dir \
+            else sorted(list(Path(depth_dir).glob('*.png')))[::stride]
+        self.pose_list = None if not traje_dir \
+            else sorted(list(Path(traje_dir).glob('*.txt')))[::stride]
+
+        w0, h0 = Image.open(self.rgb_list[0]).size
+        self.image_size = (w0, h0)
+
+        if intrinsic is None or isinstance(intrinsic, float):
+            focal = intrinsic if intrinsic is not None else\
+                np.max([h0, w0])
+            self.intrinsic = np.array([focal, focal, w0 / 2, h0 / 2])
+        elif isinstance(intrinsic, np.ndarray):
+            self.intrinsic = np.array(intrinsic)
+        else:
+            raise ValueError("intrinsic must be either None, float or np.ndarray")
+
+        self.K_origin = K_from_intr(intr=self.intrinsic)
+
+        # resize
+        if resize:
+            wr, hr = resize
+            h1 = int(h0 * np.sqrt((wr * hr) / (h0 * w0)))
+            w1 = int(w0 * np.sqrt((wr * hr) / (h0 * w0)))
+            self.intrinsic[0::2] *= (w1 / w0)
+            self.intrinsic[1::2] *= (h1 / h0)
+            self.resize = (w1, h1)
+            self.image_size = (w1, h1)
+        else:
+            self.resize = None
+
+    def read_image(self, path: Union[str, Path]) -> np.ndarray:
+        image = cv2.imread(str(path), cv2.IMREAD_UNCHANGED)
+        # # convert rgb
+        # if len(image.shape) == 3 and image.shape[-1] == 3:
+        #     image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+        if self.distort is not None:
+            image = cv2.undistort(image, self.K_origin, self.distort)
+        if self.resize:
+            image = cv2.resize(image, self.resize)
+        return image
+
+    def __len__(self) -> int:
+        return len(self.rgb_list)
+
+    def __getitem__(self, idx) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
+        rgb = self.read_image(self.rgb_list[idx])
+        depth = None if not self.depth_list else \
+            self.read_image(self.depth_list[idx]).astype(np.float32) / self.depth_scale
+        pose = None if not self.pose_list else np.loadtxt(self.pose_list[idx])
+
+        return rgb, depth, pose, self.intrinsic
+
+