gaussian-splatting 1.17.11__cp310-cp310-win_amd64.whl → 1.18.0__cp310-cp310-win_amd64.whl

gaussian_splatting/camera.py

@@ -133,7 +133,7 @@ def build_camera(
     )
 
 
-def dict2camera(camera_dict, load_depth=False, device="cuda", custom_data: dict = {}):
+def dict2camera(camera_dict, load_mask=True, load_depth=True, device="cuda", custom_data: dict = {}):
     C2W = torch.zeros((4, 4), device=device)
     C2W[:3, 3] = torch.tensor(camera_dict['position'], dtype=torch.float, device=device)
     C2W[:3, :3] = torch.tensor(camera_dict['rotation'], dtype=torch.float, device=device)
@@ -141,6 +141,12 @@ def dict2camera(camera_dict, load_depth=False, device="cuda", custom_data: dict
     Rt = torch.linalg.inv(C2W)
     T = Rt[:3, 3]
     R = Rt[:3, :3]
+    if load_mask and ('ground_truth_image_mask_path' not in camera_dict or camera_dict['ground_truth_image_mask_path'] is None):
+        logging.warning(f"Value of key 'ground_truth_image_mask_path' is not a valid path, skipping mask loading.")
+    if load_depth and ('ground_truth_depth_path' not in camera_dict or camera_dict['ground_truth_depth_path'] is None):
+        logging.warning(f"Value of key 'ground_truth_depth_path' is not a valid path, skipping depth loading.")
+    if load_depth and ('ground_truth_depth_mask_path' not in camera_dict or camera_dict['ground_truth_depth_mask_path'] is None):
+        logging.warning(f"Value of key 'ground_truth_depth_mask_path' is not a valid path, skipping depth mask loading.")
     return build_camera(
         image_width=camera_dict['width'],
         image_height=camera_dict['height'],
@@ -149,7 +155,7 @@ def dict2camera(camera_dict, load_depth=False, device="cuda", custom_data: dict
         R=R,
         T=T,
         image_path=camera_dict['ground_truth_image_path'] if 'ground_truth_image_path' in camera_dict else None,
-        image_mask_path=camera_dict['ground_truth_image_mask_path'] if 'ground_truth_image_mask_path' in camera_dict else None,
+        image_mask_path=camera_dict['ground_truth_image_mask_path'] if (load_mask and 'ground_truth_image_mask_path' in camera_dict) else None,
         depth_path=camera_dict['ground_truth_depth_path'] if (load_depth and 'ground_truth_depth_path' in camera_dict) else None,
         depth_mask_path=camera_dict['ground_truth_depth_mask_path'] if (load_depth and 'ground_truth_depth_mask_path' in camera_dict) else None,
         device=device,

gaussian_splatting/dataset/camera_trainable.py

@@ -81,8 +81,8 @@ class TrainableCameraDataset(CameraDataset):
             json.dump(cameras, f, indent=2)
 
     @classmethod
-    def from_json(cls, path, load_depth=False):
-        cameras = JSONCameraDataset(path, load_depth=load_depth)
+    def from_json(cls, path, load_mask=True, load_depth=True):
+        cameras = JSONCameraDataset(path, load_mask=load_mask, load_depth=load_depth)
         exposures = [(torch.tensor(camera['exposure'], dtype=torch.float) if 'exposure' in camera else torch.eye(3, 4)) for camera in cameras.json_cameras]
         return cls(cameras, exposures)
 
@@ -91,8 +91,8 @@ class FixedTrainableCameraDataset(JSONCameraDataset):
     # Same as TrainableCameraDataset, but is fixed
     # Used for loading cameras saved by TrainableCameraDataset
 
-    def __init__(self, path, load_depth=False):
-        super().__init__(path, load_depth=load_depth)
+    def __init__(self, path, load_mask=True, load_depth=True):
+        super().__init__(path, load_mask=load_mask, load_depth=load_depth)
         self.load_exposures()
 
     def to(self, device):

gaussian_splatting/dataset/colmap/dataset.py

@@ -27,7 +27,7 @@ class ColmapCamera(NamedTuple):
     depth_mask_path: str
 
 
-def parse_colmap_camera(cameras, images, image_dir, depth_dir=None) -> List[ColmapCamera]:
+def parse_colmap_camera(cameras, images, image_dir, load_mask=True, depth_dir=None) -> List[ColmapCamera]:
     parsed_cameras = []
     for _, key in enumerate(cameras):
         extr = cameras[key]
@@ -49,9 +49,11 @@ def parse_colmap_camera(cameras, images, image_dir, depth_dir=None) -> List[Colm
             raise ValueError("Colmap camera model not handled: only undistorted datasets (PINHOLE or SIMPLE_PINHOLE cameras) supported!")
 
         image_path = os.path.join(image_dir, extr.name)
-        image_mask_path = os.path.join(image_dir, os.path.splitext(extr.name)[0] + '_mask.tiff')
-        if not os.path.exists(image_mask_path):
-            image_mask_path = os.path.splitext(image_mask_path)[0] + '.png'
+        image_mask_path = None
+        if load_mask:
+            image_mask_path = os.path.join(image_dir, os.path.splitext(extr.name)[0] + '_mask.tiff')
+            if not os.path.exists(image_mask_path):
+                image_mask_path = os.path.splitext(image_mask_path)[0] + '.png'
         depth_path, depth_mask_path = None, None
         if depth_dir is not None:
             depth_path = os.path.join(depth_dir, os.path.splitext(extr.name)[0] + '.tiff')
@@ -72,7 +74,7 @@ def parse_colmap_camera(cameras, images, image_dir, depth_dir=None) -> List[Colm
     return parsed_cameras
 
 
-def read_colmap_cameras(colmap_folder, load_depth=False) -> List[ColmapCamera]:
+def read_colmap_cameras(colmap_folder, load_mask=True, load_depth=True) -> List[ColmapCamera]:
     path = colmap_folder
     image_dir = os.path.join(path, "images")
     try:
@@ -86,13 +88,13 @@ def read_colmap_cameras(colmap_folder, load_depth=False) -> List[ColmapCamera]:
         cam_extrinsics = read_images_text(cameras_extrinsic_file)
         cam_intrinsics = read_cameras_text(cameras_intrinsic_file)
     depth_dir = os.path.join(path, "depths") if load_depth else None
-    return parse_colmap_camera(cam_extrinsics, cam_intrinsics, image_dir, depth_dir)
+    return parse_colmap_camera(cam_extrinsics, cam_intrinsics, image_dir, load_mask=load_mask, depth_dir=depth_dir)
 
 
 class ColmapCameraDataset(CameraDataset):
-    def __init__(self, colmap_folder, load_depth=False):
+    def __init__(self, colmap_folder, load_mask=True, load_depth=True):
         super().__init__()
-        self.raw_cameras = read_colmap_cameras(colmap_folder, load_depth=load_depth)
+        self.raw_cameras = read_colmap_cameras(colmap_folder, load_mask=load_mask, load_depth=load_depth)
         self.cameras = [build_camera(**cam._asdict()) for cam in self.raw_cameras]
 
     def to(self, device):
@@ -106,5 +108,5 @@ class ColmapCameraDataset(CameraDataset):
         return self.cameras[idx]
 
 
-def ColmapTrainableCameraDataset(colmap_folder, load_depth=False):
-    return TrainableCameraDataset(ColmapCameraDataset(colmap_folder, load_depth=load_depth))
+def ColmapTrainableCameraDataset(colmap_folder, load_mask=True, load_depth=True):
+    return TrainableCameraDataset(ColmapCameraDataset(colmap_folder, load_mask=load_mask, load_depth=load_depth))

gaussian_splatting/dataset/dataset.py

@@ -35,9 +35,10 @@ class CameraDataset:
 
 
 class JSONCameraDataset(CameraDataset):
-    def __init__(self, path, load_depth=False):
+    def __init__(self, path, load_mask=True, load_depth=True):
         with open(path, 'r') as f:
             self.json_cameras = json.load(f)
+        self.load_mask = load_mask
         self.load_depth = load_depth
         self.load_cameras()
 
@@ -51,7 +52,7 @@ class JSONCameraDataset(CameraDataset):
         return self.load_cameras(device=device)
 
     def load_cameras(self, device=None):
-        self.cameras = [dict2camera(camera, load_depth=self.load_depth, device=device) for camera in self.json_cameras]
+        self.cameras = [dict2camera(camera, load_mask=self.load_mask, load_depth=self.load_depth, device=device) for camera in self.json_cameras]
         return self
 
 

gaussian_splatting/prepare.py

@@ -6,11 +6,11 @@ from .trainer import *
 from .trainer.extensions import ScaleRegularizeTrainerWrapper
 
 
-def prepare_dataset(source: str, device: str, trainable_camera: bool = False, load_camera: str = None, load_depth=False) -> CameraDataset:
+def prepare_dataset(source: str, device: str, trainable_camera: bool = False, load_camera: str = None, load_mask=True, load_depth=True) -> CameraDataset:
     if trainable_camera:
-        dataset = (TrainableCameraDataset.from_json(load_camera, load_depth=load_depth) if load_camera else ColmapTrainableCameraDataset(source, load_depth=load_depth)).to(device)
+        dataset = (TrainableCameraDataset.from_json(load_camera, load_mask=load_mask, load_depth=load_depth) if load_camera else ColmapTrainableCameraDataset(source, load_mask=load_mask, load_depth=load_depth)).to(device)
     else:
-        dataset = (FixedTrainableCameraDataset(load_camera, load_depth=load_depth) if load_camera else ColmapCameraDataset(source, load_depth=load_depth)).to(device)
+        dataset = (FixedTrainableCameraDataset(load_camera, load_mask=load_mask, load_depth=load_depth) if load_camera else ColmapCameraDataset(source, load_mask=load_mask, load_depth=load_depth)).to(device)
     return dataset
 
 

gaussian_splatting/render.py

@@ -12,8 +12,11 @@ from gaussian_splatting.utils.lpipsPyTorch import lpips
 from gaussian_splatting.prepare import prepare_dataset, prepare_gaussians
 
 
-def prepare_rendering(sh_degree: int, source: str, device: str, trainable_camera: bool = False, load_ply: str = None, load_camera: str = None, load_depth=False) -> Tuple[CameraDataset, GaussianModel]:
-    dataset = prepare_dataset(source=source, device=device, trainable_camera=trainable_camera, load_camera=load_camera, load_depth=load_depth)
+def prepare_rendering(
+        sh_degree: int, source: str, device: str,
+        trainable_camera: bool = False, load_ply: str = None, load_camera: str = None,
+        load_depth=False) -> Tuple[CameraDataset, GaussianModel]:
+    dataset = prepare_dataset(source=source, device=device, trainable_camera=trainable_camera, load_camera=load_camera, load_mask=False, load_depth=load_depth)
     gaussians = prepare_gaussians(sh_degree=sh_degree, source=source, device=device, trainable_camera=trainable_camera, load_ply=load_ply)
     return dataset, gaussians
 
@@ -101,5 +104,5 @@ if __name__ == "__main__":
     with torch.no_grad():
         dataset, gaussians = prepare_rendering(
             sh_degree=args.sh_degree, source=args.source, device=args.device, trainable_camera=args.mode == "camera",
-            load_ply=load_ply, load_camera=args.load_camera, load_depth=True)
+            load_ply=load_ply, load_camera=args.load_camera, load_depth=args.save_depth_pcd)
         rendering(dataset, gaussians, save, save_pcd=args.save_depth_pcd, rescale_depth_gt=not args.no_rescale_depth_gt)

gaussian_splatting/train.py

@@ -12,8 +12,12 @@ from gaussian_splatting.trainer import AbstractTrainer
 from gaussian_splatting.prepare import basemodes, shliftmodes, prepare_dataset, prepare_gaussians, prepare_trainer
 
 
-def prepare_training(sh_degree: int, source: str, device: str, mode: str, trainable_camera: bool = False, load_ply: str = None, load_camera: str = None, load_depth=False, with_scale_reg=False, configs={}) -> Tuple[CameraDataset, GaussianModel, AbstractTrainer]:
-    dataset = prepare_dataset(source=source, device=device, trainable_camera=trainable_camera, load_camera=load_camera, load_depth=load_depth)
+def prepare_training(
+        sh_degree: int, source: str, device: str, mode: str,
+        trainable_camera: bool = False, load_ply: str = None, load_camera: str = None,
+        load_mask=False, load_depth=False,
+        with_scale_reg=False, configs={}) -> Tuple[CameraDataset, GaussianModel, AbstractTrainer]:
+    dataset = prepare_dataset(source=source, device=device, trainable_camera=trainable_camera, load_camera=load_camera, load_mask=load_mask, load_depth=load_depth)
     gaussians = prepare_gaussians(sh_degree=sh_degree, source=source, device=device, trainable_camera=trainable_camera, load_ply=load_ply)
     trainer = prepare_trainer(gaussians=gaussians, dataset=dataset, mode=mode, trainable_camera=trainable_camera, load_ply=load_ply, with_scale_reg=with_scale_reg, configs=configs)
     return dataset, gaussians, trainer
@@ -65,6 +69,7 @@ if __name__ == "__main__":
     parser.add_argument("-i", "--iteration", default=30000, type=int)
     parser.add_argument("-l", "--load_ply", default=None, type=str)
     parser.add_argument("--load_camera", default=None, type=str)
+    parser.add_argument("--no_image_mask", action="store_true")
     parser.add_argument("--no_depth_data", action="store_true")
     parser.add_argument("--with_scale_reg", action="store_true")
     parser.add_argument("--mode", choices=sorted(list(set(list(basemodes.keys()) + list(shliftmodes.keys())))), default="base")
@@ -78,7 +83,9 @@ if __name__ == "__main__":
     configs = {o.split("=", 1)[0]: eval(o.split("=", 1)[1]) for o in args.option}
     dataset, gaussians, trainer = prepare_training(
         sh_degree=args.sh_degree, source=args.source, device=args.device, mode=args.mode, trainable_camera="camera" in args.mode,
-        load_ply=args.load_ply, load_camera=args.load_camera, load_depth=not args.no_depth_data, with_scale_reg=args.with_scale_reg, configs=configs)
+        load_ply=args.load_ply, load_camera=args.load_camera,
+        load_mask=not args.no_image_mask, load_depth=not args.no_depth_data,
+        with_scale_reg=args.with_scale_reg, configs=configs)
     dataset.save_cameras(os.path.join(args.destination, "cameras.json"))
     torch.cuda.empty_cache()
     training(

{gaussian_splatting-1.17.11.dist-info → gaussian_splatting-1.18.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: gaussian_splatting
-Version: 1.17.11
+Version: 1.18.0
 Summary: Refactored python training and inference code for 3D Gaussian Splatting
 Home-page: https://github.com/yindaheng98/gaussian-splatting
 Author: yindaheng98

{gaussian_splatting-1.17.11.dist-info → gaussian_splatting-1.18.0.dist-info}/RECORD

@@ -1,20 +1,20 @@
 gaussian_splatting/__init__.py,sha256=CiOZMcyPTAaKtEuMZUhEda_Ad4_RUhmIstB-A3iuOJY,131
-gaussian_splatting/camera.py,sha256=LB44l5ok36jhkWXja2oIvnuBio3pfE5m18aCEFKl2jc,8477
+gaussian_splatting/camera.py,sha256=vo7mu6lyFpIhDqOAgNiJuPan8t_nDJn5cJkAYygLFcA,9243
 gaussian_splatting/camera_trainable.py,sha256=nI6hFFRV2ev7VwLlKUbzEdN9zUmngYZAANGLr1p1yBA,3841
 gaussian_splatting/gaussian_model.py,sha256=_Dy_dDa2prALhVgg428a-O8-8PODg3c_JPkOJJ8X4o8,13275
-gaussian_splatting/prepare.py,sha256=SdljFF2TU-ETkIY2ceHEz9tGA-o4HfHxGsFKIhyB1Ek,3042
-gaussian_splatting/render.py,sha256=2dP7Yr5e5uLJAs9MwRJSBY9eZQjaDtf9XB-HdtPmxPc,5930
-gaussian_splatting/train.py,sha256=VAFlPHB5ETH8BiYhe0LWY2UWXyB16tSRPKUaV5VPp5A,5196
+gaussian_splatting/prepare.py,sha256=GWPRpufg5larcKGNwlRtN22xVs9fp3ptMu11rQeySX8,3141
+gaussian_splatting/render.py,sha256=sh67INAUhEy5lfkmuB6RBTqe08NF-8MMR7QXTy4Ogg8,5990
+gaussian_splatting/train.py,sha256=no945bMVVRylvkkllQIURDplSwm8EydvUhn-wD9Cn2k,5388
 gaussian_splatting/dataset/__init__.py,sha256=-runuT-61P0YVpfV_WXqwUZM1oY0N012YH13Bt3rzSU,138
-gaussian_splatting/dataset/camera_trainable.py,sha256=D8bqeLVcnQ4qrsy3U1s6BVzdG4KCCPGR-Tygj6nLp-Q,4969
-gaussian_splatting/dataset/dataset.py,sha256=mlcIS0pJNdUIT-RObcQNYwgxxTOFJ0OYg3AlhoQ4Mww,2315
+gaussian_splatting/dataset/camera_trainable.py,sha256=Kd8v-_ZJ9dLIQ2QyVOXbmouYf5QjbgOgHNRHVpkgCms,5041
+gaussian_splatting/dataset/dataset.py,sha256=0tmIZ5P7kOEdABiEAXPznkRN91e5rcT5VsAzOLoOuEM,2392
 gaussian_splatting/dataset/colmap/__init__.py,sha256=YEYT2k2WJSqrkkZq4KAJYS9UMgqU6W6TJaeHLRc1CM4,213
-gaussian_splatting/dataset/colmap/dataset.py,sha256=Lq2b3hMdtOmdqPjvEjR6CLukAR7dZBEKMz8yzDD2Bgo,4519
+gaussian_splatting/dataset/colmap/dataset.py,sha256=0UBQ6ynOqElHZSphJ-MSbYQqCwwYZaAXl1y9AY5YKuY,4720
 gaussian_splatting/dataset/colmap/params_init.py,sha256=6_6gZ0Wl4aZrps2PJ_U234sxW5D-vOTfwioVa1FWC-E,1802
 gaussian_splatting/dataset/colmap/read_write_model.py,sha256=TenI7ai5UV7Ksg2vAXvJWnYFwOOo1tlS_633RfCLuQU,23137
-gaussian_splatting/diff_gaussian_rasterization/_C.cp310-win_amd64.pyd,sha256=YN0Oev5kgszT2ukGdSi_Fzkdfff3ugTmOyN3C4f84vs,1287680
+gaussian_splatting/diff_gaussian_rasterization/_C.cp310-win_amd64.pyd,sha256=xBBzubbNtBsa59qXqReMer7oFeshmdxBcZtMMNaueJk,1287680
 gaussian_splatting/diff_gaussian_rasterization/__init__.py,sha256=a9D0IZiPx-Mk1795hSq54T-NYT4MtEN_MZrxeMhw0Eo,6705
-gaussian_splatting/simple_knn/_C.cp310-win_amd64.pyd,sha256=XXsPotAF7TR-3ik38t_erNhxU4iekR6cTtgduLau8bU,1156608
+gaussian_splatting/simple_knn/_C.cp310-win_amd64.pyd,sha256=El47ul_aP-C8jETADhbfzvtNZUg8OlmVkaoIKSCpW1k,1156608
 gaussian_splatting/trainer/__init__.py,sha256=962fEY8A0spSQn5de_d_LkPOjA1PYKrLbuAkxwZo7mI,940
 gaussian_splatting/trainer/abc.py,sha256=kpYnJjLOhsyhE-V2J79EC9nih6MYBcXkmK9cHUA-3ao,4022
 gaussian_splatting/trainer/base.py,sha256=gO1x4m82xrZNl8NZVw2CWYqIvZJIMUWmBtPZQPeyxJ0,3370
@@ -45,8 +45,8 @@ gaussian_splatting/utils/lpipsPyTorch/modules/__init__.py,sha256=47DEQpj8HBSa-_T
 gaussian_splatting/utils/lpipsPyTorch/modules/lpips.py,sha256=YScu0oXIEstCCjJVRItS_R_csUw70sBMFuP8Syl2UdI,1187
 gaussian_splatting/utils/lpipsPyTorch/modules/networks.py,sha256=kqIebq7dAhHypTXweFVEf_RDbN7_Zv7O3MlD-CfRvpg,2788
 gaussian_splatting/utils/lpipsPyTorch/modules/utils.py,sha256=TDcem3E3HqDNN2MT8qlOL_BKVHeO4HRE77JxF-kOWk8,915
-gaussian_splatting-1.17.11.dist-info/licenses/LICENSE.md,sha256=bMuRQKn0u485mx8JBBTJ5Simc-aWHaQsxmoB6jsg5oE,4752
-gaussian_splatting-1.17.11.dist-info/METADATA,sha256=YhuQHoRNZ-SttnBQZ0IumPV8JWSZ4MzMe_hRSm3dUW4,17184
-gaussian_splatting-1.17.11.dist-info/WHEEL,sha256=KUuBC6lxAbHCKilKua8R9W_TM71_-9Sg5uEP3uDWcoU,101
-gaussian_splatting-1.17.11.dist-info/top_level.txt,sha256=uaYrPYXRHhpybgCnsoazTcdhpzZGnLT_vd5eoRzBWWI,19
-gaussian_splatting-1.17.11.dist-info/RECORD,,
+gaussian_splatting-1.18.0.dist-info/licenses/LICENSE.md,sha256=bMuRQKn0u485mx8JBBTJ5Simc-aWHaQsxmoB6jsg5oE,4752
+gaussian_splatting-1.18.0.dist-info/METADATA,sha256=z6Q1TP39OROC_mlfGLSBgZQUip5neZOFgqQKx1Gbj_c,17183
+gaussian_splatting-1.18.0.dist-info/WHEEL,sha256=KUuBC6lxAbHCKilKua8R9W_TM71_-9Sg5uEP3uDWcoU,101
+gaussian_splatting-1.18.0.dist-info/top_level.txt,sha256=uaYrPYXRHhpybgCnsoazTcdhpzZGnLT_vd5eoRzBWWI,19
+gaussian_splatting-1.18.0.dist-info/RECORD,,