Update app.py

app.py

@@ -1,31 +1,73 @@
 import torch
 import torch.nn as nn
+import torch.nn.functional as F
 from torchvision import transforms
 from PIL import Image
 import gradio as gr
 
-model = torch.load("squeezenet.pth", map_location=torch.device('cpu'))
-model.eval()
+class FireModule(nn.Module):
+    def __init__(self, in_channels, s1x1, e1x1, e3x3):
+        super(FireModule, self).__init__()
+        self.squeeze = nn.Conv2d(in_channels=in_channels, out_channels=s1x1, kernel_size=1, stride=1)
+        self.expand1x1 = nn.Conv2d(in_channels=s1x1, out_channels=e1x1, kernel_size=1)
+        self.expand3x3 = nn.Conv2d(in_channels=s1x1, out_channels=e3x3, kernel_size=3, padding=1)
+        
+    def forward(self, x):
+        x = F.relu(self.squeeze(x)) 
+        x1 = self.expand1x1(x)   
+        x2 = self.expand3x3(x)      
+        x = F.relu(torch.cat((x1, x2), dim=1))  
+        return x
+
+class SqueezeNet(nn.Module):
+    def __init__(self, out_channels):
+        super(SqueezeNet, self).__init__()
+        self.conv1 = nn.Conv2d(in_channels=3, out_channels=96, kernel_size=7, stride=2)
+        self.max_pool1 = nn.MaxPool2d(kernel_size=3, stride=2)
+        self.fire2 = FireModule(in_channels=96, s1x1=16, e1x1=64, e3x3=64)
+        self.fire3 = FireModule(in_channels=128, s1x1=16, e1x1=64, e3x3=64)  
+        self.fire4 = FireModule(in_channels=128, s1x1=32, e1x1=128, e3x3=128)
+        self.max_pool2 = nn.MaxPool2d(kernel_size=3, stride=2)
+        self.fire5 = FireModule(in_channels=256, s1x1=32, e1x1=128, e3x3=128)
+        self.fire6 = FireModule(in_channels=256, s1x1=48, e1x1=192, e3x3=192)
+        self.fire7 = FireModule(in_channels=384, s1x1=48, e1x1=192, e3x3=192)
+        self.fire8 = FireModule(in_channels=384, s1x1=64, e1x1=256, e3x3=256)
+        self.max_pool3 = nn.MaxPool2d(kernel_size=3, stride=2)
+        self.fire9 = FireModule(in_channels=512, s1x1=64, e1x1=256, e3x3=256)
+        self.conv10 = nn.Conv2d(in_channels=512, out_channels=out_channels, kernel_size=1, stride=1)
+        self.avgpool = nn.AvgPool2d(kernel_size=12, stride=1)
+        
+    def forward(self, x):
+        x = self.max_pool1(self.conv1(x))
+        x = self.max_pool2(self.fire4(self.fire3(self.fire2(x))))
+        x = self.max_pool3(self.fire8(self.fire7(self.fire6(self.fire5(x)))))
+        x = self.avgpool(self.conv10(self.fire9(x)))
+        return torch.flatten(x, start_dim=1)
+
+# Initialize the model and load weights
+model = SqueezeNet(out_channels=1)  # Adjust output channels if needed
+model.load_state_dict(torch.load("squeezenet.pth", map_location=torch.device('cpu')))
+model.eval()  
 
 transform = transforms.Compose([
-    transforms.Resize((128, 128)),
-    transforms.ToTensor(),
-    transforms.Normalize([0.5], [0.5])
+    transforms.Resize((128, 128)),   # Resize to match model's input size
+    transforms.ToTensor(),           # Convert to tensor
+    transforms.Normalize([0.5], [0.5])  # Normalize based on training dataset
 ])
 
 def classify_brain_tumor(image):
-    image = transform(image).unsqueeze(0)
+    image = transform(image).unsqueeze(0)  # Preprocess and add batch dimension
     with torch.no_grad():
         output = model(image)
-        _, predicted = torch.max(output, 1)
-    return "Tumor" if predicted.item() == 1 else "No Tumor"
+        prediction = torch.sigmoid(output).item()  # Apply sigmoid for binary classification
+    return "Tumor" if prediction >= 0.5 else "No Tumor"
 
 interface = gr.Interface(
     fn=classify_brain_tumor,
     inputs=gr.inputs.Image(type="pil"),
     outputs="text",
     title="Brain Tumor Classification",
-    description="Upload an MRI image to classify if it has a tumor or not. The Model is SqueezeNet."
+    description="Upload an MRI image to classify if it has a tumor or not."
 )
 
-interface.launch()
+interface.launch()