luigi12345's picture
m
f192f0b
raw
history blame
8.53 kB
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
# --- 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()
# Mapping for class labels
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]
# Calculate segmentation confidence based on probability distribution
# For each pixel classified as cup/disc, check how confident the model is
cup_mask = pred_disc_cup == 2
disc_mask = pred_disc_cup == 1
# Get confidence only for pixels predicted as cup/disc
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
# --- Utility Functions ---
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"
# --- Streamlit Interface ---
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)
# Sidebar settings
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)
# Image display section
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("---")
# Classification Results
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")
# Segmentation Results
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}")
# Clinical Metrics
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("---")
# ... rest of the code remains the same ...