Update app.py

app.py

@@ -5,11 +5,9 @@ from torchvision import transforms
 from PIL import Image
 import numpy as np
 import nibabel as nib
-from functools import lru_cache
-
 
 # -----------------------------
-# (1) Minimal SIREN definition
+# 1️⃣ Minimal SIREN definition
 # -----------------------------
 class SirenLayer(torch.nn.Module):
     def __init__(self, dim_in, dim_out, w0=30.0):
@@ -38,67 +36,57 @@ class Siren(torch.nn.Module):
     def forward(self, x):
         for layer in self.net:
             x = layer(x)
-        return torch.sigmoid(self.final(x))  # output ∈ [0,1]
+        return torch.sigmoid(self.final(x))
 
 
-# ----------------------------------
-# (2) Cached model loader
-# ----------------------------------
-@lru_cache(maxsize=1)
+# -----------------------------
+# 2️⃣ Load trained model
+# -----------------------------
 def load_model():
-    repo_id = "medimaging/physics-informed-sirenMRI-model"
-    filename = "sirenMRI_full_model.pt"
-
-    model_path = hf_hub_download(repo_id=repo_id, filename=filename)
+    model_path = hf_hub_download(
+        repo_id="medimaging/physics-informed-sirenMRI-model",
+        filename="sirenMRI_full_model.pt"
+    )
     obj = torch.load(model_path, map_location="cpu")
 
-    # Try to interpret the object
     if isinstance(obj, torch.nn.Module):
         model = obj
     elif isinstance(obj, dict):
-        if "all_slice_models" in obj:
+        if "model_state_dict" in obj:
+            model = Siren()
+            model.load_state_dict(obj["model_state_dict"])
+        elif "all_slice_models" in obj:
             first_slice = list(obj["all_slice_models"].values())[0]
-            model = Siren(dim_in=2, dim_hidden=256, dim_out=1, num_layers=3)
+            model = Siren()
             model.load_state_dict(first_slice)
-        elif "model_state_dict" in obj:
-            model = Siren(dim_in=2, dim_hidden=256, dim_out=1, num_layers=3)
-            model.load_state_dict(obj["model_state_dict"])
         else:
-            model = Siren(dim_in=2, dim_hidden=256, dim_out=1, num_layers=3)
+            model = Siren()
             model.load_state_dict(obj)
     else:
-        raise ValueError("Unrecognized model format in the .pt file.")
+        raise ValueError("❌ Unknown model format")
 
     model.eval()
     return model
 
 
-# ----------------------------------
-# (3) Pre/post-processing
-# ----------------------------------
-preprocess = transforms.Compose([
-    transforms.Resize((256, 256)),
-    transforms.ToTensor(),
-])
-postprocess = transforms.ToPILImage()
+model = load_model()
+
 
+# -----------------------------
+# 3️⃣ Helper functions
+# -----------------------------
+preprocess = transforms.Compose([transforms.Resize((256, 256)), transforms.ToTensor()])
+postprocess = transforms.ToPILImage()
 
-# ----------------------------------
-# (4) Utility functions
-# ----------------------------------
 def psnr(img1, img2):
     mse = torch.mean((img1 - img2) ** 2)
-    if mse == 0:
-        return float("inf")
-    return 20 * torch.log10(1.0 / torch.sqrt(mse))
-
+    return float("inf") if mse == 0 else 20 * torch.log10(1.0 / torch.sqrt(mse))
 
-def estimate_compression_ratio(image, model_params):
-    img_bytes = image.size[0] * image.size[1] * 3
+def estimate_compression_ratio(original_img, model_params):
+    img_bytes = original_img.size[0] * original_img.size[1] * 3
     model_bytes = model_params * 4
     return round(img_bytes / model_bytes, 2)
 
-
 def extract_middle_slice(nii_path):
     nii = nib.load(nii_path)
     data = nii.get_fdata()
@@ -109,23 +97,15 @@ def extract_middle_slice(nii_path):
     return Image.fromarray(slice_img)
 
 
-# ----------------------------------
-# (5) Core pipeline
-# ----------------------------------
+# -----------------------------
+# 4️⃣ Main function
+# -----------------------------
 def compress_and_reconstruct(file):
