gaussian-splatting 1.18.1__cp312-cp312-win_amd64.whl → 1.19.1__cp312-cp312-win_amd64.whl

gaussian_splatting/prepare.py

@@ -8,9 +8,17 @@ from .trainer.extensions import ScaleRegularizeTrainerWrapper
 
 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_mask=load_mask, load_depth=load_depth) if load_camera else ColmapTrainableCameraDataset(source, load_mask=load_mask, 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_mask=load_mask, load_depth=load_depth) if load_camera else ColmapCameraDataset(source, load_mask=load_mask, 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/trainer/abc.py

@@ -33,6 +33,10 @@ class AbstractTrainer(ABC):
     def schedulers(self) -> Dict[str, Callable[[int], float]]:
         raise ValueError("Schedulers is not set")
 
+    @abstractmethod
+    def preprocess(self, camera: Camera) -> Camera:
+        pass
+
     @abstractmethod
     def loss(self, out: dict, camera: Camera) -> torch.Tensor:
         pass
@@ -49,6 +53,7 @@ class AbstractTrainer(ABC):
 
     def step(self, camera: Camera):
         self.update_learning_rate()
+        camera = self.preprocess(camera)
         out = self.model(camera)
         loss = self.loss(out, camera)
         loss.backward()
@@ -100,6 +105,9 @@ class TrainerWrapper(AbstractTrainer):
     def schedulers(self) -> Dict[str, Callable[[int], float]]:
         return self.base_trainer.schedulers
 
+    def preprocess(self, camera: Camera) -> Camera:
+        return self.base_trainer.preprocess(camera)
+
     def loss(self, out: dict, camera: Camera) -> torch.Tensor:
         return self.base_trainer.loss(out, camera)
 

gaussian_splatting/trainer/base.py

@@ -22,12 +22,23 @@ class BaseTrainer(AbstractTrainer):
             opacity_lr=0.025,
             scaling_lr=0.005,
             rotation_lr=0.001,
-            mask_mode="none",  # "ignore", "random", "bg_color"
+            mask_mode="none",
+            # "none"=do not use mask
+            # "ignore"=loss of the masked area will be set to 0
+            # "bg_color"=fill the masked area of ground truth with the bg_color for rendering
+            bg_color=None,
+            # None=do not change bg_color
+            # "random"=set bg_color to random color
+            # tuple(float, float, float)=set bg_color to the given color
     ):
         super().__init__()
         self.lambda_dssim = lambda_dssim
-        assert mask_mode in ["none", "ignore", "random", "bg_color"], f"Unknown mask policy: {mask_mode}"
+        assert mask_mode in ["none", "ignore", "bg_color"], f"Unknown mask policy: {mask_mode}"
+        assert bg_color is None or bg_color == "random" or (
+            isinstance(bg_color, tuple) and len(bg_color) == 3 and all(isinstance(c, float) for c in bg_color)
+        ), f"bg_color must be 'random' or a RGB value tuple(float, float, float), got {bg_color}"
         self.mask_mode = mask_mode
