init

app.py

@@ -18,6 +18,7 @@ from monai.inferers import sliding_window_inference
 from unetr2d import UNETR2D
 import time
 from skimage import io, segmentation, morphology, measure, exposure
+import tifffile as tif
 
 
 def visualize_instance_seg_mask(mask):
@@ -130,19 +131,27 @@ def get_seg(pre_img_data, model_name, custom_model_path, threshold):
     return test_pred_mask
 
 
-def predict(img, threshold):
-    print('##########', img)
-    seg_labels = get_seg(preprocess(img), 'swinunetr', './best_Dice_model.pth', float(threshold))
-    seg_labels = visualize_instance_seg_mask(seg_labels)
-    return seg_labels, './text.tif'
+def predict(img, threshold=0.5):
+    print('##########', img.name)
+    img_name = img.name
+    if img_name.endswith('.tif') or img_name.endswith('.tiff'):
+        img_data = tif.imread(img_name)
+    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)
+
+    tif.imwrite(join(os.getcwd(), 'segmentation.tiff'), seg_labels, compression='zlib')
+
+    return img_data, seg_rgb, join(os.getcwd(), 'segmentation.tiff')
 
 
 demo = gr.Interface(
     predict, 
     # inputs=[gr.Image()],
     # inputs="file",
-    inputs=[gr.File(), gr.Number()]
-    outputs=[gr.Image(), gr.File()],
+    inputs=[gr.File(label="input image"), gr.Slider(minimum=0, maximum=1, value=0.5, step=0.1, label="threshold")],
+    outputs=[gr.Image(label="image"), gr.Image(label="segmentation"), gr.File(label="download segmentation")],
     title="NeurIPS CellSeg Demo",
     # examples=[["cell_00225.png"]]
 )