Update app.py

app.py

@@ -8,8 +8,8 @@ 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"))
+    # Load classification model (corrected filename)
+    classification_model = torch.load("classifier.pt", map_location=torch.device("cpu"))
     classification_model.eval()
     
     # Load segmentation model
@@ -22,7 +22,7 @@ classification_model, segmentation_model = load_models()
 
 # Define preprocessing function for classification
 def preprocess_image(image):
-    # Convert to grayscale & extract green channel
+    # Convert to numpy array & extract green channel
     image = np.array(image)
     green_channel = image[:, :, 1]
 
@@ -38,7 +38,7 @@ def preprocess_image(image):
     
     return transform(img_clahe).unsqueeze(0)
 
-# Define function for segmentation preprocessing (resize + normalize)
+# Define function for segmentation preprocessing
 def preprocess_segmentation(image):
     transform = transforms.Compose([
         transforms.Resize((512, 512)),
@@ -83,7 +83,18 @@ if uploaded_file:
         segmentation_mask = segmentation_output.squeeze().cpu().numpy()
         segmentation_mask = (segmentation_mask > 0.5).astype(np.uint8) * 255
 
+        # Convert mask to image
+        mask_image = Image.fromarray(segmentation_mask)
+
         # Display segmentation result
-        st.image(segmentation_mask, caption="Segmented Lesions", use_column_width=True)
+        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"
+        )
     else:
         st.write("No DR detected. Segmentation not required.")