gaussian-splatting 1.18.1__tar.gz → 1.19.0__tar.gz

{gaussian_splatting-1.18.1/gaussian_splatting.egg-info → gaussian_splatting-1.19.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: gaussian_splatting
-Version: 1.18.1
+Version: 1.19.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.18.1 → gaussian_splatting-1.19.0}/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-1.18.1 → gaussian_splatting-1.19.0}/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-1.18.1 → gaussian_splatting-1.19.0}/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,9 +103,6 @@ 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)

{gaussian_splatting-1.18.1 → gaussian_splatting-1.19.0/gaussian_splatting.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: gaussian_splatting
-Version: 1.18.1
+Version: 1.19.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.18.1 → gaussian_splatting-1.19.0}/setup.py

@@ -44,7 +44,7 @@ if os.name == 'nt':
 
 setup(
     name="gaussian_splatting",
-    version='1.18.1',
+    version='1.19.0',
     author='yindaheng98',
     author_email='yindaheng98@gmail.com',
     url='https://github.com/yindaheng98/gaussian-splatting',