Spaces:
Sleeping
Sleeping
| 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 | |
| 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) | |
| model = SqueezeNet(out_channels=1) | |
| model.load_state_dict(torch.load("squeezenet.pth", map_location=torch.device('cpu'))) | |
| model.eval() | |
| transform = transforms.Compose([transforms.Resize((224,224)), | |
| transforms.RandomHorizontalFlip(0.2), | |
| transforms.RandomVerticalFlip(0.1), | |
| transforms.RandomAutocontrast(0.2), | |
| transforms.RandomAdjustSharpness(0.3), | |
| transforms.ToTensor() | |
| ]) | |
| def classify_brain_tumor(image): | |
| image = transform(image).unsqueeze(0) | |
| with torch.no_grad(): | |
| output = model(image) | |
| prediction = torch.sigmoid(output).item() | |
| return "Tumor" if prediction >= 0.5 else "No Tumor" | |
| interface = gr.Interface( | |
| fn=classify_brain_tumor, | |
| inputs=gr.Image(type="pil"), | |
| outputs="text", | |
| title="Brain Tumor Classification", | |
| description="Upload an MRI image to classify if it has a tumor or not." | |
| ) | |
| interface.launch() | |