Fix tensor dimension mismatch in segmentation

- Fix RuntimeError: output with shape [1024, 1024] doesn't match broadcast shape [1, 1024, 1024] - Add dimension check and squeeze operation in _create_colored_segmentation - Ensure mask is 2D before broadcasting with colored_image - This fixes the tensor broadcasting issue in ConceptSegmentationNode
2025-09-17 14:10:34 +09:00
parent 9b8728edbe
commit 1040d02ea0
1 changed files with 4 additions and 0 deletions
@@ -291,6 +291,10 @@ class ConceptSegmentationNode:
                color = colors[concept]
                mask = (segmentation_mask == i + 1).float()
                
+                # Ensure mask is 2D for broadcasting
+                if len(mask.shape) == 3:
+                    mask = mask.squeeze(0)  # Remove batch dimension
+                
                for c in range(3):
                    colored_image[0, :, :, c] += mask * color[c]