|
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 streamlit as st |
|
from PIL import Image |
|
import io |
|
import zipfile |
|
|
|
|
|
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 |
|
|
|
self.cls_extractor = AutoImageProcessor.from_pretrained(cls_model_path) |
|
self.cls_model = Swinv2ForImageClassification.from_pretrained(cls_model_path).to(device).eval() |
|
|
|
self.seg_extractor = AutoImageProcessor.from_pretrained(seg_model_path) |
|
self.seg_model = SegformerForSemanticSegmentation.from_pretrained(seg_model_path).to(device).eval() |
|
|
|
self.cls_id2label = self.cls_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 |
|
probs = F.softmax(outputs, dim=-1) |
|
disease_idx = probs.cpu()[0, :].numpy().argmax() |
|
confidence = probs.cpu()[0, disease_idx].item() * 100 |
|
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 |
|
) |
|
seg_probs = F.softmax(upsampled_logits, dim=1) |
|
pred_disc_cup = upsampled_logits.argmax(dim=1)[0] |
|
|
|
|
|
|
|
cup_mask = pred_disc_cup == 2 |
|
disc_mask = pred_disc_cup == 1 |
|
|
|
|
|
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): |
|
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) |
|
except: |
|
vcdr = np.nan |
|
|
|
mask = (disc_cup > 0).astype(np.uint8) |
|
x, y, w, h = cv2.boundingRect(mask) |
|
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) |
|
|
|
cropped_image = image[y:y+h, x:x+w] if w >= 50 and h >= 50 else image.copy() |
|
_, 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 |
|
|
|
|
|
def simple_vcdr(mask): |
|
disc_area = np.sum(mask == 1) |
|
cup_area = np.sum(mask == 2) |
|
if disc_area == 0: |
|
return np.nan |
|
vcdr = cup_area / disc_area |
|
return vcdr |
|
|
|
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 |
|
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" |
|
|
|
|
|
def main(): |
|
st.set_page_config(layout="wide", page_title="Glaucoma Screening Tool") |
|
|
|
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) |
|
|
|
|
|
st.sidebar.markdown("### Upload Settings") |
|
uploaded_files = st.sidebar.file_uploader("Upload Retinal Images", |
|
type=['png', 'jpeg', 'jpg'], |
|
accept_multiple_files=True, |
|
help="Support multiple images in PNG, JPEG formats") |
|
|
|
st.sidebar.markdown("### Analysis Settings") |
|
st.sidebar.info("π Set confidence threshold to filter results") |
|
confidence_threshold = st.sidebar.slider( |
|
"Classification Confidence Threshold (%)", |
|
0, 100, 70, |
|
help="Images with confidence above this threshold will be marked as reliable predictions") |
|
|
|
if uploaded_files: |
|
for uploaded_file in uploaded_files: |
|
image = Image.open(uploaded_file).convert('RGB') |
|
image_np = np.array(image).astype(np.uint8) |
|
|
|
st.markdown(f"## π Analysis Results: {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) |
|
|
|
|
|
st.subheader("Original Image") |
|
st.image(image_np, use_column_width=True) |
|
st.subheader("Segmentation Overlay") |
|
st.image(disc_cup_image, use_column_width=True) |
|
st.subheader("Region of Interest") |
|
st.image(cropped_image, use_column_width=True) |
|
|
|
st.markdown("---") |
|
|
|
|
|
st.markdown("### π Classification") |
|
diagnosis = model.cls_id2label[disease_idx] |
|
is_confident = cls_conf >= confidence_threshold |
|
|
|
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.write(f"Classification Confidence: {cls_conf:.1f}%") |
|
if not is_confident: |
|
st.warning("β οΈ Below confidence threshold") |
|
|
|
|
|
st.markdown("### π― Segmentation Quality") |
|
st.write(f"Optic Cup Confidence: {cup_conf:.1f}%") |
|
st.write(f"Optic Disc Confidence: {disc_conf:.1f}%") |
|
|
|
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}") |
|
|
|
|
|
st.markdown("### π Clinical Metrics") |
|
st.write(f"Cup-to-Disc Ratio (CDR): {vcdr:.3f}") |
|
|
|
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)") |
|
|
|
st.markdown("---") |
|
|
|
|