Upload 30 files

Dockerfile

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

app.py

@@ -0,0 +1,66 @@
+from src.model.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
+import config.configure as config
+from src.pipelines.predict import predict_mask
+
+model = UNet(3, 1, [64, 128, 256, 512])
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+
+model.load_state_dict(torch.load(config.SAVE_MODEL_PATH, 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
+    mask = predict_mask(data=input_tensor, device=device, model=model, inference=True)
+    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()

config/.kaggle/kaggle.json

@@ -0,0 +1 @@
+{"username":"YOUR_KAGGLE_USERNAME","key":"YOUR_API"}

config/configure.py

@@ -0,0 +1,5 @@
+import os
+mask_images_path = os.getcwd() + "/data/lgg-mri-segmentation/kaggle_3m/*/*_mask.tif"
+SAVE_MODEL_PATH = os.getcwd() + "/src/model/model/best_model.pth" 
+SAVE_DATA_PATH = os.getcwd() + "/data"
+DATASET_NAME = "mateuszbuda/lgg-mri-segmentation"

main.py

@@ -0,0 +1,62 @@
+import os
+import torch.nn as nn
+import torch
+import matplotlib.pyplot as plt
+import torchsummary
+import torchview
+import config.configure as config
+from src import logger
+from src.data.data_ingestion import DataIngestion
+from src.data.data_preprocess import data_loaders
+from src.pipelines.training import model_fit
+from src.model.unet import UNet
+## graphviiz
+STAGE_NAME = "Data Ingestion stage"
+try:
+   logger.info(f">>>>>>>>  Starting {STAGE_NAME}  <<<<<<<<")
+   data_ingestion = DataIngestion()
+   data_ingestion.download()
+except Exception as e:
+        logger.exception(e)
+        raise e
+
+STAGE_NAME = 'Training'
+BATCH_SIZE = 32
+NUM_WORKERS = 3
+EPOCHS = 50
+PATH = config.SAVE_MODEL_PATH
+
+try:
+    logger.info(f'Preparing DataLoders')
+
+    # getting the dataloaders
+    train_loader, valid_loader = data_loaders(batch_size=BATCH_SIZE, num_workers=NUM_WORKERS, train_split=True)
+
+    # fitting the model
+    loss_fn = nn.BCEWithLogitsLoss()
+    in_channels = 3
+    out_channels = 1
+    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    features = [64, 128, 256, 512]
+    model = UNet(in_channels=in_channels, out_channels=out_channels, features=features)
+    optimizer = torch.optim.AdamW(model.parameters(),lr=1e-4)
+    
+
+    # starting the training stage
+    logger.info(f"Strating {STAGE_NAME} Stage \n\n ==============")
+
+    summary = model_fit(
+            epochs=EPOCHS,
+            model=model,
+            device=device,
+            train_loader=train_loader,
+            valid_loader=valid_loader,
+            criterion=loss_fn,
+            optimizer=optimizer,
+            PATH=PATH
+        )
+    
+    
+except Exception as e:
+    logger.exception(e)
+    raise e

requirements.txt

@@ -0,0 +1,18 @@
+# Core Libraries
+numpy>=1.21.0
+pandas>=1.3.3
+kaggle>=1.6.3
+# Data Visualization
+matplotlib>=3.4.3
+seaborn>=0.13.0
+pillow>=9.3.0
+
+tqdm>=4.66.0
+
+# Web Application Framework
+streamlit>=1.28.0
+
+scikit-learn>=0.24.2
+torch==2.1.2
+torchvision>=0.16.2
+albumentations>=1.3.1

src/__init__.py

@@ -0,0 +1,20 @@
+import logging
+import os
+# Configure the logging settings
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+
+# Create a logger with the name of the current module or script
+logger = logging.getLogger(__name__)
+
+# Create a FileHandler and set the log file path
+if not os.path.exists("logs/"):
+    os.makedirs('logs/')
+log_file_path = 'logs/my_log_file.log'
+file_handler = logging.FileHandler(log_file_path)
+
+# Create a formatter and set it for the FileHandler
+formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
+file_handler.setFormatter(formatter)
+
+# Add the FileHandler to the logger
+logger.addHandler(file_handler)

src/data/data_ingestion.py

@@ -0,0 +1,22 @@
+import os
+from src import logger
+import config.configure as config
+print(os.getcwd())
+os.environ["KAGGLE_CONFIG_DIR"] = "config/.kaggle/"
+import kaggle
+class DataIngestion():
+    def __init__(self):
+        pass
+    def download(self):
+        logger.info(f"Downloading dataset into data/")
+        dataset_name = config.DATASET_NAME
+
+        # Replace '/path/to/your/folder' with the path to the folder where you want to save the dataset
+        output_folder = config.SAVE_DATA_PATH
+        
+        # Download the dataset to the specified folder
+        kaggle.api.dataset_download_files(dataset_name, path=output_folder, unzip=True)
+        logger.info(f"Download Completed")
+
+if __name__ == '__main__':
+    data = DataIngestion()

src/data/data_preprocess.py

@@ -0,0 +1,113 @@
+import os
+import glob
+import pandas as pd
+import numpy as np
+import torch
+from torch import nn
+from torch.utils.data import Dataset, DataLoader
+import albumentations as A
+from albumentations.pytorch import ToTensorV2
+from PIL import Image
+from sklearn.model_selection import train_test_split
+from config.configure import mask_images_path
+from src import logger
+def get_dataframe(path: str) -> pd.DataFrame:
+    """
+    Create a DataFrame containing image paths, mask paths, and labels.
+
+    Args:
+        path (str): path [mask_images]
+
+    Returns:
+        pd.DataFrame: DataFrame with image paths, mask paths, and labels.
+    """
+
+    image_masks = glob.glob(path)
+    image_paths = [file_path.replace("_mask", '') for file_path in image_masks]
+
+    def labels(mask_path):
+        label = []
+        for mask in mask_path:
+            img = Image.open(mask)
+            label.append(1) if np.array(img).sum() > 0 else label.append(0)
+        return label
+
+    mask_labels = labels(image_masks)
+
+    df = pd.DataFrame({
+        'image_path': image_paths,
+        'mask_path': image_masks,
+        'label': mask_labels
+    })
+
+    return df
+
+class MRIDataset(Dataset):
+    def __init__(self, paths, transform):
+        """
+        Custom dataset for MRI images.
+
+        Args:
+            paths (pd.DataFrame): DataFrame containing mask paths.
+            transform: Data augmentation and transformation pipeline.
+        """
+        self.paths = paths
+        self.transform = transform
+
+    def __len__(self):
+        return len(self.paths)
+
+    def __getitem__(self, idx):
+        image_path, mask_path = self.paths.iloc[idx]
+        image = Image.open(image_path)
+        mask = Image.open(mask_path)
+
+        image = np.array(image).astype(np.float32) / 255.
+        mask = np.array(mask).astype(np.float32) / 255.
+
+        if self.transform:
+            transformed = self.transform(image=image, mask=mask)
+            return transformed['image'], transformed['mask'].unsqueeze(0)
+        else:
+            transformed = ToTensorV2()(image=image, mask=mask)
+            return transformed['image'], transformed['mask'].unsqueeze(0)
+
+
+def data_loaders(batch_size,num_workers, train_split=False) -> DataLoader:
+    
+    logger.info(f"Preprocessing Data")
+    df = get_dataframe(mask_images_path)
+
+    train_transforms = A.Compose([
+    A.Resize(224, 224, p=1.0),
+    A.RandomBrightnessContrast(p=0.2),
+    A.HorizontalFlip(p=0.5),
+    A.VerticalFlip(p=0.5),
+    ToTensorV2(),
+    ])
+
+    # Only reshape val and test data
+    val_transforms = A.Compose([
+        A.Resize(224, 224, p=1.0),
+        ToTensorV2(),
+    ])
+
+    # splitting the dataset
+    train_x, val_x, train_y, val_y = train_test_split(df.drop('label',axis=1), df.label,test_size=0.3)
+    val_x , test_x, val_y, test_y = train_test_split(val_x, val_y, test_size = 0.2)
+
+    train_data = MRIDataset(train_x, train_transforms)
+    val_data = MRIDataset(val_x, val_transforms)
+    test_data = MRIDataset(test_x[test_y == 1], val_transforms)
+
+
+    # train_loader = DataLoader(train_data, batch_size=32, shuffle=True)
+
+    if train_split:
+        train_loader = DataLoader(train_data, batch_size=batch_size, shuffle=True, num_workers=num_workers)
+        val_loader = DataLoader(val_data, batch_size=batch_size, shuffle=True, num_workers=num_workers)
+
+        return train_loader, val_loader
+    else:
+        test_loader = DataLoader(test_data, batch_size=32, shuffle=True)
+        return test_loader 

src/model/model/best_model.pth

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

src/model/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)

src/pipelines/predict.py

@@ -0,0 +1,70 @@
+import torch
+import torch.nn as nn
+from tqdm import tqdm
+from typing import *
+from src import logger
+
+def predict_mask(
+        data: Any,
+        device: Any,
+        model: nn.Module,
+        inference: bool,
+        valid_loader=None,
+        criterion=None,
+):
+    """
+    predicts mask for the image
+    Args:
+        data (Any): image data for predicting
+        model (nn.Module): model for training
+        device (0/'cud'/'cpu'/Any): name of device
+        inference (bool): Whether to evaluate or predict
+        valid_Loader (nn.Module): test loader for training
+        criterion (nn.Module): loss criteria
+
+    Example:
+    >>>     train(
+    >>>     data = torch.FloatTensor,
+    >>>     model=model,
+    >>>     device=0/'cuda'/'cpu'
+    >>>     ingerence=0
+    >>>     valid_loader= test_loader
+    >>>     criterion= fn_loss
+    """
+
+    if inference:
+
+        with torch.no_grad():
+            image = data.type(torch.FloatTensor).to(device)
+            model = model.to(device)
+            pred = model(image)
+            pred = torch.sigmoid(pred)
+            mask = (pred > 0.6).float()
+            
+            return mask.cpu().detach()
+    else:
+        with torch.no_grad():
+            val_Loss = 0
+            val_Dicescore = 0
+            model.eval()   
+            for x, y in tqdm(valid_loader):
+                x = x.type(torch.cuda.FloatTensor).to(device)
+                y = y.type(torch.cuda.FloatTensor).to(device)
+
+                predict = model(x)
+                loss = criterion(predict, y)
+                val_Loss += loss.item()
+                
+                predict = torch.sigmoid(predict)
+                predict = (predict > 0.5).float() 
+
+                dice_score = (2 * (y*predict).sum() + 1e-8)/((y+predict).sum() + 1e-8)
+                try:
+                    val_Dicescore += dice_score.cpu().item()
+                except:
+                    val_Dicescore += dice_score
+
+            val_Loss /= len(valid_loader)
+            val_Dicescore /= len(valid_loader)
+
+        logger.info(f"Test Loss: {val_Loss}  - Dice Score: {val_Dicescore}")

src/pipelines/training.py

@@ -0,0 +1,120 @@
+import torch
+from tqdm import tqdm
+import torch.nn as nn
+from src import logger
+from typing import *
+import warnings
+warnings.filterwarnings('ignore')
+def model_fit(
+    epochs: int,
+    model: nn.Module,
+    device: Any,
+    train_loader: Any,
+    valid_loader: Any,
+    criterion: nn.Module,
+    optimizer: nn.Module,
+    PATH: str
+):
+    """
+
+     Args:
+         epochs (int): # of epochs
+         model (nn.Module): model for training
+         device (Union[int, str]): number or name of device
+         train_loader (Any): pytorch loader for trainset
+         valid_loader (Any): pytorch loader for testset
+         criterion (nn.Module): loss critiria
+         optimizer (nn.Module): optimizer for model training
+         path (str): path for saving model
+
+
+    Example:
+    >>>     train(
+    >>>     epochs=25,
+    >>>     model=model,
+    >>>     device=0/'cuda'/'cpu',
+    >>>     train_loader=train_loader,
+    >>>     valid_loader=valid_loader,
+    >>>     criterion=fn_loss,
+    >>>     optimizer=optimizer)
+    """
+
+
+    best_DICESCORE = 0
+    model.to(device)
+    summary = {
+        'train_loss' : [],
+        'train_dice' : [],
+        'valid_loss' : [],
+        'valid_dice' : []
+    }
+    for epoch in range(epochs):
+        logger.info(f"EPOCH {epoch}/{epochs}")
+        train_Loss = 0
+        train_Dicescore = 0
+        model.train()
+        for x, y in tqdm(train_loader):
+            x = x.type(torch.FloatTensor).to(device)
+            y = y.type(torch.FloatTensor).to(device)
+            
+            predict = model(x)
+            loss =  criterion(predict, y)
+            train_Loss += loss.item()
+            
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+            
+            predict = torch.sigmoid(predict)
+            predict = (predict > 0.5).float() 
+            
+            dice_score = (2 * (y*predict).sum() + 1e-8)/((y+predict).sum() + 1e-8)
+
+            try:
+                train_Dicescore += dice_score.cpu().item()
+            except:
+                train_Dicescore += dice_score
+            
+        train_Loss /= len(train_loader)
+        train_Dicescore /= len(train_loader)
+    
+
+
+        with torch.no_grad():
+            val_Loss = 0
+            val_Dicescore = 0
+            model.eval()   
+            for x, y in tqdm(valid_loader):
+                x = x.type(torch.FloatTensor).to(device)
+                y = y.type(torch.FloatTensor).to(device)
+
+                predict = model(x)
+                loss = criterion(predict, y)
+                val_Loss += loss.item()
+                
+                predict = torch.sigmoid(predict)
+                predict = (predict > 0.5).float() 
+
+                dice_score = (2 * (y*predict).sum() + 1e-8)/((y+predict).sum() + 1e-8)
+                try:
+                    val_Dicescore += dice_score.cpu().item()
+                except:
+                    val_Dicescore += dice_score
+
+            val_Loss /= len(valid_loader)
+            val_Dicescore /= len(valid_loader)
+
+
+        logger.info(f"Loss: {train_Loss}  - Dice Score: {train_Dicescore} - Validation Loss: {val_Loss} - Validation Dice Score: {val_Dicescore}")
+
+        if val_Dicescore > best_DICESCORE:
+            best_DICESCORE  = val_Dicescore
+            torch.save(model, PATH)
+
+        summary['train_loss'] = train_Loss
+        summary['train_dice'] = train_Dicescore
+        summary['valid_loss'] = val_Loss 
+        summary['valid_dice'] = val_Dicescore
+
+
+    return summary