Update app.py

app.py

@@ -1,104 +1,178 @@
-import torch
-import streamlit as st
-import numpy as np
 import cv2
+import numpy as np
 from PIL import Image
-import torchvision.transforms as transforms
-import torchvision.models as models
-import io
-
-# Define the classification model structure (ResNet152 with modified FC layer)
-class ResNetClassifier(torch.nn.Module):
-    def __init__(self, num_classes=5):
-        super(ResNetClassifier, self).__init__()
-        self.model = models.resnet152(pretrained=False)
-        self.model.fc = torch.nn.Sequential(
-            torch.nn.Linear(self.model.fc.in_features, 512),
-            torch.nn.ReLU(),
-            torch.nn.Linear(512, num_classes)
-        )
-    
-    def forward(self, x):
-        return self.model(x)
+import torch
+import torch.nn as nn
+from torchvision import models, transforms
+import streamlit as st
+from typing import Tuple
 
-# Load models
-@st.cache_resource
-def load_models():
-    # Load classification model
-    classification_model = ResNetClassifier()
-    classifier_checkpoint = torch.load("classifier.pt", map_location=torch.device("cpu"))
-    classification_model.load_state_dict(classifier_checkpoint["model_state_dict"])
-    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
+# Device configuration
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
-classification_model, segmentation_model = load_models()
+# Constants
+CLASS_NAMES = ["No DR", "Mild", "Moderate", "Severe", "Proliferative DR"]
+LESION_COLORS = {
+    0: [0, 0, 0],      # Background (black)
+    1: [255, 255, 0],  # Bright lesions (yellow)
+    2: [255, 0, 0]     # Red lesions (red)
+}
 
