Enhance fibril segmentation app with model selection, improved UI, and device handling; add UNet model weights

app.py

@@ -81,6 +81,118 @@
 # # Last Updated: 08 July 2025
 # # Improvements: Added examples, better UI, and device handling
 
+# import os
+# import torch
+# import numpy as np
+# from PIL import Image
+# import albumentations as A
+# from albumentations.pytorch import ToTensorV2
+# import segmentation_models_pytorch as smp
+# import gradio as gr
+
+# # ─── Configuration ─────────────────────────────────────────
+# CONFIG = {
+#     "model_path": "./model/encoder_resnet34_decoder_UnetPlusPlus_fibril_seg_model.pth",
+#     "img_size": 512
+# }
+
+# # ─── Device Setup ──────────────────────────────────────────
+# device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+# print(f"✅ Using device: {device}")
+
+# # ─── Load Model ────────────────────────────────────────────
+# model = smp.UnetPlusPlus(
+#     encoder_name='resnet34',
+#     encoder_depth=5,
+#     encoder_weights='imagenet',
+#     decoder_channels=(256, 128, 64, 32, 16),
+#     in_channels=1,
+#     classes=1,
+#     activation=None
+# ).to(device)
+
+# model.load_state_dict(torch.load(CONFIG["model_path"], map_location=device))
+# model.eval()
+
+# # ─── Transform Function ────────────────────────────────────
+# def get_transform(size):
+#     return A.Compose([
+#         A.Resize(size, size),
+#         A.Normalize(mean=(0.5,), std=(0.5,)),
+#         ToTensorV2()
+#     ])
+
+# transform = get_transform(CONFIG["img_size"])
+
+# # ─── Prediction Function ───────────────────────────────────
+# def predict(image):
+#     image = image.convert("L")  # Ensure grayscale
+#     img_np = np.array(image)
+#     img_tensor = transform(image=img_np)["image"].unsqueeze(0).to(device)
+
+#     with torch.no_grad():
+#         pred = torch.sigmoid(model(img_tensor))
+#         mask = (pred > 0.5).float().cpu().squeeze().numpy()
+
+#     mask_img = Image.fromarray((mask * 255).astype(np.uint8))
+#     return mask_img
+
+# # ─── Gradio UI (Improved) ──────────────────────────────────
+# examples = [
+#     ["examples/example1.jpg"],
+#     ["examples/example2.jpg"],
+#     ["examples/example3.jpg"],
+#     ["examples/example4.jpg"],
+#     ["examples/example5.jpg"],
+#     ["examples/example6.jpg"],
+#     ["examples/example7.jpg"]
+# ]
+
+# css = """
+# .gradio-container {
+#     max-width: 950px;
+#     margin: auto;
+# }
+# .gr-button {
+#     background-color: #4a90e2;
+#     color: white;
+#     border-radius: 5px;
+# }
+# .gr-button:hover {
+#     background-color: #357ABD;
+# }
+# """
+
+# with gr.Blocks(css=css) as demo:
+#     gr.Markdown("## 🧬 Fibril Segmentation with UNet++")
+#     gr.Markdown("Upload a **grayscale microscopy image**, and this model will predict the **segmentation mask of fibrillar structures**.\n\nModel: ResNet34 encoder + UNet++ decoder")
+
+#     with gr.Row():
+#         input_img = gr.Image(label="Upload Microscopy Image", type="pil")
+#         output_mask = gr.Image(label="Predicted Segmentation Mask", type="pil")
+
+#     submit_btn = gr.Button("Segment Image")
+
+#     submit_btn.click(fn=predict, inputs=input_img, outputs=output_mask)
+
+#     gr.Examples(
+#         examples=examples,
+#         inputs=input_img,
+#         label="Try with Example Images",
+#         cache_examples=False
+#     )
+
+# # ─── Launch App ────────────────────────────────────────────
+# if __name__ == "__main__":
+#     demo.launch()
+
+
+
+#  +++++++++++++++ Final Version: 1.2.0 ++++++++++++++++++++++
+# Last Updated: 08 July 2025
+# Improvements: Added model selection, better UI, and device handling
+
+
 import os
 import torch
 import numpy as np
@@ -90,29 +202,23 @@ from albumentations.pytorch import ToTensorV2
 import segmentation_models_pytorch as smp
 import gradio as gr
 
-# ─── Configuration ─────────────────────────────────────────
-CONFIG = {
-    "model_path": "./model/encoder_resnet34_decoder_UnetPlusPlus_fibril_seg_model.pth",
-    "img_size": 512
-}
-
 # ─── Device Setup ──────────────────────────────────────────
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 print(f"✅ Using device: {device}")
 
