Update app.py

app.py

@@ -7,6 +7,7 @@ from transformers import AutoImageProcessor, Swinv2ForImageClassification, Segfo
 import matplotlib.pyplot as plt
 import streamlit as st
 from PIL import Image
+import io
 
 # --- GlaucomaModel Class ---
 class GlaucomaModel(object):
@@ -34,7 +35,7 @@ class GlaucomaModel(object):
             # Softmax for probabilities
             probs = F.softmax(outputs, dim=-1)
             disease_idx = probs.cpu()[0, :].numpy().argmax()
-            confidence = probs.cpu()[0, disease_idx].item()
+            confidence = probs.cpu()[0, disease_idx].item() * 100  # Scale to percentage
         return disease_idx, confidence
 
     def optic_disc_cup_pred(self, image):
@@ -49,8 +50,8 @@ class GlaucomaModel(object):
         # 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()
-        disc_confidence = seg_probs[0, 1, :, :].mean().item()
+        cup_confidence = seg_probs[0, 2, :, :].mean().item() * 100  # Scale to percentage
+        disc_confidence = seg_probs[0, 1, :, :].mean().item() * 100  # Scale to percentage
         return pred_disc_cup.numpy().astype(np.uint8), cup_confidence, disc_confidence
 
     def process(self, image):
@@ -66,20 +67,16 @@ class GlaucomaModel(object):
         # Mask for optic disc and cup
         mask = (disc_cup > 0).astype(np.uint8)
         
-        # Get bounding box of the optic cup + disc and add padding
+        # Get bounding box of the optic cup + disc and add dynamic padding
         x, y, w, h = cv2.boundingRect(mask)
-        padding = 20  # Add padding to avoid edge distortion
+        padding = max(50, int(0.2 * max(w, h)))  # Dynamic padding (20% of width or height)
         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
-        min_size = 50  # Define a minimum size for the cropped image
-        if w < min_size or h < min_size:
-            cropped_image = image.copy()  # Fallback: if bounding box too small, return original image
-        else:
-            cropped_image = image[y:y+h, x:x+w]
+        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)
@@ -95,9 +92,7 @@ def simple_vcdr(mask):
     """
     disc_area = np.sum(mask == 1)
     cup_area = np.sum(mask == 2)
-    
-    # Avoid division by zero
-    if disc_area == 0:
+    if disc_area == 0:  # Avoid division by zero
         return np.nan
     vcdr = cup_area / disc_area
     return vcdr
@@ -115,15 +110,12 @@ def add_mask(image, mask, classes, colors, alpha=0.5):
     overlay = image.copy()
     for class_id, color in zip(classes, colors):
         overlay[mask == class_id] = color
-    # Blend the overlay with the original image
     output = cv2.addWeighted(overlay, alpha, image, 1 - alpha, 0)
     return output, overlay
 
 # --- Streamlit Interface ---
 def main():
-    # Wide mode in Streamlit
     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')
 
@@ -153,12 +145,10 @@ def main():
             st.sidebar.write("Please upload an image")
         else:
             with st.spinner('Loading model...'):
-                # Load the model on available device
                 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
@@ -176,17 +166,27 @@ def main():
                 ax.axis('off')
                 cols[3].pyplot(fig)
 
+                # Make cropped image downloadable
+                buf = io.BytesIO()
+                Image.fromarray(cropped_image).save(buf, format="PNG")
+                st.sidebar.download_button(
+                    label="Download Cropped Image",
+                    data=buf.getvalue(),
+                    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*100:.02f}% confidence)|
-                |Optic Cup Segmentation Confidence|{cup_confidence*100:.02f}%|
-                |Optic Disc Segmentation Confidence|{disc_confidence*100:.02f}%|
+                |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)
 
 if __name__ == '__main__':
-    main()
+