m

app.py

@@ -4,12 +4,10 @@ import numpy as np
 import torch.nn.functional as F
 from torch import nn
 from transformers import AutoImageProcessor, Swinv2ForImageClassification, SegformerForSemanticSegmentation
-import matplotlib.pyplot as plt
 import streamlit as st
 from PIL import Image
 import io
 import zipfile
-import os
 
 # --- GlaucomaModel Class ---
 class GlaucomaModel(object):
@@ -48,8 +46,16 @@ class GlaucomaModel(object):
         )
         seg_probs = F.softmax(upsampled_logits, dim=1)
         pred_disc_cup = upsampled_logits.argmax(dim=1)[0]
-        cup_confidence = seg_probs[0, 2, :, :].mean().item() * 100
-        disc_confidence = seg_probs[0, 1, :, :].mean().item() * 100
+        
+        # Calculate segmentation confidence based on probability distribution
+        # For each pixel classified as cup/disc, check how confident the model is
+        cup_mask = pred_disc_cup == 2
+        disc_mask = pred_disc_cup == 1
+        
+        # Get confidence only for pixels predicted as cup/disc
+        cup_confidence = seg_probs[0, 2, cup_mask].mean().item() * 100 if cup_mask.any() else 0
+        disc_confidence = seg_probs[0, 1, disc_mask].mean().item() * 100 if disc_mask.any() else 0
+        
         return pred_disc_cup.numpy().astype(np.uint8), cup_confidence, disc_confidence
 
     def process(self, image):
@@ -90,64 +96,133 @@ def add_mask(image, mask, classes, colors, alpha=0.5):
     output = cv2.addWeighted(overlay, alpha, image, 1 - alpha, 0)
     return output, overlay
 
+def get_confidence_level(confidence):
+    if confidence >= 90:
+        return "Very High"
+    elif confidence >= 75:
+        return "High"
+    elif confidence >= 60:
+        return "Moderate"
+    elif confidence >= 45:
+        return "Low"
+    else:
+        return "Very Low"
+
 # --- Streamlit Interface ---
 def main():
-    st.set_page_config(layout="wide")
-    st.title("Batch Glaucoma Screening from Retinal Fundus Images")
+    st.set_page_config(layout="wide", page_title="Glaucoma Screening Tool")
     
-    # Explanation for the confidence threshold
-    st.sidebar.write("**Confidence Threshold** (optional): Set a threshold to filter images based on the model's confidence in glaucoma classification.")
-    confidence_threshold = st.sidebar.slider("Confidence Threshold (%)", 0, 100, 70)
-    uploaded_files = st.sidebar.file_uploader("Upload Images", type=['png', 'jpeg', 'jpg'], accept_multiple_files=True)
+    # Header with better styling
+    st.markdown("""
+        <h1 style='text-align: center;'>Glaucoma Screening from Retinal Fundus Images</h1>
+        <p style='text-align: center; color: gray;'>Upload retinal images for automated glaucoma detection and optic disc/cup segmentation</p>
+        """, unsafe_allow_html=True)
+    
+    # Sidebar with better organization
+    with st.sidebar:
+        st.markdown("### Upload Settings")
+        uploaded_files = st.file_uploader("Upload Retinal Images", 
+            type=['png', 'jpeg', 'jpg'], 
+            accept_multiple_files=True,
+            help="Support multiple images in PNG, JPEG formats")
+        
+        st.markdown("### Analysis Settings")
+        st.info("📊 Set confidence threshold to filter results")
+        confidence_threshold = st.slider(
+            "Classification Confidence Threshold (%)", 
+            0, 100, 70,
+            help="Images with confidence above this threshold will be marked as reliable predictions")
     
-    confident_images = []
-    download_confident_images = []
-
     if uploaded_files:
         for uploaded_file in uploaded_files:
             image = Image.open(uploaded_file).convert('RGB')
             image_np = np.array(image).astype(np.uint8)
             
-            with st.spinner(f'Processing {uploaded_file.name}...'):
+            with st.spinner(f'🔄 Processing {uploaded_file.name}...'):
                 model = GlaucomaModel(device=torch.device("cuda:0" if torch.cuda.is_available() else "cpu"))
                 disease_idx, disc_cup_image, vcdr, cls_conf, cup_conf, disc_conf, cropped_image = model.process(image_np)
                 
