Update app.py

app.py

@@ -1,99 +1,131 @@
-# Copyright (C) 2023, Xu Sun.
-
-# This program is licensed under the Apache License version 2.
-# See LICENSE or go to <https://www.apache.org/licenses/LICENSE-2.0.txt> for full license details.
-
+import cv2
 import torch
 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
-from glaucoma import GlaucomaModel
-
-run_device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-
 
+# --- GlaucomaModel Class ---
+class GlaucomaModel(object):
+    def __init__(self, 
+                 cls_model_path="pamixsun/swinv2_tiny_for_glaucoma_classification", 
+                 seg_model_path='pamixsun/segformer_for_optic_disc_cup_segmentation',
+                 device=torch.device('cpu')):
+        self.device = device
+        # Classification model for glaucoma
+        self.cls_extractor = AutoImageProcessor.from_pretrained(cls_model_path)
+        self.cls_model = Swinv2ForImageClassification.from_pretrained(cls_model_path).to(device).eval()
+        # 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
+        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()
+        return disease_idx, confidence
+
+    def optic_disc_cup_pred(self, image):
+        inputs = self.seg_extractor(images=image.copy(), return_tensors="pt")
+        with torch.no_grad():
+            inputs.to(self.device)
+            outputs = self.seg_model(**inputs)
+        logits = outputs.logits.cpu()
+        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()
+        disc_confidence = seg_probs[0, 1, :, :].mean().item()
+        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)  # Assuming simple_vcdr() is defined elsewhere
+        except:
+            vcdr = np.nan
+        _, 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
+
+# --- Streamlit Interface ---
 def main():
-    # Wide mode
+    # Wide mode in Streamlit
     st.set_page_config(layout="wide")
 
-    # Designing the interface
     st.title("Glaucoma Screening from Retinal Fundus Images")
-    # For newline
-    st.write('\n')
-    # Author info
     st.write('Developed by X. Sun. Find more info about me: https://pamixsun.github.io')
-    # For newline
-    st.write('\n')
-    # Instructions
-    st.markdown("*Hint: click on the top-right corner of an image to enlarge it!*")
-    # Set the columns
+
+    # Set columns for the interface
     cols = st.beta_columns((1, 1, 1))
     cols[0].subheader("Input image")
     cols[1].subheader("Optic disc and optic cup")
-    cols[2].subheader("Class activation map")
-    
-    # set the visualization figure
-    fig, ax = plt.subplots()
+    cols[2].subheader("Classification Map")
 
-    # Sidebar
-    # File selection
+    # File uploader
     st.sidebar.title("Image selection")
-    # Disabling warning
     st.set_option('deprecation.showfileUploaderEncoding', False)
-    # Choose your own image
     uploaded_file = st.sidebar.file_uploader("Upload image", type=['png', 'jpeg', 'jpg'])
+
     if uploaded_file is not None:
-        # read the upload image
+        # Read and display uploaded image
         image = Image.open(uploaded_file).convert('RGB')
         image = np.array(image).astype(np.uint8)
-        # page_idx = 0
+        fig, ax = plt.subplots()
         ax.imshow(image)
         ax.axis('off')
         cols[0].pyplot(fig)
 
-    # For newline
-    st.sidebar.write('\n')
-
-    # actions
     if st.sidebar.button("Analyze image"):
-
         if uploaded_file is None:
             st.sidebar.write("Please upload an image")
-
         else:
             with st.spinner('Loading model...'):
-                # load 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...'):
-                # Forward the image to the model and get results
-                disease_idx, disc_cup_image, cam, vcdr = model.process(image)
+                # Get predictions from the model
+                disease_idx, disc_cup_image, vcdr, cls_confidence, cup_confidence, disc_confidence = model.process(image)
 
-                # plot the optic disc and optic cup image
+                # Display optic disc and cup image
                 ax.imshow(disc_cup_image)
                 ax.axis('off')
                 cols[1].pyplot(fig)
 
-                # plot the stitched image
-                ax.imshow(cam)
+                # Display classification map
+                ax.imshow(image)
                 ax.axis('off')
                 cols[2].pyplot(fig)
 
-                # Display JSON
-                st.subheader("  Screening results:")
-                st.write('\n')
-                
+                # 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]}|
+                |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}%|
                 """
                 st.markdown(final_results_as_table)
 
-
 if __name__ == '__main__':
     main()