Update app.py

app.py

@@ -57,12 +57,67 @@ class GlaucomaModel(object):
         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)  # Assuming simple_vcdr() is defined elsewhere
+            vcdr = simple_vcdr(disc_cup)  # Calculate vertical cup-to-disc ratio
         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 padding
+        x, y, w, h = cv2.boundingRect(mask)
+        padding = 20  # Add padding to avoid edge distortion
+        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]
+
+        # 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
+
+        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)
+    
+    # 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
+    # Blend the overlay with the original image
+    output = cv2.addWeighted(overlay, alpha, image, 1 - alpha, 0)
+    return output, overlay
 
 # --- Streamlit Interface ---
 def main():
@@ -73,10 +128,11 @@ def main():
     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))
+    cols = st.beta_columns((1, 1, 1, 1))
     cols[0].subheader("Input image")
     cols[1].subheader("Optic disc and optic cup")
-    cols[2].subheader("Classification Map")
+    cols[2].subheader("Class activation map")
+    cols[3].subheader("Cropped Image")
 
     # File uploader
     st.sidebar.title("Image selection")
@@ -103,7 +159,7 @@ def main():
 
             with st.spinner('Analyzing...'):
                 # Get predictions from the model
-                disease_idx, disc_cup_image, vcdr, cls_confidence, cup_confidence, disc_confidence = model.process(image)
+                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)
@@ -115,6 +171,11 @@ def main():
                 ax.axis('off')
                 cols[2].pyplot(fig)
 
+                # Display the cropped image
+                ax.imshow(cropped_image)
+                ax.axis('off')
+                cols[3].pyplot(fig)
+
                 # Display results with confidence
                 st.subheader("Screening results:")
                 final_results_as_table = f"""