Update app.py

app.py

@@ -1,7 +1,8 @@
 import gradio as gr
 import os
-os.environ["CUDA_VISIBLE_DEVICES"] = "-1"  # Force CPU if needed
 import torch
+import tensorflow as tf
+tf.config.set_visible_devices([], 'GPU')
 import numpy as np
 from PIL import Image
 from PIL import Image as PILImage
@@ -13,31 +14,38 @@ 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 train_tools.models import MEDIARFormer
+from core.MEDIAR import Predictor
 from cellpose import models as cellpose_models, io as cellpose_io, plot as cellpose_plot
+from transformers import SegformerImageProcessor, SegformerForSemanticSegmentation
+
+# === Setup for GPU or CPU ===
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print(f"Using device: {device}")
 
 # 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.load_state_dict(torch.load("trained_model_200.pt", map_location=device))
+model_segformer.to(device)
 model_segformer.eval()
 
-# StarDist model
+# Load StarDist model (CPU-only, no GPU support)
 model_stardist = StarDist2D.from_pretrained('2D_versatile_fluo')
 
-# Cellpose model
-model_cellpose = cellpose_models.CellposeModel(gpu=False)
+# Load Cellpose model with GPU if available
+model_cellpose = cellpose_models.CellposeModel(gpu=torch.cuda.is_available())
 
-# Handle SegFormer prediction
+# SegFormer Inference
 def infer_segformer(image):
     image = image.convert("RGB")
     inputs = processor_segformer(images=image, return_tensors="pt")
+    inputs = {k: v.to(device) for k, v in inputs.items()}
+
     with torch.no_grad():
         logits = model_segformer(**inputs).logits
     pred_mask = torch.argmax(logits, dim=1)[0].cpu().numpy()
@@ -49,7 +57,7 @@ def infer_segformer(image):
         color_mask[pred_mask == c] = colors[c]
     return image, Image.fromarray(color_mask)
 
-# Handle StarDist prediction
+# StarDist Inference
 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))
@@ -57,12 +65,11 @@ def infer_stardist(image):
     overlay = (overlay[..., :3] * 255).astype(np.uint8)
     return image, Image.fromarray(overlay)
 
-# Handle MEDIAR prediction
+# MEDIAR Inference
 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 = {
@@ -74,11 +81,12 @@ def infer_mediar(image, temp_dir="temp_mediar"):
     }
 
     model = MEDIARFormer(**model_args)
-    weights = torch.load("from_phase1.pth", map_location="cpu")
+    weights = torch.load("MEDIAR_Weights/from_phase1.pth", map_location=device)
     model.load_state_dict(weights, strict=False)
+    model.to(device)
     model.eval()
 
-    predictor = Predictor(model, "cpu", temp_dir, temp_dir, algo_params={"use_tta": False})
+    predictor = Predictor(model, device.type, temp_dir, temp_dir, algo_params={"use_tta": False})
     predictor.img_names = ["input_image.tiff"]
     _ = predictor.conduct_prediction()
 
@@ -93,13 +101,13 @@ def infer_mediar(image, temp_dir="temp_mediar"):
     buf.seek(0)
 
     return image, Image.open(buf)
-# Handle Cellpose prediction
+
+# Cellpose Inference
 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)
@@ -112,18 +120,12 @@ def infer_cellpose(image, temp_dir="temp_cellpose"):
 
     return image, Image.open(output_overlay)
 
-# Wrapper function
+# Main segmentation dispatcher
 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...
+    ext = image.format.lower() if hasattr(image, 'format') and image.format else None
+    if model_name == "Cellpose" and ext not in ["tif", "tiff", None]:
+        return None, f"❌ Cellpose only supports `.tif` or `.tiff` images."
+
     if model_name == "SegFormer":
         return infer_segformer(image)
     elif model_name == "StarDist":
@@ -135,6 +137,7 @@ def segment(model_name, image):
     else:
         return None, f"❌ Unknown model: {model_name}"
 
+# === Gradio UI ===
 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.")
@@ -156,7 +159,7 @@ with gr.Blocks(title="Cell Segmentation Explorer") as app:
     def handle_submit(model_name, img):
         if img is None:
             return None
-        _, result = segment(model_name, img)  # Only return the mask (segmentation result)
+        _, result = segment(model_name, img)
         return result
 
     submit_btn.click(
@@ -171,18 +174,12 @@ with gr.Blocks(title="Cell Segmentation Explorer") as app:
         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"
-    ]
-
+    original_sample_paths = ["Sample Images/img1.png", "Sample Images/img2.png", "Sample Images/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"
@@ -192,7 +189,7 @@ with gr.Blocks(title="Cell Segmentation Explorer") as app:
     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
+            def load_full_image(idx=i):
                 return PILImage.open(original_sample_paths[idx])
 
             sample_img = gr.Image(value=img_path, type="pil", interactive=True, show_label=False)
@@ -203,5 +200,4 @@ with gr.Blocks(title="Cell Segmentation Explorer") as app:
             )
             sample_image_components.append(sample_img)
 
-
 app.launch()