use mask in 2dgs (#497)

nerfstudio-project · Nov 23, 2024 · 7951619 · 7951619
1 parent 0f42139
commit 7951619
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/examples/simple_trainer_2dgs.py b/examples/simple_trainer_2dgs.py
@@ -577,6 +577,10 @@ def train(self):
                 step=step,
                 info=info,
             )
+            masks = data["mask"].to(device) if "mask" in data else None
+            if masks is not None:
+                pixels = pixels * masks[..., None]
+                colors = colors * masks[..., None]
 
             # loss
             l1loss = F.l1_loss(colors, pixels)