multiple improvements

easyvolcap/dataloaders/datasets/volumetric_video_dataset.py

@@ -645,7 +645,6 @@ def carve_using_bytes(H, W, K, R, T, latent_index):
             log(magenta(f'Accumulated vhull bound of frames: {self.frame_sample}: '),
                 line(torch.stack([self.vhull_bounds.min(dim=0)[0][0],
                                   self.vhull_bounds.max(dim=0)[0][1]])))
-        log(self.vhull_only)
         if self.vhull_only:
             exit(0)
 

easyvolcap/utils/loss_utils.py

@@ -26,7 +26,8 @@ class ImgLossType(Enum):
     L1 = auto()
     L2 = auto()
     SSIM = auto()
-
+    MSSSIM = auto()
+    WL1 = auto()
 
 class DptLossType(Enum):
     SMOOTHL1 = auto()
@@ -413,6 +414,10 @@ def l1(x: torch.Tensor, y: torch.Tensor):
     return l1_reg(x - y)
 
 
+def wl1(x: torch.Tensor, y: torch.Tensor, w: torch.Tensor):
+    return l1_reg(w * (x - y))
+
+
 def l2(x: torch.Tensor, y: torch.Tensor):
     return l2_reg(x - y)
 
@@ -638,6 +643,11 @@ def compute_ssim(x: torch.Tensor, y: torch.Tensor):
     return ssim(x, y, data_range=1.0, win_size=11, win_sigma=1.5, K=(0.01, 0.03))
 
 
+def compute_msssim(x: torch.Tensor, y: torch.Tensor):
+    from pytorch_msssim import ms_ssim
+    return ms_ssim(x, y, data_range=1.0, win_size=11, win_sigma=1.5, K=(0.01, 0.03))
+
+
 # from MonoSDF
 def compute_scale_and_shift(prediction, target, mask):
     # System matrix: A = [[a_00, a_01], [a_10, a_11]]