fix app_mask

daniel03c1 · Nov 4, 2022 · 034619c · 034619c
1 parent 8111899
commit 034619c
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/TensoRF/models/tensoRF.py b/TensoRF/models/tensoRF.py
@@ -137,6 +137,14 @@ def compute_appfeature(self, points):
             plane = min_max_quantize(self.app_plane[idx], self.grid_bit)
             line = min_max_quantize(self.app_line[idx], self.grid_bit)
 
+            if self.use_mask:
+                mask = torch.sigmoid(self.app_plane_mask[idx])
+                plane = (plane * (mask >= 0.5) - plane * mask).detach() \
+                      + plane * mask
+                mask = torch.sigmoid(self.app_line_mask[idx])
+                line = (line * (mask >= 0.5) - line * mask).detach() \
+                     + line * mask
+
             if self.use_dwt:
                 plane = inverse(plane, self.dwt_level)