Update app.py

app.py

@@ -11,13 +11,74 @@ from typing import Tuple
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
 # Constants
-CLASS_NAMES = ["No DR", "Mild", "Moderate", "Severe", "Proliferative DR"]
+CLASS_NAMES = ["No_DR", "Mild", "Moderate", "Severe", "Proliferate_DR"]
 LESION_COLORS = {
     0: [0, 0, 0],      # Background (black)
     1: [255, 255, 0],  # Bright lesions (yellow)
     2: [255, 0, 0]     # Red lesions (red)
 }
 
+# ====================== UNET ARCHITECTURE ======================
+class UNet(nn.Module):
+    def __init__(self, input_channels=3, num_classes=5):
+        super(UNet, self).__init__()
+        
+        def conv_block(in_channels, out_channels):
+            return nn.Sequential(
+                nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
+                nn.ReLU(inplace=True),
+                nn.BatchNorm2d(out_channels),
+                nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
+                nn.ReLU(inplace=True),
+                nn.BatchNorm2d(out_channels),
+            )
+
+        self.encoder1 = conv_block(input_channels, 32)
+        self.pool1 = nn.MaxPool2d(2)
+        self.encoder2 = conv_block(32, 64)
+        self.pool2 = nn.MaxPool2d(2)
+        self.encoder3 = conv_block(64, 128)
+        self.pool3 = nn.MaxPool2d(2)
+
+        self.bottleneck = conv_block(128, 256)
+
+        self.up3 = nn.ConvTranspose2d(256, 128, kernel_size=2, stride=2)
+        self.decoder3 = conv_block(256, 128)
+
+        self.up2 = nn.ConvTranspose2d(128, 64, kernel_size=2, stride=2)
+        self.decoder2 = conv_block(128, 64)
+
+        self.up1 = nn.ConvTranspose2d(64, 32, kernel_size=2, stride=2)
+        self.decoder1 = conv_block(64, 32)
+
+        self.final_conv = nn.Conv2d(32, num_classes, kernel_size=1)
+
+    def forward(self, x):
+        enc1 = self.encoder1(x)
+        x = self.pool1(enc1)
+        
+        enc2 = self.encoder2(x)
+        x = self.pool2(enc2)
+        
+        enc3 = self.encoder3(x)
+        x = self.pool3(enc3)
+        
+        x = self.bottleneck(x)
+        
+        x = self.up3(x)
+        x = torch.cat([x, enc3], dim=1)
+        x = self.decoder3(x)
+        
+        x = self.up2(x)
+        x = torch.cat([x, enc2], dim=1)
+        x = self.decoder2(x)
+        
+        x = self.up1(x)
+        x = torch.cat([x, enc1], dim=1)
+        x = self.decoder1(x)
+        
+        return self.final_conv(x)
+
 # ====================== CLASSIFIER ======================
 def create_classifier_model():
     model = models.resnet152(pretrained=False)
@@ -38,7 +99,6 @@ def load_classifier():
     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))
@@ -54,12 +114,12 @@ def get_classifier_transform():
 # ====================== SEGMENTATION ======================
 @st.cache_resource
 def load_segmenter():
-    model = torch.load('best_unet_model.pth', map_location=device)
+    model = UNet().to(device)
+    model.load_state_dict(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)
@@ -76,20 +136,18 @@ def get_segmenter_transform():
     ])
 
 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
+    final_mask[(pred_class == 1) | (pred_class == 4)] = 1
+    final_mask[(pred_class == 2) | (pred_class == 3)] = 2
     return final_mask, probs
 
 # ====================== 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)
+                          interpolation=cv2.INTER_NEAREST)
     
     overlay = original_np.copy()
     for class_idx, color in LESION_COLORS.items():
@@ -97,10 +155,27 @@ def create_lesion_overlay(original: Image.Image, mask: np.ndarray) -> Image.Imag
     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)
 
+def segment_image(image: Image.Image, model: nn.Module) -> dict:
+    processed_img = preprocess_segmenter(image)
+    img_pil = Image.fromarray(processed_img)
+    transform = get_segmenter_transform()
+    image_tensor = transform(img_pil).unsqueeze(0).to(device)
+    
+    with torch.no_grad():
+        output = model(image_tensor)
+    
+    final_mask, class_probs = process_segmentation_output(output)
+    total_pixels = final_mask.size
+    return {
+        'mask': final_mask,
+        'probs': class_probs,
+        'bright_area': (np.sum(final_mask == 1) / total_pixels * 100),
+        'red_area': (np.sum(final_mask == 2) / total_pixels * 100)
+    }
+
 # ====================== MAIN APP ======================
 def main():
     st.set_page_config(layout="wide")
@@ -116,13 +191,12 @@ def main():
         col1, col2 = st.columns(2)
         
         with col1:
-            st.image(original_image, caption="Original Image", use_column_width=True)
+            st.image(original_image, caption="Original Image", use_container_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)
+        img_tensor = get_classifier_transform()(Image.fromarray(clf_processed)).unsqueeze(0).to(device)
         
         with torch.no_grad():
             logps = classifier(img_tensor)
@@ -151,9 +225,9 @@ def main():
                                         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)
+                    st.image(overlay, caption="Lesion Overlay", use_container_width=True)
+                    st.image(heat_bright, caption="Bright Lesion Probability", use_container_width=True)
+                    st.image(heat_red, caption="Red Lesion Probability", use_container_width=True)
 
                 # Metrics
                 st.write("**Lesion Analysis:**")