-                # Confidence-based grouping
-                is_confident = cls_conf >= confidence_threshold
-                if is_confident:
-                    confident_images.append(uploaded_file.name)
-                    download_confident_images.append((cropped_image, uploaded_file.name))
-                
-                # Display Results
-                st.subheader(f"Results for {uploaded_file.name}")
-                cols = st.beta_columns(4)  # Use st.beta_columns for compatibility with older Streamlit
-                cols[0].image(image_np, caption="Input Image", use_column_width=True)
-                cols[1].image(disc_cup_image, caption="Disc/Cup Segmentation", use_column_width=True)
-                cols[2].image(image_np, caption="Class Activation Map", use_column_width=True)
-                cols[3].image(cropped_image, caption="Cropped Image", use_column_width=True)
-
-                # Display confidence and metrics
-                st.write(f"**Vertical cup-to-disc ratio:** {vcdr:.04f}")
-                st.write(f"**Category:** {model.cls_id2label[disease_idx]} ({cls_conf:.02f}% confidence)")
-                st.write(f"**Optic Cup Segmentation Confidence:** {cup_conf:.02f}%")
-                st.write(f"**Optic Disc Segmentation Confidence:** {disc_conf:.02f}%")
-                st.write(f"**Confidence Group:** {'Confident' if is_confident else 'Not Confident'}")
+                # Create expandable section for each image
+                with st.expander(f"📊 Analysis Results: {uploaded_file.name}", expanded=True):
+                    # Image display section
+                    cols = st.columns(3)
+                    cols[0].image(image_np, caption="Original Image", use_column_width=True)
+                    cols[1].image(disc_cup_image, caption="Segmentation Overlay", use_column_width=True)
+                    cols[2].image(cropped_image, caption="Region of Interest", use_column_width=True)
+                    
+                    # Metrics section with clear separation
+                    st.markdown("---")
+                    metric_cols = st.columns(3)
+                    
+                    # Classification Results
+                    with metric_cols[0]:
+                        st.markdown("### 🔍 Classification")
+                        diagnosis = model.cls_id2label[disease_idx]
+                        is_confident = cls_conf >= confidence_threshold
+                        
+                        # Color-coded diagnosis
+                        if diagnosis == "Glaucoma":
+                            st.markdown(f"<div style='padding: 10px; background-color: #ffebee; border-radius: 5px;'>"
+                                      f"<h4 style='color: #c62828;'>Diagnosis: {diagnosis}</h4></div>", 
+                                      unsafe_allow_html=True)
+                        else:
+                            st.markdown(f"<div style='padding: 10px; background-color: #e8f5e9; border-radius: 5px;'>"
+                                      f"<h4 style='color: #2e7d32;'>Diagnosis: {diagnosis}</h4></div>", 
+                                      unsafe_allow_html=True)
+                        
+                        st.metric("Classification Confidence", f"{cls_conf:.1f}%")
+                        if not is_confident:
+                            st.warning("⚠️ Below confidence threshold")
+                    
+                    # Segmentation Results
+                    with metric_cols[1]:
+                        st.markdown("### 🎯 Segmentation Quality")
+                        st.metric("Optic Cup Confidence", f"{cup_conf:.1f}%")
+                        st.metric("Optic Disc Confidence", f"{disc_conf:.1f}%")
+                        
+                        # Confidence level explanation
+                        cup_level = get_confidence_level(cup_conf)
+                        disc_level = get_confidence_level(disc_conf)
+                        st.info(f"Cup Detection: {cup_level}\nDisc Detection: {disc_level}")
+                    
+                    # Clinical Metrics
+                    with metric_cols[2]:
+                        st.markdown("### 📏 Clinical Metrics")
+                        st.metric("Cup-to-Disc Ratio (CDR)", f"{vcdr:.3f}")
+                        
+                        # CDR interpretation
+                        if vcdr > 0.7:
+                            st.warning("⚠️ Elevated CDR (>0.7)")
+                        elif vcdr > 0.5:
+                            st.info("ℹ️ Borderline CDR (0.5-0.7)")
+                        else:
+                            st.success("✅ Normal CDR (<0.5)")
         
-        # Download Link for Confident Images
+        # Download section
         if download_confident_images:
+            st.sidebar.markdown("---")
+            st.sidebar.markdown("### Download Results")
             with zipfile.ZipFile("confident_cropped_images.zip", "w") as zf:
                 for cropped_image, name in download_confident_images:
                     img_buffer = io.BytesIO()
                     Image.fromarray(cropped_image).save(img_buffer, format="PNG")
                     zf.writestr(f"{name}_cropped.png", img_buffer.getvalue())
-                    
-            # Provide a markdown link to the ZIP file
-            st.sidebar.markdown(
-                f"[Download Confident Cropped Images](./confident_cropped_images.zip)", 
-                unsafe_allow_html=True
+            
+            st.sidebar.download_button(
+                label="📥 Download Analysis Results",
+                data=open("confident_cropped_images.zip", "rb"),
+                file_name="glaucoma_analysis_results.zip",
+                mime="application/zip",
+                help="Download cropped images and analysis results"
             )
     else:
-        st.sidebar.info("Upload images to begin analysis.")
+        # Welcome message when no files are uploaded
+        st.markdown("""
+            <div style='text-align: center; padding: 50px;'>
+                <h3>👋 Welcome to the Glaucoma Screening Tool</h3>
+                <p>Upload retinal fundus images using the sidebar to begin analysis</p>
+            </div>
+            """, unsafe_allow_html=True)
 
 if __name__ == '__main__':
-    main()
+    main()