Update app.py

app.py

@@ -8,6 +8,8 @@ import matplotlib.pyplot as plt
 import streamlit as st
 from PIL import Image
 import io
+import zipfile
+import os
 
 # --- GlaucomaModel Class ---
 class GlaucomaModel(object):
@@ -22,20 +24,17 @@ class GlaucomaModel(object):
         # Segmentation model for optic disc and cup
         self.seg_extractor = AutoImageProcessor.from_pretrained(seg_model_path)
         self.seg_model = SegformerForSemanticSegmentation.from_pretrained(seg_model_path).to(device).eval()
-
-        # Class activation map
+        # Mapping for class labels
         self.cls_id2label = self.cls_model.config.id2label
-        self.seg_id2label = self.seg_model.config.id2label
 
     def glaucoma_pred(self, image):
         inputs = self.cls_extractor(images=image.copy(), return_tensors="pt")
         with torch.no_grad():
             inputs.to(self.device)
             outputs = self.cls_model(**inputs).logits
-            # Softmax for probabilities
             probs = F.softmax(outputs, dim=-1)
             disease_idx = probs.cpu()[0, :].numpy().argmax()
-            confidence = probs.cpu()[0, disease_idx].item() * 100  # Scale to percentage
+            confidence = probs.cpu()[0, disease_idx].item() * 100
         return disease_idx, confidence
 
     def optic_disc_cup_pred(self, image):
@@ -47,66 +46,44 @@ class GlaucomaModel(object):
         upsampled_logits = nn.functional.interpolate(
             logits, size=image.shape[:2], mode="bilinear", align_corners=False
         )
-        # Softmax for segmentation confidence
         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  # Scale to percentage
-        disc_confidence = seg_probs[0, 1, :, :].mean().item() * 100  # Scale to percentage
+        cup_confidence = seg_probs[0, 2, :, :].mean().item() * 100
+        disc_confidence = seg_probs[0, 1, :, :].mean().item() * 100
         return pred_disc_cup.numpy().astype(np.uint8), cup_confidence, disc_confidence
 
     def process(self, image):
-        image_shape = image.shape[:2]
         disease_idx, cls_confidence = self.glaucoma_pred(image)
         disc_cup, cup_confidence, disc_confidence = self.optic_disc_cup_pred(image)
 
         try:
-            vcdr = simple_vcdr(disc_cup)  # Calculate vertical cup-to-disc ratio
+            vcdr = simple_vcdr(disc_cup)
         except:
             vcdr = np.nan
 
-        # Mask for optic disc and cup
         mask = (disc_cup > 0).astype(np.uint8)
-        
-        # Get bounding box of the optic cup + disc and add dynamic padding
         x, y, w, h = cv2.boundingRect(mask)
-        padding = max(50, int(0.2 * max(w, h)))  # Dynamic padding (20% of width or height)
+        padding = max(50, int(0.2 * max(w, h)))
         x = max(x - padding, 0)
         y = max(y - padding, 0)
         w = min(w + 2 * padding, image.shape[1] - x)
         h = min(h + 2 * padding, image.shape[0] - y)
 
-        # Ensure that the bounding box is large enough to avoid cropping errors
         cropped_image = image[y:y+h, x:x+w] if w >= 50 and h >= 50 else image.copy()
-
-        # Generate disc and cup visualization
         _, disc_cup_image = add_mask(image, disc_cup, [1, 2], [[0, 255, 0], [255, 0, 0]], 0.2)
 
         return disease_idx, disc_cup_image, vcdr, cls_confidence, cup_confidence, disc_confidence, cropped_image
 
 # --- Utility Functions ---
 def simple_vcdr(mask):
-    """
-    Simple function to calculate the vertical cup-to-disc ratio (VCDR).
-    Assumes:
-    - mask contains class 1 for optic disc and class 2 for optic cup.
-    """
     disc_area = np.sum(mask == 1)
     cup_area = np.sum(mask == 2)
-    if disc_area == 0:  # Avoid division by zero
+    if disc_area == 0:
         return np.nan
     vcdr = cup_area / disc_area
     return vcdr
 
 def add_mask(image, mask, classes, colors, alpha=0.5):