-    try:
-        model = load_model()
-    except Exception as e:
-        return None, f"❌ Model failed to load: {e}"
-
     filename = getattr(file, "name", "")
-    try:
-        if filename.endswith((".nii", ".nii.gz")):
-            image = extract_middle_slice(filename)
-        else:
-            image = Image.open(file).convert("RGB")
-    except Exception:
-        return None, "❌ Failed to read image or NIfTI file."
+    if filename.endswith(".nii") or filename.endswith(".nii.gz"):
+        image = extract_middle_slice(filename)
+    else:
+        image = Image.open(file).convert("RGB")
 
     original = preprocess(image).unsqueeze(0)
 
@@ -142,37 +122,66 @@ def compress_and_reconstruct(file):
     combined.paste(orig_img, (0, 0))
     combined.paste(recon_img, (orig_img.width, 0))
 
-    return combined, f"✅ PSNR: {score:.2f} dB | Compression Ratio: {ratio}:1"
-
-
-# ----------------------------------
-# (6) Gradio UI
-# ----------------------------------
-title = "🧠 Physics-Informed SIREN MRI Compression"
-description = """
-Upload an **MRI file (.nii / .nii.gz)** or **image (PNG/JPG)**  
-to visualize compression → reconstruction using the **Physics-Informed SIREN** model.
-
-Shows:
-- Side-by-side comparison *(Original → Reconstructed)*
-- PSNR & Compression Ratio metrics
-
-**Framework:** PyTorch + Gradio  
-**Author:** medimaging
-"""
-
-interface = gr.Interface(
-    fn=compress_and_reconstruct,
-    inputs=gr.File(label="Upload MRI or Image File (.nii, .nii.gz, .png, .jpg)"),
-    outputs=[
-        gr.Image(label="Original (Left) → Reconstructed (Right)"),
-        gr.Textbox(label="Quality Metrics", interactive=False)
-    ],
-    title=title,
-    description=description,
-    allow_flagging="never",
-    examples=None,
-)
+    metrics = f"📊 PSNR: {score:.2f} dB | 💾 Compression Ratio: {ratio}:1"
+    return combined, metrics
+
+
+# -----------------------------
+# 5️⃣ Custom Gradio UI
+# -----------------------------
+with gr.Blocks(css="""
+@keyframes gradient {
+    0% {background-position: 0% 50%;}
+    50% {background-position: 100% 50%;}
+    100% {background-position: 0% 50%;}
+}
+#title {
+    font-size: 32px;
+    font-weight: 800;
+    text-align: center;
+    background: linear-gradient(-45deg, #2563eb, #7c3aed, #06b6d4, #3b82f6);
+    background-size: 400% 400%;
+    -webkit-background-clip: text;
+    -webkit-text-fill-color: transparent;
+    animation: gradient 6s ease infinite;
+    margin-top: 20px;
+}
+#description {
+    text-align: center;
+    font-size: 18px;
+    color: #475569;
+    max-width: 720px;
+    margin: 0 auto 25px auto;
+}
+#footer {
+    text-align: center;
+    font-size: 14px;
+    color: #94a3b8;
+    margin-top: 30px;
+}
+""") as demo:
+
+    gr.HTML("<div id='title'>🧠 Compress Medical Images to ease documentation and processing</div>")
+    gr.HTML("""
+        <div id='description'>
+        Upload an <b>MRI file (.nii / .nii.gz)</b> or <b>image (.png / .jpg)</b> to visualize 
+        <b>compression → reconstruction</b> using a Physics-Informed SIREN model.<br>
+        <small>View PSNR and compression ratio below.</small>
+        </div>
+    """)
+
+    with gr.Row():
+        with gr.Column(scale=1, min_width=300):
+            file_input = gr.File(label="📁 Upload MRI or Image File")
+            submit_btn = gr.Button("🚀 Run Compression", variant="primary")
+        with gr.Column(scale=2):
+            output_image = gr.Image(label="Original (Left) → Reconstructed (Right)")
+            output_text = gr.Textbox(label="Metrics", interactive=False)
+
+    submit_btn.click(compress_and_reconstruct, inputs=file_input, outputs=[output_image, output_text])
+
+    gr.HTML("<div id='footer'>Made with ❤️ using PyTorch + Gradio</div>")
+
 
 if __name__ == "__main__":
-    interface.launch()
+    demo.launch()