gaussian-splatting 1.16.13__cp312-cp312-win_amd64.whl → 1.17.0__cp312-cp312-win_amd64.whl

gaussian_splatting/camera.py

@@ -2,7 +2,7 @@ import os
 from typing import NamedTuple, Callable
 import torch
 import logging
-from .utils import fov2focal, focal2fov, getProjectionMatrix, getWorld2View2, read_image, read_depth, read_depth_mask, matrix_to_quaternion
+from .utils import fov2focal, focal2fov, getProjectionMatrix, getWorld2View2, read_image, read_image_mask, read_depth, read_depth_mask, matrix_to_quaternion
 
 
 class Camera(NamedTuple):
@@ -18,9 +18,11 @@ class Camera(NamedTuple):
     camera_center: torch.Tensor
     quaternion: torch.Tensor
     ground_truth_image_path: str
+    ground_truth_image_mask_path: str
     ground_truth_depth_path: str
     ground_truth_depth_mask_path: str
     ground_truth_image: torch.Tensor = None
+    ground_truth_image_mask: torch.Tensor = None
     ground_truth_depth: torch.Tensor = None
     ground_truth_depth_mask: torch.Tensor = None
     postprocess: Callable[['Camera', torch.Tensor], torch.Tensor] = lambda camera, x: x
@@ -47,9 +49,9 @@ def camera2dict(camera: Camera, id):
         'fx': fov2focal(camera.FoVx, camera.image_width),
         'fy': fov2focal(camera.FoVy, camera.image_height),
         'ground_truth_image_path': camera.ground_truth_image_path.replace("\\", "/") if camera.ground_truth_image_path else None,
+        'ground_truth_image_mask_path': camera.ground_truth_image_mask_path.replace("\\", "/") if camera.ground_truth_image_mask_path else None,
         'ground_truth_depth_path': camera.ground_truth_depth_path.replace("\\", "/") if camera.ground_truth_depth_path else None,
         'ground_truth_depth_mask_path': camera.ground_truth_depth_mask_path.replace("\\", "/") if camera.ground_truth_depth_mask_path else None,
-        "img_name": os.path.basename(camera.ground_truth_image_path) if camera.ground_truth_image_path else None,
     }
     return camera_entry
 
@@ -58,7 +60,8 @@ def build_camera(
         image_height: int, image_width: int,
         FoVx: float, FoVy: float,
         R: torch.Tensor, T: torch.Tensor,
-        image_path: str = None, depth_path: str = None, depth_mask_path: str = None,
+        image_path: str = None, image_mask_path: str = None,
+        depth_path: str = None, depth_mask_path: str = None,
         device="cuda", custom_data: dict = {}
 ):
     R, T = R.to(device=device, dtype=torch.float), T.to(device=device, dtype=torch.float)
