gaussian-splatting 1.19.0__cp311-cp311-win_amd64.whl → 1.19.2__cp311-cp311-win_amd64.whl

gaussian_splatting/render.py

@@ -15,8 +15,8 @@ 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=True) -> Tuple[CameraDataset, GaussianModel]:
-    dataset = prepare_dataset(source=source, device=device, trainable_camera=trainable_camera, load_camera=load_camera, load_mask=True, load_depth=load_depth)
+        load_mask=True, load_depth=True) -> Tuple[CameraDataset, GaussianModel]:
+    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)
     return dataset, gaussians
 
@@ -57,13 +57,16 @@ def rendering(
     gt_path = os.path.join(save, "gt")
     makedirs(render_path, exist_ok=True)
     makedirs(gt_path, exist_ok=True)
-    pbar = tqdm(dataset, desc="Rendering progress")
+    pbar = tqdm(dataset, dynamic_ncols=True, desc="Rendering")
     with open(os.path.join(save, "quality.csv"), "w") as f:
         f.write("name,psnr,ssim,lpips\n")
     for idx, camera in enumerate(pbar):
         out = gaussians(camera)
         rendering = out["render"]
         gt = camera.ground_truth_image
+        if camera.ground_truth_image_mask is not None:
+            gt *= camera.ground_truth_image_mask
+            rendering *= camera.ground_truth_image_mask
         psnr_value = psnr(rendering, gt).mean().item()
         ssim_value = ssim(rendering, gt).mean().item()
         lpips_value = lpips(rendering, gt).mean().item()
@@ -96,6 +99,7 @@ if __name__ == "__main__":
     parser.add_argument("--load_camera", default=None, type=str)
     parser.add_argument("--mode", choices=["base", "camera"], default="base")
     parser.add_argument("--device", default="cuda", type=str)
+    parser.add_argument("--no_image_mask", action="store_true")
     parser.add_argument("--no_rescale_depth_gt", action="store_true")
     parser.add_argument("--save_depth_pcd", action="store_true")
     args = parser.parse_args()
@@ -104,5 +108,6 @@ 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=args.save_depth_pcd)
+            load_ply=load_ply, load_camera=args.load_camera,
+            load_mask=not args.no_image_mask, 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

@@ -31,7 +31,7 @@ def save_cfg_args(destination: str, sh_degree: int, source: str):
 
 def training(dataset: CameraDataset, gaussians: GaussianModel, trainer: AbstractTrainer, destination: str, iteration: int, save_iterations: List[int], device: str):
     shutil.rmtree(os.path.join(destination, "point_cloud"), ignore_errors=True)  # remove the previous point cloud
-    pbar = tqdm(range(1, iteration+1))
+    pbar = tqdm(range(1, iteration+1), dynamic_ncols=True, desc="Training")
     epoch = list(range(len(dataset)))
     epoch_psnr = torch.empty(3, 0, device=device)
     ema_loss_for_log = 0.0
