Upload 5 files

Dockerfile

@@ -0,0 +1,13 @@
+# lightweight python
+FROM python:3.11
+
+
+# Copy local code to the container image.
+WORKDIR /app
+COPY . /app
+
+# Install dependencies
+RUN pip install -r requirements.txt
+EXPOSE 8051
+# Run the streamlit on container startup
+CMD ["streamlit", "run", "app.py"]

app.py

@@ -0,0 +1,66 @@
+from unet import UNet
+import streamlit as st
+import torch
+from torchvision import transforms
+import albumentations as A
+from albumentations.pytorch import ToTensorV2
+from PIL import Image
+import numpy as np
+
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+model = UNet(3, 1, [64, 128, 256, 512]).to(device)
+model.load_state_dict(torch.load("best_model.pth", map_location=torch.device(device)))
+# Set up transformations for the input image
+
+
+transform =  A.Compose([
+        A.Resize(224, 224, p=1.0),
+        ToTensorV2(),
+    ])
+# Streamlit app
+def main():
+    st.title("MRI segmenation App")
+
+    # Upload image through Streamlit
+    uploaded_image = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
+
+    if uploaded_image is not None:
+        # Display the uploaded and processed images side by side
+        col1, col2 = st.columns(2)  # Using beta_columns for side-by-side layout
+
+        # Display the uploaded image in the first column
+        col1.header("Original Image")
+        col1.image(uploaded_image, caption="Uploaded Image", use_column_width=True)
+
+        # Process the image (replace this with your processing logic)
+        processed_image = generate_image(uploaded_image)
+
+        # Display the processed image in the second column
+        col2.header("Processed Image")
+        col2.image(processed_image, caption="Processed Image", use_column_width=True)
+
+# Function to generate an image using the PyTorch model
+def generate_image(uploaded_image):
+    # Load the uploaded image
+    input_image = Image.open(uploaded_image)
+
+    image = np.array(input_image).astype(np.float32) / 255.
+    # Apply transformations
+    input_tensor = transform(image=image)["image"].unsqueeze(0)
+
+    # Generate an image using the PyTorch model
+    with torch.no_grad():
+        input_tensor = input_tensor.type(torch.FloatTensor).to(device)
+        pred = model(input_tensor)
+        pred = torch.sigmoid(pred)
+        mask = (pred > 0.6).float()
+    mask = mask[0].permute(1, 2, 0)
+    image = input_tensor[0].permute(1, 2, 0)
+    
+    mask = image + mask*0.3
+    mask = mask.permute(2, 0, 1)
+    mask = transforms.ToPILImage()(mask)
+    return mask
+
+if __name__ == "__main__":
+    main()

best_model.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:80bca9fca6dd62bb74e5072bcac8a4e2d232d43b6f45e0202bf6d5a353cd2b70
+size 124203732

requirements.txt

@@ -0,0 +1,10 @@
+# Core Libraries
+numpy>=1.21.0
+
+pillow>=9.3.0
+
+# Web Application Framework
+streamlit>=1.28.0
+torch>=2.1.2
+torchvision>=0.16.2
+albumentations>=1.3.1

unet.py

@@ -0,0 +1,102 @@
+import torch
+import torch.nn as nn
+
+class DownSampling(nn.Module):
+    
+    def __init__(self, in_channels, out_channels, max_pool):
+        """
+        DownSampling block in the U-Net architecture.
+
+        Args:
+            in_channels (int): Number of input channels.
+            out_channels (int): Number of output channels.
+            max_pool (bool): Whether to use max pooling.
+        """
+        super(DownSampling, self).__init__()
+        self.max_pool = max_pool
+        self.conv1 = nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=3, stride=1, padding=1)
+        self.conv2 = nn.Conv2d(in_channels=out_channels, out_channels=out_channels, kernel_size=3, stride=1, padding=1)
+        self.batchnorm2d = nn.BatchNorm2d(out_channels)
+        self.relu = nn.ReLU()
+        self.maxpool2d = nn.MaxPool2d(kernel_size=2, stride=2)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.conv2(x)
+
+        x = self.relu(self.batchnorm2d(x))
+        skip_connection = x
+        
+        if self.max_pool:
+            next_layer = self.maxpool2d(x)
+        else:
+            return x
+        return next_layer, skip_connection
+
+class UpSampling(nn.Module):
+    def __init__(self, in_channels, out_channels):
+        """
+        UpSampling block in the U-Net architecture.
+
+        Args:
+            in_channels (int): Number of input channels.
+            out_channels (int): Number of output channels.
+        """
+        super(UpSampling, self).__init__()
+        self.up = nn.ConvTranspose2d(in_channels, out_channels=out_channels, kernel_size=2, stride=2)
+        self.conv1 = nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=3, stride=1, padding=1)
+        self.relu = nn.ReLU()
+        self.conv2 = nn.Conv2d(in_channels=out_channels, out_channels=out_channels, kernel_size=3, stride=1, padding=1)
+        self.batchnorm = nn.BatchNorm2d(out_channels)
+
+    def forward(self, x, prev_skip):
+        x = self.up(x)
+        x = torch.cat((x, prev_skip), dim=1)
+        x = self.conv1(x)
+        x = self.conv2(x)
+        next_layer = self.relu(self.batchnorm(x))
+        return next_layer
+
+class UNet(nn.Module):
+    
+    """
+        U-Net architecture.
+
+        Args:
+            in_channels (int): Number of input channels.
+            out_channels (int): Number of output channels.
+            features (list): List of feature sizes for downsampling and upsampling.
+    """
+    def __init__(self, in_channels, out_channels, features):
+        super(UNet, self).__init__()
+        self.ups = nn.ModuleList()
+        self.downs = nn.ModuleList()
+
+        for feature in features:
+            self.downs.append(DownSampling(in_channels, feature, True))
+            in_channels = feature
+
+        for feature in reversed(features):
+            self.ups.append(UpSampling(2 * feature, feature))
+
+        self.bottleneck = DownSampling(features[-1], 2 * features[-1], False)
+        self.final_conv = nn.Conv2d(features[0], out_channels, kernel_size=1)
+
+    def forward(self, x):
+        skip_connections = []
+        for down in self.downs:
+            x, skip_connection = down(x)
+            skip_connections.append(skip_connection)
+        skip_connections = skip_connections[::-1]
+        x = self.bottleneck(x)
+        for i, up in enumerate(self.ups):
+            x = up(x, skip_connections[i])
+
+        return self.final_conv(x)
+    
+if __name__ == "__main__":
+    #Example Usage
+    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    features = [64, 128, 256, 512]
+    model = UNet(1, 1, features=features).to(device)
+    print(model(torch.rand(1, 1, 512, 512)).shape)