Update app.py

app.py

@@ -1,46 +1,55 @@
 import torch
 import streamlit as st
-import torchvision.models as models
 import numpy as np
 import cv2
 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)
 
+# Load models
 @st.cache_resource
 def load_models():
-    # Load classification model architecture
-    classification_model = models.resnet152(pretrained=False)  # Use the same architecture as before
-    num_ftrs = classification_model.fc.in_features
-    classification_model.fc = torch.nn.Linear(num_ftrs, 5)  # Assuming 5 DR stages
-    classification_checkpoint = torch.load("classifier.pt", map_location=torch.device("cpu"))
-    classification_model.load_state_dict(classification_checkpoint["model_state_dict"])
+    # 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
 
-
 classification_model, segmentation_model = load_models()
 
 # Define preprocessing function for classification
 def preprocess_image(image):
-    # Convert to numpy array & extract green channel
     image = np.array(image)
-    green_channel = image[:, :, 1]
-
-    # Apply CLAHE
+    green_channel = image[:, :, 1]  # Extract green channel
+    
     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
@@ -50,7 +59,6 @@ def preprocess_segmentation(image):
         transforms.ToTensor(),
         transforms.Normalize(mean=[0.5], std=[0.5])
     ])
-    
     return transform(image).unsqueeze(0)
 
 # Streamlit UI
@@ -60,46 +68,37 @@ uploaded_file = st.file_uploader("Upload a Retinal Image", type=["jpg", "png", "
 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
-
-        # Convert mask to image
+        
+        st.image(segmentation_mask, caption="Segmented Lesions", use_column_width=True)
+        
+        # Provide download option
         mask_image = Image.fromarray(segmentation_mask)
-
-        # Display segmentation result
-        st.image(mask_image, caption="Segmented Lesions", use_column_width=True)
-
-        # Provide download button
-        st.download_button(
-            label="Download Segmented Mask",
-            data=mask_image.tobytes(),
-            file_name="segmented_mask.png",
-            mime="image/png"
-        )
+        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.")