Update app.py

app.py

@@ -13,8 +13,8 @@ from skimage.color import rgb2gray
 from csbdeep.utils import normalize
 from stardist.models import StarDist2D
 from stardist.plot import render_label
-#import MEDIARFormer
-#import Predictor
+import MEDIARFormer
+import Predictor
 from cellpose import models as cellpose_models, io as cellpose_io, plot as cellpose_plot
 
 # Load SegFormer
@@ -58,41 +58,41 @@ def infer_stardist(image):
     return image, Image.fromarray(overlay)
 
 # Handle MEDIAR prediction
-# def infer_mediar(image, temp_dir="temp_mediar"):
-#     os.makedirs(temp_dir, exist_ok=True)
-#     input_path = os.path.join(temp_dir, "input_image.tiff")
-#     output_path = os.path.join(temp_dir, "input_image_label.tiff")
-
-#     image.save(input_path)
-
-#     model_args = {
-#         "classes": 3,
-#         "decoder_channels": [1024, 512, 256, 128, 64],
-#         "decoder_pab_channels": 256,
-#         "encoder_name": 'mit_b5',
-#         "in_channels": 3
-#     }
-
-#     model = MEDIARFormer(**model_args)
-#     weights = torch.load("from_phase1.pth", map_location="cpu")
-#     model.load_state_dict(weights, strict=False)
-#     model.eval()
-
-#     predictor = Predictor(model, "cpu", temp_dir, temp_dir, algo_params={"use_tta": False})
-#     predictor.img_names = ["input_image.tiff"]
-#     _ = predictor.conduct_prediction()
-
-#     pred = imread(output_path)
-#     fig, ax = plt.subplots(figsize=(6, 6))
-#     ax.imshow(pred, cmap="cividis")
-#     ax.axis("off")
-
-#     buf = io.BytesIO()
-#     plt.savefig(buf, format="png")
-#     plt.close()
-#     buf.seek(0)
-
-#     return image, Image.open(buf)
+def infer_mediar(image, temp_dir="temp_mediar"):
+    os.makedirs(temp_dir, exist_ok=True)
+    input_path = os.path.join(temp_dir, "input_image.tiff")
+    output_path = os.path.join(temp_dir, "input_image_label.tiff")
+
+    image.save(input_path)
+
+    model_args = {
+        "classes": 3,
+        "decoder_channels": [1024, 512, 256, 128, 64],
+        "decoder_pab_channels": 256,
+        "encoder_name": 'mit_b5',
+        "in_channels": 3
+    }
+
+    model = MEDIARFormer(**model_args)
+    weights = torch.load("from_phase1.pth", map_location="cpu")
+    model.load_state_dict(weights, strict=False)
+    model.eval()
+
+    predictor = Predictor(model, "cpu", temp_dir, temp_dir, algo_params={"use_tta": False})
+    predictor.img_names = ["input_image.tiff"]
+    _ = predictor.conduct_prediction()
+
+    pred = imread(output_path)
+    fig, ax = plt.subplots(figsize=(6, 6))
+    ax.imshow(pred, cmap="cividis")
+    ax.axis("off")
+
+    buf = io.BytesIO()
+    plt.savefig(buf, format="png")
+    plt.close()
+    buf.seek(0)
+
+    return image, Image.open(buf)
 # Handle Cellpose prediction
 def infer_cellpose(image, temp_dir="temp_cellpose"):
     os.makedirs(temp_dir, exist_ok=True)
@@ -128,8 +128,8 @@ def segment(model_name, image):
         return infer_segformer(image)
     elif model_name == "StarDist":
         return infer_stardist(image)
-    # elif model_name == "MEDIAR":
-    #     return infer_mediar(image)
+    elif model_name == "MEDIAR":
+        return infer_mediar(image)
     elif model_name == "Cellpose":
         return infer_cellpose(image)
     else:
@@ -142,7 +142,7 @@ with gr.Blocks(title="Cell Segmentation Explorer") as app:
     with gr.Row():
         with gr.Column():
             model_dropdown = gr.Dropdown(
-                choices=["SegFormer", "StarDist", "Cellpose"],
+                choices=["SegFormer", "StarDist", "MEDIAR","Cellpose"],
                 label="Select Segmentation Model",
                 value="SegFormer"
             )