Create app.py

app.py

@@ -0,0 +1,207 @@
+import gradio as gr
+import os
+os.environ["CUDA_VISIBLE_DEVICES"] = "-1"  # Force CPU if needed
+import torch
+import numpy as np
+from PIL import Image
+from PIL import Image as PILImage
+from pathlib import Path
+import matplotlib.pyplot as plt
+import io
+from skimage.io import imread
+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
+from cellpose import models as cellpose_models, io as cellpose_io, plot as cellpose_plot
+
+# Load SegFormer
+from transformers import SegformerImageProcessor, SegformerForSemanticSegmentation
+processor_segformer = SegformerImageProcessor(do_reduce_labels=False)
+model_segformer = SegformerForSemanticSegmentation.from_pretrained(
+    "nvidia/segformer-b0-finetuned-ade-512-512",
+    num_labels=8,
+    ignore_mismatched_sizes=True
+)
+model_segformer.load_state_dict(torch.load("trained_model_200.pt", map_location="cpu"))
+model_segformer.eval()
+
+# StarDist model
+model_stardist = StarDist2D.from_pretrained('2D_versatile_fluo')
+
+# Cellpose model
+model_cellpose = cellpose_models.CellposeModel(gpu=False)
+
+# Handle SegFormer prediction
+def infer_segformer(image):
+    image = image.convert("RGB")
+    inputs = processor_segformer(images=image, return_tensors="pt")
+    with torch.no_grad():
+        logits = model_segformer(**inputs).logits
+    pred_mask = torch.argmax(logits, dim=1)[0].cpu().numpy()
+
+    # Colorize
+    colors = np.array([[0,0,0], [255,0,0], [0,255,0], [0,0,255], [255,255,0], [255,0,255], [0,255,255], [128,128,128]])
+    color_mask = np.zeros((pred_mask.shape[0], pred_mask.shape[1], 3), dtype=np.uint8)
+    for c in range(8):
+        color_mask[pred_mask == c] = colors[c]
+    return image, Image.fromarray(color_mask)
+
+# Handle StarDist prediction
+def infer_stardist(image):
+    image_gray = rgb2gray(np.array(image)) if image.mode == 'RGB' else np.array(image)
+    labels, _ = model_stardist.predict_instances(normalize(image_gray))
+    overlay = render_label(labels, img=image_gray)
+    overlay = (overlay[..., :3] * 255).astype(np.uint8)
+    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)
+# Handle Cellpose prediction
+def infer_cellpose(image, temp_dir="temp_cellpose"):
+    os.makedirs(temp_dir, exist_ok=True)
+    input_path = os.path.join(temp_dir, "input_image.tif")
+    output_overlay = os.path.join(temp_dir, "overlay.png")
+
+    # Save image
+    image.save(input_path)
+    img = cellpose_io.imread(input_path)
+    masks, flows, styles = model_cellpose.eval(img, batch_size=1)
+
+    fig = plt.figure(figsize=(12,5))
+    cellpose_plot.show_segmentation(fig, img, masks, flows[0])
+    plt.tight_layout()
+    fig.savefig(output_overlay)
+    plt.close(fig)
+
+    return image, Image.open(output_overlay)
+
+# Wrapper function
+def segment(model_name, image):
+    # Gradio passes a PIL.Image without filename attribute
+    # Try to check format if available, else skip check
+    ext = None
+    if hasattr(image, 'format') and image.format is not None:
+        ext = image.format.lower()
+    if model_name == "Cellpose":
+        # Accept only TIFF images for Cellpose
+        if ext not in ["tiff", "tif", None]:
+            return None, f"❌ Cellpose only supports `.tif` or `.tiff` images."
+    # ...existing code...
+    if model_name == "SegFormer":
+        return infer_segformer(image)
+    elif model_name == "StarDist":
+        return infer_stardist(image)
+    # elif model_name == "MEDIAR":
+    #     return infer_mediar(image)
+    elif model_name == "Cellpose":
+        return infer_cellpose(image)
+    else:
+        return None, f"❌ Unknown model: {model_name}"
+
+with gr.Blocks(title="Cell Segmentation Explorer") as app:
+    gr.Markdown("## Cell Segmentation Explorer")
+    gr.Markdown("Choose a segmentation model, upload an appropriate image, and view the predicted mask.")
+
+    with gr.Row():
+        with gr.Column():
+            model_dropdown = gr.Dropdown(
+                choices=["SegFormer", "StarDist", "Cellpose"],
+                label="Select Segmentation Model",
+                value="SegFormer"
+            )
+            image_input = gr.Image(type="pil", label="Uploaded Image")
+            description_box = gr.Markdown("Accepted formats: `.png`, `.jpg`, `.tif`, `.tiff`.")
+            submit_btn = gr.Button("Submit")
+            clear_btn = gr.Button("Clear")
+        with gr.Column():
+            output_image = gr.Image(label="Segmentation Result")
+
+    def handle_submit(model_name, img):
+        if img is None:
+            return None
+        _, result = segment(model_name, img)  # Only return the mask (segmentation result)
+        return result
+
+    submit_btn.click(
+        fn=handle_submit,
+        inputs=[model_dropdown, image_input],
+        outputs=output_image
+    )
+
+    clear_btn.click(
+        lambda: [None, None],
+        inputs=None,
+        outputs=[image_input, output_image]
+    )
+
+    # === SAMPLE IMAGES SECTION ===
+    gr.Markdown("---")
+    gr.Markdown("### Sample Images (click to use as input)")
+
+    # Original and resized thumbnails
+    original_sample_paths = [
+        "img1.png",
+        "img2.png",
+        "img3.png"
+    ]
+
+    resized_sample_paths = []
+    for idx, p in enumerate(original_sample_paths):
+        img = PILImage.open(p).resize((128, 128))
+        temp_path = f"/tmp/sample_resized_{idx}.png"
+        img.save(temp_path)
+        resized_sample_paths.append(temp_path)
+
+    sample_image_components = []
+    with gr.Row():
+        for i, img_path in enumerate(resized_sample_paths):
+            def load_full_image(idx=i):  # Capture loop index properly
+                return PILImage.open(original_sample_paths[idx])
+
+            sample_img = gr.Image(value=img_path, type="pil", interactive=True, show_label=False)
+            sample_img.select(
+                fn=load_full_image,
+                inputs=[],
+                outputs=image_input
+            )
+            sample_image_components.append(sample_img)
+
+
+app.launch()