-# ─── Load Model ────────────────────────────────────────────
-model = smp.UnetPlusPlus(
-    encoder_name='resnet34',
-    encoder_depth=5,
-    encoder_weights='imagenet',
-    decoder_channels=(256, 128, 64, 32, 16),
-    in_channels=1,
-    classes=1,
-    activation=None
-).to(device)
-
-model.load_state_dict(torch.load(CONFIG["model_path"], map_location=device))
-model.eval()
+# ─── Model Configurations ──────────────────────────────────
+MODEL_OPTIONS = {
+    "UNet++ (ResNet34)": {
+        "path": "./model/encoder_resnet34_decoder_UnetPlusPlus_fibril_seg_model.pth",
+        "encoder": "resnet34",
+        "architecture": "UnetPlusPlus"
+    },
+    "UNet (ResNet34)": {
+        "path": "./model/unet_fibril_seg_model.pth",
+        "encoder": "resnet34",
+        "architecture": "Unet"
+    }
+}
 
 # ─── Transform Function ────────────────────────────────────
 def get_transform(size):
@@ -122,14 +228,44 @@ def get_transform(size):
         ToTensorV2()
     ])
 
-transform = get_transform(CONFIG["img_size"])
+transform = get_transform(512)
+
+# ─── Model Loader ──────────────────────────────────────────
+def load_model(model_name):
+    config = MODEL_OPTIONS[model_name]
+    if config["architecture"] == "UnetPlusPlus":
+        model = smp.UnetPlusPlus(
+            encoder_name=config["encoder"],
+            encoder_weights="imagenet",
+            decoder_channels=(256, 128, 64, 32, 16),
+            in_channels=1,
+            classes=1,
+            activation=None
+        )
+    elif config["architecture"] == "Unet":
+        model = smp.Unet(
+            encoder_name=config["encoder"],
+            encoder_weights="imagenet",
+            decoder_channels=(256, 128, 64, 32, 16),
+            in_channels=1,
+            classes=1,
+            activation=None
+        )
+    else:
+        raise ValueError(f"Unsupported architecture: {config['architecture']}")
+
+    model.load_state_dict(torch.load(config["path"], map_location=device))
+    model.eval()
+    return model.to(device)
 
 # ─── Prediction Function ───────────────────────────────────
-def predict(image):
-    image = image.convert("L")  # Ensure grayscale
+def predict(image, model_name):
+    image = image.convert("L")
     img_np = np.array(image)
     img_tensor = transform(image=img_np)["image"].unsqueeze(0).to(device)
 
+    model = load_model(model_name)
+
     with torch.no_grad():
         pred = torch.sigmoid(model(img_tensor))
         mask = (pred > 0.5).float().cpu().squeeze().numpy()
@@ -137,7 +273,7 @@ def predict(image):
     mask_img = Image.fromarray((mask * 255).astype(np.uint8))
     return mask_img
 
-# ─── Gradio UI (Improved) ──────────────────────────────────
+# ─── Example Images ────────────────────────────────────────
 examples = [
     ["examples/example1.jpg"],
     ["examples/example2.jpg"],
@@ -148,6 +284,7 @@ examples = [
     ["examples/example7.jpg"]
 ]
 
+# ─── Custom CSS ────────────────────────────────────────────
 css = """
 .gradio-container {
     max-width: 950px;
@@ -163,9 +300,13 @@ css = """
 }
 """
 
+# ─── Gradio UI ─────────────────────────────────────────────
 with gr.Blocks(css=css) as demo:
-    gr.Markdown("## 🧬 Fibril Segmentation with UNet++")
-    gr.Markdown("Upload a **grayscale microscopy image**, and this model will predict the **segmentation mask of fibrillar structures**.\n\nModel: ResNet34 encoder + UNet++ decoder")
+    gr.Markdown("## 🧬 Fibril Segmentation Interface")
+    gr.Markdown("Choose a model and upload a grayscale microscopy image. The model will predict the **fibrillar structure mask**.")
+
+    with gr.Row():
+        model_selector = gr.Dropdown(choices=list(MODEL_OPTIONS.keys()), value="UNet++ (ResNet34)", label="Select Model")
 
     with gr.Row():
         input_img = gr.Image(label="Upload Microscopy Image", type="pil")
@@ -173,7 +314,7 @@ with gr.Blocks(css=css) as demo:
 
     submit_btn = gr.Button("Segment Image")
 
-    submit_btn.click(fn=predict, inputs=input_img, outputs=output_mask)
+    submit_btn.click(fn=predict, inputs=[input_img, model_selector], outputs=output_mask)
 
     gr.Examples(
         examples=examples,

model/unet_fibril_seg_model.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:39aabaff5e7006840147b16e67ca995fee643742b3d44b3ade469485384fd153
+size 97898267