Spaces:
Running
Running
| 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 | |
| from MEDIARFormer import MEDIARFormer | |
| from Predictor 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", "MEDIAR", "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() |