@@ -77,6 +80,15 @@ def build_camera(
         if gt_image.shape[1:] != (image_height, image_width):
             logging.warning(f"gt_image shape {gt_image.shape} does not match expected shape {image_height}x{image_width}, resizing.")
             gt_image = torch.nn.functional.interpolate(gt_image.unsqueeze(0), size=(image_height, image_width), mode='bilinear', align_corners=False).squeeze(0)
+    gt_image_mask = None
+    if image_mask_path is not None:
+        if os.path.exists(image_mask_path):
+            gt_image_mask = read_image_mask(image_mask_path).to(device)
+            if gt_image_mask.shape != (image_height, image_width):
+                logging.warning(f"gt_image_mask shape {gt_image_mask.shape} does not match expected shape {image_height}x{image_width}, resizing.")
+                gt_image_mask = torch.nn.functional.interpolate(gt_image_mask.unsqueeze(0).unsqueeze(0), size=(image_height, image_width), mode='bilinear', align_corners=False).squeeze(0).squeeze(0)
+        elif not os.path.exists(image_mask_path):
+            logging.warning(f"Image mask path {image_mask_path} does not exist, skipping mask loading.")
     gt_depth = None
     if depth_path is not None:
         if os.path.exists(depth_path):
@@ -108,6 +120,8 @@ def build_camera(
         quaternion=quaternion,
         ground_truth_image_path=image_path,
         ground_truth_image=gt_image,
+        ground_truth_image_mask_path=image_mask_path,
+        ground_truth_image_mask=gt_image_mask,
         ground_truth_depth_path=depth_path,
         ground_truth_depth=gt_depth,
         ground_truth_depth_mask_path=depth_mask_path,
@@ -132,6 +146,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,
         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/colmap/dataset.py

@@ -22,6 +22,7 @@ class ColmapCamera(NamedTuple):
     R: torch.Tensor
     T: torch.Tensor
     image_path: str
+    image_mask_path: str
     depth_path: str
     depth_mask_path: str
 
@@ -48,6 +49,9 @@ 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'
         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')
@@ -61,6 +65,7 @@ def parse_colmap_camera(cameras, images, image_dir, depth_dir=None) -> List[Colm
             R=torch.from_numpy(R).type(torch.float), T=torch.from_numpy(T).type(torch.float),
             FoVy=FovY, FoVx=FovX,
             image_path=image_path,
+            image_mask_path=image_mask_path,
             depth_path=depth_path,
             depth_mask_path=depth_mask_path,
         ))

gaussian_splatting/render.py

@@ -48,6 +48,7 @@ def build_pcd_rescale(
 def rendering(
         dataset: CameraDataset, gaussians: GaussianModel, save: str, save_pcd: bool = False,
         rescale_depth_gt: bool = True) -> None:
+    dataset.save_cameras(os.path.join(save, "cameras.json"))
     render_path = os.path.join(save, "renders")
     gt_path = os.path.join(save, "gt")
     makedirs(render_path, exist_ok=True)

gaussian_splatting/trainer/base.py

@@ -70,6 +70,10 @@ class BaseTrainer(AbstractTrainer):
     def loss(self, out: dict, camera: Camera) -> torch.Tensor:
         render = out["render"]
         gt = camera.ground_truth_image
+        mask = camera.ground_truth_image_mask
+        if mask is not None:
+            render = render * mask.unsqueeze(0)
+            gt = gt * mask.unsqueeze(0)
         Ll1 = l1_loss(render, gt)
         ssim_value = ssim(render, gt)
         loss = (1.0 - self.lambda_dssim) * Ll1 + self.lambda_dssim * (1.0 - ssim_value)

gaussian_splatting/utils/general.py

@@ -23,6 +23,10 @@ def read_image(image_path):
     return torch_image[:3, ...].clamp(0.0, 1.0)
 
 
+def read_image_mask(image_path):
+    return read_depth_mask(image_path)
+
+
 def read_png_depth(depth_path):
     cv2_depth = cv2.imread(depth_path, cv2.IMREAD_UNCHANGED)
     if cv2_depth.ndim != 2:

{gaussian_splatting-1.16.13.dist-info → gaussian_splatting-1.17.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: gaussian_splatting
-Version: 1.16.13
+Version: 1.17.0
 Summary: Refactored python training and inference code for 3D Gaussian Splatting
 Home-page: https://github.com/yindaheng98/gaussian-splatting
 Author: yindaheng98
@@ -36,6 +36,7 @@ We **refactored the original code following the standard Python package structur
 * [x] camera and 3DGS parameters joint training
 * [x] depth regularization
 * [x] local relative depth regularization
+* [x] image mask
 * [ ] integrated 2DGS (integrated [gsplat](https://github.com/nerfstudio-project/gsplat) backend)
 
 ## Install

{gaussian_splatting-1.16.13.dist-info → gaussian_splatting-1.17.0.dist-info}/RECORD

@@ -1,23 +1,23 @@
 gaussian_splatting/__init__.py,sha256=CiOZMcyPTAaKtEuMZUhEda_Ad4_RUhmIstB-A3iuOJY,131
-gaussian_splatting/camera.py,sha256=aiIUnGgMvyY6h4SjuuzDW1AM5QsZhxZMF2Fsq913jHs,7056
+gaussian_splatting/camera.py,sha256=jQdJZaHJABXshZ9LH9U9FC_-fi0bEUZo4m3b71PImbA,8218
 gaussian_splatting/camera_trainable.py,sha256=XHynfc-9Z3d8N00u4rUlyAglwTadT_BDoV4ffOBLcLU,3862
 gaussian_splatting/gaussian_model.py,sha256=Jg0RJXNeCMNjQpZStag2H8JpvTeukjgqnrHVHJRSDcY,13296
 gaussian_splatting/prepare.py,sha256=SdljFF2TU-ETkIY2ceHEz9tGA-o4HfHxGsFKIhyB1Ek,3042
-gaussian_splatting/render.py,sha256=KFqMn1Uo76dOYbwmoBtU4_oVADaQtLKhbw6TmDUT7lk,5830
+gaussian_splatting/render.py,sha256=SyvRmiZ1OzYTZf7lQTRPjncQ0vL9297oWF3SgxEfJnA,5892
 gaussian_splatting/train.py,sha256=VAFlPHB5ETH8BiYhe0LWY2UWXyB16tSRPKUaV5VPp5A,5196
 gaussian_splatting/dataset/__init__.py,sha256=-runuT-61P0YVpfV_WXqwUZM1oY0N012YH13Bt3rzSU,138
 gaussian_splatting/dataset/camera_trainable.py,sha256=pHtjXXCAlNbbtTvj56-UogM8BHPjc_CNCwFWxixDYkY,4935
 gaussian_splatting/dataset/dataset.py,sha256=mlcIS0pJNdUIT-RObcQNYwgxxTOFJ0OYg3AlhoQ4Mww,2315
 gaussian_splatting/dataset/colmap/__init__.py,sha256=YEYT2k2WJSqrkkZq4KAJYS9UMgqU6W6TJaeHLRc1CM4,213
-gaussian_splatting/dataset/colmap/dataset.py,sha256=utatJesp1QMC739IKA9hRU2-EEqzQUtuRW3rJZ8x3ug,4223
+gaussian_splatting/dataset/colmap/dataset.py,sha256=Lq2b3hMdtOmdqPjvEjR6CLukAR7dZBEKMz8yzDD2Bgo,4519
 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.cp312-win_amd64.pyd,sha256=iEt4wOtQmfAasHiYSzx9sRjHyR1GjaH6O4Wb3f0JRtw,1300992
+gaussian_splatting/diff_gaussian_rasterization/_C.cp312-win_amd64.pyd,sha256=aVVUrdsCqUbuR2OQoYW3R8YHw1A4_2JySQHsn3OnWcY,1300992
 gaussian_splatting/diff_gaussian_rasterization/__init__.py,sha256=a9D0IZiPx-Mk1795hSq54T-NYT4MtEN_MZrxeMhw0Eo,6705
-gaussian_splatting/simple_knn/_C.cp312-win_amd64.pyd,sha256=mwTbEgOsY6DZJsDhL3g60i2qfv9_pVgq96mJe0nSymI,1168896
+gaussian_splatting/simple_knn/_C.cp312-win_amd64.pyd,sha256=bYSNOgKdXezQ_xcSEHR-I4JhJL5JIp6vgZEUEgsxYP8,1168896
 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=EQI_0q2t53XJyePE061_nIfACPb1-vfsFlmFoSM5OYw,2628
+gaussian_splatting/trainer/base.py,sha256=GrenIGWMXtF-2-b3b1WCA-gQV30TpMQMxBZmPPVKFSg,2795
 gaussian_splatting/trainer/camera_trainable.py,sha256=TBQXn2f578qeizPz6tgqFm-GRvttv9duuB1xx7_J9TQ,4567
 gaussian_splatting/trainer/combinations.py,sha256=7NX4fXdDOx8ri1_mgAaWNx-YVdo5XsqMlr9qy-Ll2MM,5329
 gaussian_splatting/trainer/depth.py,sha256=PxWBSNxzoQcRfCFI_yJnJMS6s8qFWn81CXK6O6ffXL0,7059
@@ -32,7 +32,7 @@ gaussian_splatting/trainer/densifier/trainer.py,sha256=b1MD-BvlfkVMn2Jodus7LA4Vc
 gaussian_splatting/trainer/extensions/__init__.py,sha256=pWCBbZiXZ8bNND2hdrmSejuvJKLf26xHIxIxz-i--Bs,100
 gaussian_splatting/trainer/extensions/scale_reg.py,sha256=bN8TV_woH64IEjgOLTRcLwJL19MLpWwGZPMbn1o-Xy8,2040
 gaussian_splatting/utils/__init__.py,sha256=p7csDXqwQdXRb0JQs_BsOTBo6ogGrCqW9NEDMfBBC5A,187
-gaussian_splatting/utils/general.py,sha256=lAK6_8shqJZc--YeBP1g5Z-7WYX9qVO52ern7g3NcfA,2672
+gaussian_splatting/utils/general.py,sha256=SXV98H9m63HAk97QIaHn99Ev-qx0m-5lWjE8KuVg4iA,2750
 gaussian_splatting/utils/graphics.py,sha256=ZcI03aYO21kSqX0ofrz2L64tMP4jsNR1poubU6n35fw,1636
 gaussian_splatting/utils/loss.py,sha256=exGQFzH3OO71xnPj_99MxsHo80EzCxumGULJR5AW64Y,2255
 gaussian_splatting/utils/metrics.py,sha256=R5SecfLUqt6_zNOAGXF6HsO5wA2DraFzK3inV7pz-Ss,573
@@ -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.16.13.dist-info/licenses/LICENSE.md,sha256=bMuRQKn0u485mx8JBBTJ5Simc-aWHaQsxmoB6jsg5oE,4752
-gaussian_splatting-1.16.13.dist-info/METADATA,sha256=2bXYWeXP8-pvMB8ULXS3CFzG9M3kIKz5qPRDDRQnDfc,17029
-gaussian_splatting-1.16.13.dist-info/WHEEL,sha256=8UP9x9puWI0P1V_d7K2oMTBqfeLNm21CTzZ_Ptr0NXU,101
-gaussian_splatting-1.16.13.dist-info/top_level.txt,sha256=uaYrPYXRHhpybgCnsoazTcdhpzZGnLT_vd5eoRzBWWI,19
-gaussian_splatting-1.16.13.dist-info/RECORD,,
+gaussian_splatting-1.17.0.dist-info/licenses/LICENSE.md,sha256=bMuRQKn0u485mx8JBBTJ5Simc-aWHaQsxmoB6jsg5oE,4752
+gaussian_splatting-1.17.0.dist-info/METADATA,sha256=h5nRgV0Ew4gcpFQ9OuD3lb2K8iA2spi7P6o0VYH6BCo,17046
+gaussian_splatting-1.17.0.dist-info/WHEEL,sha256=8UP9x9puWI0P1V_d7K2oMTBqfeLNm21CTzZ_Ptr0NXU,101
+gaussian_splatting-1.17.0.dist-info/top_level.txt,sha256=uaYrPYXRHhpybgCnsoazTcdhpzZGnLT_vd5eoRzBWWI,19
+gaussian_splatting-1.17.0.dist-info/RECORD,,