Fix mask redundant dims

lldacing · Sep 24, 2024 · 03ff6c2 · 03ff6c2
1 parent e3a7de4
commit 03ff6c2
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/birefnetNode.py b/birefnetNode.py
@@ -199,11 +199,14 @@ def rem_bg(self, model, images):
             image = apply_mask_to_image(image.cpu(), mask.cpu())
 
             _images.append(image)
-            _masks.append(mask)
+            _masks.append(mask.squeeze(0))
 
         out_images = torch.cat(_images, dim=0)
         out_masks = torch.cat(_masks, dim=0)
 
+        if out_masks.shape[0] == 1:
+            out_masks = out_masks.squeeze(0)
+
         return out_images, out_masks