-    """
-    Adds a transparent mask to the original image.
-    Args:
-    - image: the original RGB image
-    - mask: the predicted segmentation mask
-    - classes: a list of class indices to apply masks for (e.g., [1, 2])
-    - colors: a list of colors for each class (e.g., [[0, 255, 0], [255, 0, 0]] for green and red)
-    - alpha: transparency level (default = 0.5)
-    """
     overlay = image.copy()
     for class_id, color in zip(classes, colors):
         overlay[mask == class_id] = color
@@ -116,79 +93,59 @@ def add_mask(image, mask, classes, colors, alpha=0.5):
 # --- Streamlit Interface ---
 def main():
     st.set_page_config(layout="wide")
-    st.title("Glaucoma Screening from Retinal Fundus Images")
-    st.write('Developed by X. Sun. Find more info about me: https://pamixsun.github.io')
-
-    # Set columns for the interface
-    cols = st.beta_columns((1, 1, 1, 1))
-    cols[0].subheader("Input image")
-    cols[1].subheader("Optic disc and optic cup")
-    cols[2].subheader("Class activation map")
-    cols[3].subheader("Cropped Image")
-
-    # File uploader
-    st.sidebar.title("Image selection")
-    st.set_option('deprecation.showfileUploaderEncoding', False)
-    uploaded_file = st.sidebar.file_uploader("Upload image", type=['png', 'jpeg', 'jpg'])
-
-    if uploaded_file is not None:
-        # Read and display uploaded image
-        image = Image.open(uploaded_file).convert('RGB')
-        image = np.array(image).astype(np.uint8)
-        fig, ax = plt.subplots()
-        ax.imshow(image)
-        ax.axis('off')
-        cols[0].pyplot(fig)
-
-    if st.sidebar.button("Analyze image"):
-        if uploaded_file is None:
-            st.sidebar.write("Please upload an image")
-        else:
-            with st.spinner('Loading model...'):
-                run_device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-                model = GlaucomaModel(device=run_device)
-
-            with st.spinner('Analyzing...'):
-                # Get predictions from the model
-                disease_idx, disc_cup_image, vcdr, cls_confidence, cup_confidence, disc_confidence, cropped_image = model.process(image)
-
-                # Display optic disc and cup image
-                ax.imshow(disc_cup_image)
-                ax.axis('off')
-                cols[1].pyplot(fig)
-
-                # Display classification map
-                ax.imshow(image)
-                ax.axis('off')
-                cols[2].pyplot(fig)
-
-                # Display the cropped image
-                ax.imshow(cropped_image)
-                ax.axis('off')
-                cols[3].pyplot(fig)
-
-                # Make cropped image downloadable by converting it to bytes
-                buf = io.BytesIO()
-                Image.fromarray(cropped_image).save(buf, format="PNG")
-                byte_img = buf.getvalue()
-                st.sidebar.download_button(
-                    label="Download Cropped Image",
-                    data=byte_img,
-                    file_name="cropped_image.png",
-                    mime="image/png"
-                )
-
-                # Display results with confidence
-                st.subheader("Screening results:")
-                final_results_as_table = f"""
-                |Parameters|Outcomes|
-                |---|---|
-                |Vertical cup-to-disc ratio|{vcdr:.04f}|
-                |Category|{model.cls_id2label[disease_idx]} ({cls_confidence:.02f}% confidence)|
-                |Optic Cup Segmentation Confidence|{cup_confidence:.02f}%|
-                |Optic Disc Segmentation Confidence|{disc_confidence:.02f}%|
-                """
-                st.markdown(final_results_as_table)
+    st.title("Batch Glaucoma Screening from Retinal Fundus Images")
+    
+    st.sidebar.title("Settings")
+    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)
+    
+    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}...'):
+                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
+                with st.expander(f"Results for {uploaded_file.name}", expanded=False):
+                    cols = st.columns(4)
+                    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)
+
+                    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'}")
+        
+        # Download Button for Confident Images
+        if download_confident_images:
+            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
+            )
+    else:
+        st.sidebar.info("Upload images to begin analysis.")
 
 if __name__ == '__main__':
     main()