Debug app.py

app.py

@@ -37,6 +37,7 @@ from skimage.util.dtype import dtype_range
 from skimage._shared.utils import _supported_float_type
 from scipy.ndimage import find_objects, binary_fill_holes
 
+import random
 
 ########################### Data Loading Modules #########################################################
 DTYPE_RANGE = dtype_range.copy()
@@ -1102,6 +1103,22 @@ def compute_masks(
     
     return mask, p
 
+def visualize_instance_seg_mask(mask):
+    image = np.zeros((mask.shape[0], mask.shape[1], 3))
+    labels = np.unique(mask)
+    label2color = {label: (random.randint(50, 255), random.randint(50, 255), random.randint(50, 255)) for label in labels if label > 0}
+    label2color[0] = (0, 0, 0)
+    for label in labels:
+        image[mask==label, :] = label2color[label]
+    # for i in range(image.shape[0]):
+    #     for j in range(image.shape[1]):
+    #         if np.max(label2color[mask[i, j]]) > 0:
+    #             print('####', np.max(label2color[mask[i, j]]), np.min(label2color[mask[i, j]]))
+    #         image[i, j, :] = label2color[mask[i, j]]
+    # image = image / 255
+    image = image.astype(np.uint8)
+    return image
+
 def predict(img):
     # Dataset parameters
     ### for huggingface space
@@ -1239,7 +1256,7 @@ def predict(img):
     tif.imwrite(file_path, pred_mask, compression="zlib")
     print(np.max(pred_mask))
     # return img_data, seg_rgb, join(os.getcwd(), 'segmentation.tiff')
-    return origin_img, pred_mask, file_path
+    return origin_img, visualize_instance_seg_mask(pred_mask), file_path
 
 demo = gr.Interface(
     predict,