debug

app.py

@@ -26,12 +26,15 @@ def visualize_instance_seg_mask(mask):
     labels = np.unique(mask)
     label2color = {label: (random.randint(0, 1), random.randint(0, 255), random.randint(0, 255)) for label in labels if label > 0}
     label2color[0] = (0, 0, 0)
-    # for label in labels:
-    #     image[mask==label, :]
-    for i in range(image.shape[0]):
-        for j in range(image.shape[1]):
-            image[i, j, :] = label2color[mask[i, j]]
-    image = image / 255
+    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
 
 
@@ -141,11 +144,12 @@ def predict(img, threshold=0.5):
     else:
         img_data = io.imread(img_name)
     seg_labels = get_seg(preprocess(img_data), 'swinunetr', './best_Dice_model.pth', float(threshold))
-    # seg_rgb = visualize_instance_seg_mask(seg_labels)
-    seg_rgb = seg_labels
+    seg_rgb = visualize_instance_seg_mask(seg_labels)
 
     tif.imwrite(join(os.getcwd(), 'segmentation.tiff'), seg_labels, compression='zlib')
 
+    print(np.max(img_data), np.min(img_data))
+    print(np.max(seg_rgb), np.min(seg_rgb))
     return img_data, seg_rgb, join(os.getcwd(), 'segmentation.tiff')