Small fix.

luxonis · Sep 9, 2024 · 97895eb · 97895eb
1 parent ec57834
commit 97895eb
Showing 1 changed file with 5 additions and 11 deletions.
diff --git a/depthai_nodes/ml/parsers/segmentation.py b/depthai_nodes/ml/parsers/segmentation.py
@@ -85,17 +85,11 @@ def run(self):
                         segmentation_mask,
                     )
                 )
-                class_map = (
-                    np.max(segmentation_mask, axis=0)
-                    .reshape(segmentation_mask.shape[1], segmentation_mask.shape[2], 1)
-                    .astype(np.uint8)
-                )
-            else:
-                class_map = (
-                    np.argmax(segmentation_mask, axis=0)
-                    .reshape(segmentation_mask.shape[1], segmentation_mask.shape[2], 1)
-                    .astype(np.uint8)
-                )
+            class_map = (
+                np.argmax(segmentation_mask, axis=0)
+                .reshape(segmentation_mask.shape[1], segmentation_mask.shape[2], 1)
+                .astype(np.uint8)
+            )
 
             imgFrame = create_segmentation_message(class_map)
             imgFrame.setTimestamp(output.getTimestamp())