fix category mask

app.py

@@ -19,7 +19,8 @@ def segment(input_image, category):
     image = mp.Image(image_format=mp.ImageFormat.SRGB, data=np.asarray(input_image))
     segmentation_result = segmenter.segment(image)
     category_mask = segmentation_result.category_mask
-    category_mask = category_mask == MASK_CATEGORY.index(category)
+    category_mask_np = category_mask.numpy_view()
+    target_mask = category_mask_np == MASK_CATEGORY.index(category)
 
     # Generate solid color images for showing the output segmentation mask.
     image_data = image.numpy_view()
@@ -28,7 +29,7 @@ def segment(input_image, category):
     bg_image = np.zeros(image_data.shape, dtype=np.uint8)
     bg_image[:] = BG_COLOR
 
-    dilated_mask = binary_dilation(category_mask.numpy_view(), iterations=4)
+    dilated_mask = binary_dilation(target_mask, iterations=4)
     condition = np.stack((dilated_mask,) * 3, axis=-1) > 0.2
 
     output_image = np.where(condition, fg_image, bg_image)