From 67adbfdcae9c64e2e0355cca103d1ec7d676db7b Mon Sep 17 00:00:00 2001
From: ihmily <114978440+ihmily@users.noreply.github.com>
Date: Wed, 24 Dec 2025 17:03:20 +0800
Subject: [PATCH] =?UTF-8?q?Fix=20mask=20broadcasting=20error=20for=20Media?=
 =?UTF-8?q?Pipe=20=E2=89=A50.10.30?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 py/person_mask_Ultra.py    | 9 ++++++++-
 py/person_mask_ultra_v2.py | 9 ++++++++-
 2 files changed, 16 insertions(+), 2 deletions(-)

diff --git a/py/person_mask_Ultra.py b/py/person_mask_Ultra.py
index 4f803ad..6291d8f 100644
--- a/py/person_mask_Ultra.py
+++ b/py/person_mask_Ultra.py
@@ -115,7 +115,14 @@ class PersonMaskUltra:
                     mask_arrays.append(mask_background_array)
                 else:
                     for i, mask in enumerate(masks):
-                        condition = np.stack((mask.numpy_view(),) * image_shape[-1], axis=-1) > confidence
+                        mask_2d = mask.numpy_view()
+                        if mask_2d.ndim == 3 and mask_2d.shape[2] == 1:
+                            mask_2d = mask_2d.squeeze(axis=2)
+                        elif mask_2d.ndim != 2:
+                            raise ValueError(f"Unexpected mask shape: {mask_2d.shape}")
+                        condition = np.stack((mask_2d,) * image_shape[-1], axis=-1) > confidence
+                        if condition.ndim == 4 and condition.shape[2] == 1:
+                            condition = condition.squeeze(2)
                         mask_array = np.where(condition, mask_foreground_array, mask_background_array)
                         mask_arrays.append(mask_array)
                 # Merge our masks taking the maximum from each
diff --git a/py/person_mask_ultra_v2.py b/py/person_mask_ultra_v2.py
index 8cdbd64..a725842 100644
--- a/py/person_mask_ultra_v2.py
+++ b/py/person_mask_ultra_v2.py
@@ -131,7 +131,14 @@ class PersonMaskUltraV2:
                     mask_arrays.append(mask_background_array)
                 else:
                     for i, mask in enumerate(masks):
-                        condition = np.stack((mask.numpy_view(),) * image_shape[-1], axis=-1) > confidence
+                        mask_2d = mask.numpy_view()
+                        if mask_2d.ndim == 3 and mask_2d.shape[2] == 1:
+                            mask_2d = mask_2d.squeeze(axis=2)
+                        elif mask_2d.ndim != 2:
+                            raise ValueError(f"Unexpected mask shape: {mask_2d.shape}")
+                        condition = np.stack((mask_2d,) * image_shape[-1], axis=-1) > confidence
+                        if condition.ndim == 4 and condition.shape[2] == 1:
+                            condition = condition.squeeze(2)
                         mask_array = np.where(condition, mask_foreground_array, mask_background_array)
                         mask_arrays.append(mask_array)
                 # Merge our masks taking the maximum from each