+        self.bg_color = bg_color
         params = [
             {'params': [model._xyz], 'lr': position_lr_init * scene_extent, "name": "xyz"},
             {'params': [model._features_dc], 'lr': feature_lr, "name": "f_dc"},
@@ -70,6 +81,17 @@ class BaseTrainer(AbstractTrainer):
     def schedulers(self) -> Dict[str, Callable[[int], float]]:
         return self._schedulers
 
+    def preprocess(self, camera: Camera) -> Camera:
+        if self.bg_color == "random":
+            camera = camera._replace(bg_color=torch.rand_like(camera.bg_color))
+        elif isinstance(self.bg_color, tuple):
+            camera = camera._replace(bg_color=torch.tensor(self.bg_color, device=camera.bg_color.device, dtype=camera.bg_color.dtype))
+        elif self.bg_color is None:
+            pass
+        else:
+            raise ValueError(f"bg_color must be 'random' or a tuple(int, int, int), got {self.bg_color}")
+        return camera
+
     def loss(self, out: dict, camera: Camera) -> torch.Tensor:
         render = out["render"]
         gt = camera.ground_truth_image
@@ -81,12 +103,11 @@ class BaseTrainer(AbstractTrainer):
                 assert mask is not None, "Mask is required for 'ignore' mask policy"
                 render = render * mask.unsqueeze(0)
                 gt = gt * mask.unsqueeze(0)
-            case "random":
-                assert mask is not None, "Mask is required for 'random' mask policy"
-                gt = gt * mask.unsqueeze(0) + (1 - mask.unsqueeze(0)) * torch.rand_like(gt)
             case "bg_color":
                 assert mask is not None, "Mask is required for 'bg_color' mask policy"
-                gt = gt * mask.unsqueeze(0) + (1 - mask.unsqueeze(0)) * camera.bg_color.unsqueeze(-1).unsqueeze(-1)
+                # bg_color after postprocess
+                bg_color = camera.postprocess(camera, camera.bg_color.unsqueeze(-1).unsqueeze(-1)).clamp(0.0, 1.0)
+                gt = gt * mask.unsqueeze(0) + (1 - mask.unsqueeze(0)) * bg_color
             case _:
                 raise ValueError(f"Unknown mask policy: {self.mask_mode}")
         Ll1 = l1_loss(render, gt)

{gaussian_splatting-1.18.1.dist-info → gaussian_splatting-1.19.1.dist-info}/METADATA

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

{gaussian_splatting-1.18.1.dist-info → gaussian_splatting-1.19.1.dist-info}/RECORD

@@ -2,7 +2,7 @@ gaussian_splatting/__init__.py,sha256=CiOZMcyPTAaKtEuMZUhEda_Ad4_RUhmIstB-A3iuOJ
 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=GWPRpufg5larcKGNwlRtN22xVs9fp3ptMu11rQeySX8,3141
+gaussian_splatting/prepare.py,sha256=rgwdDhPU-sS57ZgTLKwzognbS1hfR2JtjZqlC4FqPbI,3241
 gaussian_splatting/render.py,sha256=MSvJiOpyIr5IgJeaB6xj2GDcK4Jb6mcj7CWVXb-ajIk,5988
 gaussian_splatting/train.py,sha256=a3vHZf5sX5eVsxv394pmwTSon1mPevQVPGAWwYHqaY8,5386
 gaussian_splatting/dataset/__init__.py,sha256=-runuT-61P0YVpfV_WXqwUZM1oY0N012YH13Bt3rzSU,138
@@ -12,12 +12,12 @@ gaussian_splatting/dataset/colmap/__init__.py,sha256=YEYT2k2WJSqrkkZq4KAJYS9UMgq
 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.cp312-win_amd64.pyd,sha256=aj88_ALoC0xx2JFvlS2mQIrEETMZOJY31u-oDn7o6MY,1300992
+gaussian_splatting/diff_gaussian_rasterization/_C.cp312-win_amd64.pyd,sha256=pfC4F0RO4kLxzGvBxIN67JDrk_nUs-I_oMlf-iLpOa8,1300992
 gaussian_splatting/diff_gaussian_rasterization/__init__.py,sha256=a9D0IZiPx-Mk1795hSq54T-NYT4MtEN_MZrxeMhw0Eo,6705
-gaussian_splatting/simple_knn/_C.cp312-win_amd64.pyd,sha256=KoC2pG2Ie1pm9XaIeRH1UD3gWnC1gk3GrLOvbidg_4Y,1168896
+gaussian_splatting/simple_knn/_C.cp312-win_amd64.pyd,sha256=I1jbI1-fs3z6PvzBI_DPu-i3fUxIWd1J3VGp98Q5c00,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=A2nQDgFbGod4orVbLwKqbCm_FBE-HNjtfQjCc9z3xGk,3713
+gaussian_splatting/trainer/abc.py,sha256=_gcqmEobhSOdZnMyNb2oKS6cZJ-Mg3oYL4xJ5Y3_oic,4262
+gaussian_splatting/trainer/base.py,sha256=fngLruQ9hMSNLFbc_5woG7jm6cidpoZ0dzk_zImRaE4,4851
 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
@@ -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.18.1.dist-info/licenses/LICENSE.md,sha256=bMuRQKn0u485mx8JBBTJ5Simc-aWHaQsxmoB6jsg5oE,4752
-gaussian_splatting-1.18.1.dist-info/METADATA,sha256=IJjDkFlUcJxSIdQKjjW6KhZZHmRKgHHdLRFCTuVI_OQ,17183
-gaussian_splatting-1.18.1.dist-info/WHEEL,sha256=8UP9x9puWI0P1V_d7K2oMTBqfeLNm21CTzZ_Ptr0NXU,101
-gaussian_splatting-1.18.1.dist-info/top_level.txt,sha256=uaYrPYXRHhpybgCnsoazTcdhpzZGnLT_vd5eoRzBWWI,19
-gaussian_splatting-1.18.1.dist-info/RECORD,,
+gaussian_splatting-1.19.1.dist-info/licenses/LICENSE.md,sha256=bMuRQKn0u485mx8JBBTJ5Simc-aWHaQsxmoB6jsg5oE,4752
+gaussian_splatting-1.19.1.dist-info/METADATA,sha256=UDNOnxvs-3k2zfDht3UiojA_pjcTsAvB8j-LW0ROObQ,17183
+gaussian_splatting-1.19.1.dist-info/WHEEL,sha256=8UP9x9puWI0P1V_d7K2oMTBqfeLNm21CTzZ_Ptr0NXU,101
+gaussian_splatting-1.19.1.dist-info/top_level.txt,sha256=uaYrPYXRHhpybgCnsoazTcdhpzZGnLT_vd5eoRzBWWI,19
+gaussian_splatting-1.19.1.dist-info/RECORD,,