@@ -45,8 +45,13 @@ def training(dataset: CameraDataset, gaussians: GaussianModel, trainer: Abstract
         idx = epoch[epoch_idx]
         loss, out = trainer.step(dataset[idx])
         with torch.no_grad():
+            ground_truth_image = dataset[idx].ground_truth_image
+            rendered_image = out["render"]
+            if dataset[idx].ground_truth_image_mask is not None:
+                ground_truth_image *= dataset[idx].ground_truth_image_mask
+                rendered_image *= dataset[idx].ground_truth_image_mask
             ema_loss_for_log = 0.4 * loss.item() + 0.6 * ema_loss_for_log
-            epoch_psnr = torch.concat([epoch_psnr, psnr(out["render"], dataset[idx].ground_truth_image)], dim=1)
+            epoch_psnr = torch.concat([epoch_psnr, psnr(rendered_image, ground_truth_image)], dim=1)
             if step % 10 == 0:
                 pbar.set_postfix({'epoch': step // len(dataset), 'loss': ema_loss_for_log, 'psnr': avg_psnr_for_log, 'n': gaussians._xyz.shape[0]})
         if step in save_iterations:

gaussian_splatting/trainer/base.py

@@ -105,7 +105,9 @@ class BaseTrainer(AbstractTrainer):
                 gt = gt * mask.unsqueeze(0)
             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.19.0.dist-info → gaussian_splatting-1.19.2.dist-info}/METADATA

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

{gaussian_splatting-1.19.0.dist-info → gaussian_splatting-1.19.2.dist-info}/RECORD

@@ -3,8 +3,8 @@ gaussian_splatting/camera.py,sha256=vo7mu6lyFpIhDqOAgNiJuPan8t_nDJn5cJkAYygLFcA,
 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=rgwdDhPU-sS57ZgTLKwzognbS1hfR2JtjZqlC4FqPbI,3241
-gaussian_splatting/render.py,sha256=MSvJiOpyIr5IgJeaB6xj2GDcK4Jb6mcj7CWVXb-ajIk,5988
-gaussian_splatting/train.py,sha256=a3vHZf5sX5eVsxv394pmwTSon1mPevQVPGAWwYHqaY8,5386
+gaussian_splatting/render.py,sha256=RhZoILWxtkuDEhiL2cQsvZWNBmkyDAROYdvrQaFhPkc,6295
+gaussian_splatting/train.py,sha256=RdPg53jf2C8SBosTod5GmoDcIIk83hhdAIcd1prvTcM,5735
 gaussian_splatting/dataset/__init__.py,sha256=-runuT-61P0YVpfV_WXqwUZM1oY0N012YH13Bt3rzSU,138
 gaussian_splatting/dataset/camera_trainable.py,sha256=Kd8v-_ZJ9dLIQ2QyVOXbmouYf5QjbgOgHNRHVpkgCms,5041
 gaussian_splatting/dataset/dataset.py,sha256=0tmIZ5P7kOEdABiEAXPznkRN91e5rcT5VsAzOLoOuEM,2392
@@ -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.cp311-win_amd64.pyd,sha256=oyjIZwIYxtZpEXbhpWSJmj3U3jA32LlKd9OAxLAgI4U,1295360
+gaussian_splatting/diff_gaussian_rasterization/_C.cp311-win_amd64.pyd,sha256=n3TrGEe3-zOhJj-tNZiY0CBgPns4_PpTYa5FTcO8d0E,1295360
 gaussian_splatting/diff_gaussian_rasterization/__init__.py,sha256=a9D0IZiPx-Mk1795hSq54T-NYT4MtEN_MZrxeMhw0Eo,6705
-gaussian_splatting/simple_knn/_C.cp311-win_amd64.pyd,sha256=ll9GppZDR03FGo6SiP8_Rm24BcVGuhve11Bj_xX-XlA,1164288
+gaussian_splatting/simple_knn/_C.cp311-win_amd64.pyd,sha256=SVEJlbMeELkNeeSb_rxSbisEhMFEA3fmWWPOaHecDRU,1164288
 gaussian_splatting/trainer/__init__.py,sha256=962fEY8A0spSQn5de_d_LkPOjA1PYKrLbuAkxwZo7mI,940
 gaussian_splatting/trainer/abc.py,sha256=_gcqmEobhSOdZnMyNb2oKS6cZJ-Mg3oYL4xJ5Y3_oic,4262
-gaussian_splatting/trainer/base.py,sha256=GfPifKjelRoCOXLJP6_tk1jgFWTJloBaCyUiVCcE9u4,4724
+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.19.0.dist-info/licenses/LICENSE.md,sha256=bMuRQKn0u485mx8JBBTJ5Simc-aWHaQsxmoB6jsg5oE,4752
-gaussian_splatting-1.19.0.dist-info/METADATA,sha256=L0gv8LTr-DW4WFGl_jIFl568SuvIqlhsp863yh5LSZI,17183
-gaussian_splatting-1.19.0.dist-info/WHEEL,sha256=JLOMsP7F5qtkAkINx5UnzbFguf8CqZeraV8o04b0I8I,101
-gaussian_splatting-1.19.0.dist-info/top_level.txt,sha256=uaYrPYXRHhpybgCnsoazTcdhpzZGnLT_vd5eoRzBWWI,19
-gaussian_splatting-1.19.0.dist-info/RECORD,,
+gaussian_splatting-1.19.2.dist-info/licenses/LICENSE.md,sha256=bMuRQKn0u485mx8JBBTJ5Simc-aWHaQsxmoB6jsg5oE,4752
+gaussian_splatting-1.19.2.dist-info/METADATA,sha256=LlY7Xiln84Uz24GZUl14IqxRsWL3nUXOe0jLZHNdAdY,17183
+gaussian_splatting-1.19.2.dist-info/WHEEL,sha256=JLOMsP7F5qtkAkINx5UnzbFguf8CqZeraV8o04b0I8I,101
+gaussian_splatting-1.19.2.dist-info/top_level.txt,sha256=uaYrPYXRHhpybgCnsoazTcdhpzZGnLT_vd5eoRzBWWI,19
+gaussian_splatting-1.19.2.dist-info/RECORD,,