Update app.py

app.py

@@ -7,28 +7,28 @@ from PIL import Image
 import cv2
 
 ############################################
-# ========== UNET MODEL ====================
+# ========== ORIGINAL TRAINING UNET =========
 ############################################
 
 class DoubleConv(nn.Module):
     def __init__(self, in_channels, out_channels):
         super().__init__()
-        self.conv = nn.Sequential(
-            nn.Conv2d(in_channels, out_channels, 3, padding=1),
-            nn.ReLU(inplace=True),
-            nn.Conv2d(out_channels, out_channels, 3, padding=1),
+        self.conv_op = nn.Sequential(
+            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
             nn.ReLU(inplace=True),
+            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
+            nn.ReLU(inplace=True)
         )
 
     def forward(self, x):
-        return self.conv(x)
+        return self.conv_op(x)
 
 
 class DownSample(nn.Module):
     def __init__(self, in_channels, out_channels):
         super().__init__()
         self.conv = DoubleConv(in_channels, out_channels)
-        self.pool = nn.MaxPool2d(2)
+        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
 
     def forward(self, x):
         down = self.conv(x)
@@ -39,12 +39,12 @@ class DownSample(nn.Module):
 class UpSample(nn.Module):
     def __init__(self, in_channels, out_channels):
         super().__init__()
-        self.up = nn.ConvTranspose2d(in_channels, in_channels // 2, 2, 2)
+        self.up = nn.ConvTranspose2d(in_channels, in_channels//2, kernel_size=2, stride=2)
         self.conv = DoubleConv(in_channels, out_channels)
 
     def forward(self, x1, x2):
         x1 = self.up(x1)
-        x = torch.cat([x1, x2], dim=1)
+        x = torch.cat([x1, x2], 1)
         return self.conv(x)
 
 
@@ -52,34 +52,34 @@ class UNet(nn.Module):
     def __init__(self, in_channels=3, num_classes=1):
         super().__init__()
 
-        self.down1 = DownSample(in_channels, 64)
-        self.down2 = DownSample(64, 128)
-        self.down3 = DownSample(128, 256)
-        self.down4 = DownSample(256, 512)
+        self.down_convolution_1 = DownSample(in_channels, 64)
+        self.down_convolution_2 = DownSample(64, 128)
+        self.down_convolution_3 = DownSample(128, 256)
+        self.down_convolution_4 = DownSample(256, 512)
 
-        self.bottleneck = DoubleConv(512, 1024)
+        self.bottle_neck = DoubleConv(512, 1024)
 
-        self.up1 = UpSample(1024, 512)
-        self.up2 = UpSample(512, 256)
-        self.up3 = UpSample(256, 128)
-        self.up4 = UpSample(128, 64)
+        self.up_convolution_1 = UpSample(1024, 512)
+        self.up_convolution_2 = UpSample(512, 256)
+        self.up_convolution_3 = UpSample(256, 128)
+        self.up_convolution_4 = UpSample(128, 64)
 
-        self.out = nn.Conv2d(64, num_classes, kernel_size=1)
+        self.out = nn.Conv2d(in_channels=64, out_channels=num_classes, kernel_size=1)
 
     def forward(self, x):
-        d1, p1 = self.down1(x)
-        d2, p2 = self.down2(p1)
-        d3, p3 = self.down3(p2)
-        d4, p4 = self.down4(p3)
+        down_1, p1 = self.down_convolution_1(x)
+        down_2, p2 = self.down_convolution_2(p1)
+        down_3, p3 = self.down_convolution_3(p2)
+        down_4, p4 = self.down_convolution_4(p3)
 
-        b = self.bottleneck(p4)
+        b = self.bottle_neck(p4)
 
-        u1 = self.up1(b, d4)
-        u2 = self.up2(u1, d3)
-        u3 = self.up3(u2, d2)
-        u4 = self.up4(u3, d1)
+        up_1 = self.up_convolution_1(b, down_4)
+        up_2 = self.up_convolution_2(up_1, down_3)
+        up_3 = self.up_convolution_3(up_2, down_2)
+        up_4 = self.up_convolution_4(up_3, down_1)
 
-        return self.out(u4)
+        return self.out(up_4)
 
 
 ############################################
@@ -88,7 +88,7 @@ class UNet(nn.Module):
 
 device = torch.device("cpu")
 
-model = UNet()
+model = UNet(in_channels=3, num_classes=1)
 model.load_state_dict(torch.load("my_checkpoint.pth", map_location=device))
 model.eval()
 
@@ -112,7 +112,7 @@ def dice_coefficient(pred, target, epsilon=1e-7):
     return ((2. * intersection + epsilon) / (union + epsilon)).item()
 
 ############################################
-# ========== PREPROCESS TIFF ===============
+# ========== TIFF SAFE LOADER ==============
 ############################################
 
 def load_image(file):
@@ -126,7 +126,6 @@ def load_image(file):
     img_pil = Image.fromarray(img_np).convert("RGB")
     return img_pil, img_np
 
-
 ############################################
 # ========== PREDICTION ====================
 ############################################
@@ -137,7 +136,6 @@ def predict(image_file, mask_file=None):
         return None, "Please upload an image."
 
     image_pil, original_np = load_image(image_file)
-
     input_tensor = transform(image_pil).unsqueeze(0)
 
     with torch.no_grad():
@@ -145,19 +143,17 @@ def predict(image_file, mask_file=None):
         output = torch.sigmoid(output)
 
     pred_mask = output.squeeze().numpy()
-    pred_mask_binary = (pred_mask > 0.5).astype(np.uint8)
+    pred_binary = (pred_mask > 0.5).astype(np.uint8)
 
-    # Resize mask to original image size
-    pred_mask_resized = cv2.resize(
-        pred_mask_binary,
+    # Resize mask back to original size
+    pred_resized = cv2.resize(
+        pred_binary,
         (original_np.shape[1], original_np.shape[0])
     )
 
-    # Create red overlay
     overlay = original_np.copy()
-    overlay[pred_mask_resized == 1] = [255, 0, 0]
+    overlay[pred_resized == 1] = [255, 0, 0]
 
-    # If mask provided → compute Dice
     if mask_file is not None:
         mask_pil, _ = load_image(mask_file)
         mask_tensor = transform(mask_pil.convert("L"))
@@ -166,20 +162,18 @@ def predict(image_file, mask_file=None):
 
     return overlay, "Prediction complete."
 
-
 ############################################
 # ========== GRADIO UI =====================
 ############################################
 
 description = """
-# 🧠 Brain Tumor Segmentation using UNet
-
-This model was trained on:
+# 🧠 Brain Tumor Segmentation (UNet)
 
-🔗 https://www.kaggle.com/datasets/mateuszbuda/lgg-mri-segmentation
+Dataset used for training:
+https://www.kaggle.com/datasets/mateuszbuda/lgg-mri-segmentation
 
-Upload a `.tif` MRI image to predict tumor segmentation.
-Optionally upload the true mask to compute Dice score.
+Upload a `.tif` MRI image.
+Optionally upload the ground-truth mask to compute Dice score.
 """
 
 demo = gr.Interface(