Create app.py

app.py

@@ -0,0 +1,89 @@
+import torch
+import streamlit as st
+import numpy as np
+import cv2
+from PIL import Image
+import torchvision.transforms as transforms
+
+# Load models
+@st.cache_resource
+def load_models():
+    # Load classification model
+    classification_model = torch.load("resnet152_final.pt", map_location=torch.device("cpu"))
+    classification_model.eval()
+    
+    # Load segmentation model
+    segmentation_model = torch.load("best_unet_model.pth", map_location=torch.device("cpu"))
+    segmentation_model.eval()
+    
+    return classification_model, segmentation_model
+
+classification_model, segmentation_model = load_models()
+
+# Define preprocessing function for classification
+def preprocess_image(image):
+    # Convert to grayscale & extract green channel
+    image = np.array(image)
+    green_channel = image[:, :, 1]
+
+    # Apply CLAHE
+    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
+    img_clahe = clahe.apply(green_channel)
+
+    # Convert to tensor
+    transform = transforms.Compose([
+        transforms.ToTensor(),
+        transforms.Normalize(mean=[0.5], std=[0.5])
+    ])
+    
+    return transform(img_clahe).unsqueeze(0)
+
+# Define function for segmentation preprocessing (resize + normalize)
+def preprocess_segmentation(image):
+    transform = transforms.Compose([
+        transforms.Resize((512, 512)),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=[0.5], std=[0.5])
+    ])
+    
+    return transform(image).unsqueeze(0)
+
+# Streamlit UI
+st.title("Diabetic Retinopathy Detection, Classification & Segmentation")
+
+uploaded_file = st.file_uploader("Upload a Retinal Image", type=["jpg", "png", "jpeg"])
+if uploaded_file:
+    image = Image.open(uploaded_file)
+    st.image(image, caption="Uploaded Image", use_column_width=True)
+
+    # Preprocess image for classification
+    input_tensor = preprocess_image(image)
+
+    # Run classification
+    with torch.no_grad():
+        output = classification_model(input_tensor)
+        predicted_class = torch.argmax(output).item()
+
+    # Display result
+    dr_stages = ["No DR", "Mild", "Moderate", "Severe", "Proliferative DR"]
+    st.write(f"**Diabetic Retinopathy Stage:** {dr_stages[predicted_class]}")
+
+    # If DR detected, proceed to segmentation
+    if predicted_class > 0:
+        st.write("Lesion segmentation in progress...")
+
+        # Preprocess for segmentation
+        segmentation_input = preprocess_segmentation(image)
+
+        # Run segmentation
+        with torch.no_grad():
+            segmentation_output = segmentation_model(segmentation_input)
+        
+        # Convert output to mask
+        segmentation_mask = segmentation_output.squeeze().cpu().numpy()
+        segmentation_mask = (segmentation_mask > 0.5).astype(np.uint8) * 255
+
+        # Display segmentation result
+        st.image(segmentation_mask, caption="Segmented Lesions", use_column_width=True)
+    else:
+        st.write("No DR detected. Segmentation not required.")