-# Define preprocessing function for classification
-def preprocess_image(image):
-    image = np.array(image)
-    green_channel = image[:, :, 1]  # Extract green channel
-    
+# ====================== CLASSIFIER ======================
+def create_classifier_model():
+    model = models.resnet152(pretrained=False)
+    num_ftrs = model.fc.in_features
+    model.fc = nn.Sequential(
+        nn.Linear(num_ftrs, 512),
+        nn.ReLU(),
+        nn.Linear(512, 5),
+        nn.LogSoftmax(dim=1)
+    return model
+
+@st.cache_resource
+def load_classifier():
+    model = create_classifier_model().to(device)
+    checkpoint = torch.load('classifier.pt', map_location=device)
+    model.load_state_dict(checkpoint['model_state_dict'])
+    model.eval()
+    return model
+
+def preprocess_classifier(image: Image.Image) -> np.ndarray:
+    """Green channel + CLAHE preprocessing"""
+    img_np = np.array(image)
+    green_channel = img_np[:, :, 1]
     clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
-    img_clahe = clahe.apply(green_channel)
-    
-    transform = transforms.Compose([
+    return np.stack([clahe.apply(green_channel)]*3, axis=-1)
+
+def get_classifier_transform():
+    return transforms.Compose([
+        transforms.Resize((224, 224)),
         transforms.ToTensor(),
-        transforms.Normalize(mean=[0.5], std=[0.5])
+        transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
     ])
-    return transform(img_clahe).unsqueeze(0)
 
-# Define function for segmentation preprocessing
-def preprocess_segmentation(image):
-    transform = transforms.Compose([
+# ====================== SEGMENTATION ======================
+@st.cache_resource
+def load_segmenter():
+    model = torch.load('best_unet_model.pth', map_location=device)
+    model.eval()
+    return model
+
+def preprocess_segmenter(image: Image.Image) -> np.ndarray:
+    """LAB + CLAHE + Median filtering"""
+    img_np = np.array(image)
+    img_filtered = cv2.medianBlur(img_np, 3)
+    lab = cv2.cvtColor(img_filtered, cv2.COLOR_RGB2LAB)
+    l, a, b = cv2.split(lab)
+    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
+    lab_clahe = cv2.merge((clahe.apply(l), a, b))
+    return cv2.cvtColor(lab_clahe, cv2.COLOR_LAB2RGB)
+
+def get_segmenter_transform():
+    return transforms.Compose([
         transforms.Resize((512, 512)),
         transforms.ToTensor(),
-        transforms.Normalize(mean=[0.5], std=[0.5])
+        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
     ])
-    return transform(image).unsqueeze(0)
 
-# Streamlit UI
-st.title("Diabetic Retinopathy Detection & Segmentation")
+def process_segmentation_output(output: torch.Tensor) -> Tuple[np.ndarray, np.ndarray]:
+    """Convert 5-class output to 3-class mask"""
+    probs = torch.softmax(output, dim=1).cpu().numpy().squeeze()
+    pred_class = np.argmax(probs, axis=0)
+    final_mask = np.zeros_like(pred_class, dtype=np.uint8)
+    final_mask[(pred_class == 1) | (pred_class == 4)] = 1  # Bright
+    final_mask[(pred_class == 2) | (pred_class == 3)] = 2  # Red
+    return final_mask, probs
 
-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)
+# ====================== VISUALIZATION ======================
+def create_lesion_overlay(original: Image.Image, mask: np.ndarray) -> Image.Image:
+    """Color-coded lesion overlay"""
+    original_np = np.array(original)
+    mask_resized = cv2.resize(mask, (original_np.shape[1], original_np.shape[0]), 
+                            interpolation=cv2.INTER_NEAREST)
     
-    # Run classification
-    with torch.no_grad():
-        output = classification_model(input_tensor)
-        predicted_class = torch.argmax(output).item()
-    
-    dr_stages = ["No DR", "Mild", "Moderate", "Severe", "Proliferative DR"]
-    st.write(f"**Diabetic Retinopathy Stage:** {dr_stages[predicted_class]}")
+    overlay = original_np.copy()
+    for class_idx, color in LESION_COLORS.items():
+        overlay[mask_resized == class_idx] = color
+    return Image.fromarray(cv2.addWeighted(overlay, 0.4, original_np, 0.6, 0))
+
+def create_heatmap(prob_map: np.ndarray, original_size: Tuple[int, int]) -> np.ndarray:
+    """Probability heatmap visualization"""
+    resized = cv2.resize(prob_map, original_size)
+    return cv2.applyColorMap((resized * 255).astype(np.uint8), cv2.COLORMAP_JET)
+
+# ====================== MAIN APP ======================
+def main():
+    st.set_page_config(layout="wide")
+    st.title("Diabetic Retinopathy Analysis")
     
-    # If DR detected, proceed to segmentation
-    if predicted_class > 0:
-        st.write("Lesion segmentation in progress...")
+    uploaded_file = st.file_uploader("Upload retinal image", type=["jpg", "jpeg", "png"])
+    if not uploaded_file:
+        st.info("Please upload an image")
+        return
+
+    try:
+        original_image = Image.open(uploaded_file).convert('RGB')
+        col1, col2 = st.columns(2)
         
-        segmentation_input = preprocess_segmentation(image)
+        with col1:
+            st.image(original_image, caption="Original Image", use_column_width=True)
+
+        # Classification
+        classifier = load_classifier()
+        clf_processed = preprocess_classifier(original_image)
+        clf_transform = get_classifier_transform()
+        img_tensor = clf_transform(Image.fromarray(clf_processed)).unsqueeze(0).to(device)
         
         with torch.no_grad():
-            segmentation_output = segmentation_model(segmentation_input)
-        
-        segmentation_mask = segmentation_output.squeeze().cpu().numpy()
-        segmentation_mask = (segmentation_mask > 0.5).astype(np.uint8) * 255
-        
-        st.image(segmentation_mask, caption="Segmented Lesions", use_column_width=True)
-        
-        # Provide download option
-        mask_image = Image.fromarray(segmentation_mask)
-        buf = io.BytesIO()
-        mask_image.save(buf, format="PNG")
-        byte_im = buf.getvalue()
-        st.download_button("Download Segmentation Mask", data=byte_im, file_name="segmentation_mask.png", mime="image/png")
-    else:
-        st.write("No DR detected. Segmentation not required.")
+            logps = classifier(img_tensor)
+            ps = torch.exp(logps)
+            pred_class = torch.argmax(ps).item()
+            probabilities = ps[0].cpu().numpy() * 100
+
+        st.subheader("Classification Results")
+        if pred_class == 0:
+            st.success(f"**Prediction:** {CLASS_NAMES[pred_class]}")
+        else:
+            st.error(f"**Prediction:** {CLASS_NAMES[pred_class]}")
+            st.write("**Confidence Levels:**")
+            for name, prob in zip(CLASS_NAMES, probabilities):
+                st.progress(int(prob))
+                st.write(f"{name}: {prob:.1f}%")
+
+            # Segmentation
+            segmenter = load_segmenter()
+            with st.spinner("Detecting lesions..."):
+                seg_results = segment_image(original_image, segmenter)
+                overlay = create_lesion_overlay(original_image, seg_results['mask'])
+                heat_bright = create_heatmap(seg_results['probs'][1] + seg_results['probs'][4], 
+                                           original_image.size)
+                heat_red = create_heatmap(seg_results['probs'][2] + seg_results['probs'][3], 
+                                        original_image.size)
+
+                with col2:
+                    st.image(overlay, caption="Lesion Overlay", use_column_width=True)
+                    st.image(heat_bright, caption="Bright Lesion Probability", use_column_width=True)
+                    st.image(heat_red, caption="Red Lesion Probability", use_column_width=True)
+
+                # Metrics
+                st.write("**Lesion Analysis:**")
+                cols = st.columns(3)
+                cols[0].metric("Bright Lesions", f"{seg_results['bright_area']:.2f}%")
+                cols[1].metric("Red Lesions", f"{seg_results['red_area']:.2f}%")
+                cols[2].metric("Total Affected", 
+                              f"{seg_results['bright_area'] + seg_results['red_area']:.2f}%")
+
+                # Download
+                st.download_button(
+                    "Download Mask",
+                    cv2.imencode('.png', seg_results['mask'] * 85)[1].tobytes(),
+                    "dr_mask.png",
+                    "image/png"
+                )
+
+    except Exception as e:
+        st.error(f"Error processing image: {str(e)}")
+
+if __name__ == "__main__":
+    main()