#@: Python Imports from __future__ import annotations import os, warnings from typing import Any, NewType, Container, Optional import zipfile #@: data analysis import numpy as np import pandas as pd import matplotlib.pyplot as plt import seaborn as sns from skimage import io, transform from PIL import Image #@: DL import torch import torch.nn as nn import torch.nn.functional as F import torchvision from torchvision import transforms from torchvision.models.detection.faster_rcnn import FastRCNNPredictor from torchvision.models.detection.mask_rcnn import MaskRCNNPredictor #@: typing scripts _path = NewType('_path', Any) _transforms = NewType('_transforms', Any) _img = NewType('_img', Any) _mask = NewType('_mask', Any) _loader = NewType('_loader', Any) _model = NewType('_model', Any) _plot = NewType('_plot', Any) _optimizer = NewType('_optimizer', Any) _criterion = NewType('_criterion', Any) _scheduler = NewType('_scheduler', Any) #@: class Dataset class NuclieDataset(torch.utils.data.Dataset): def __init__(self, root: _path, image_transforms: Optional[Container[_transforms]] = None, mask_transforms: Optional[Container[_transforms]] = None) -> None: self.root = root self.folders: Any = os.listdir(root) if image_transforms: self.image_transforms = image_transforms if mask_transforms: self.mask_transforms = mask_transforms def __len__(self) -> int: return len(self.folders) def __getitem__(self, index: int) -> tuple[_img, _mask]: image_folder: Any = os.path.join(self.root, self.folders[index], 'images/') mask_folder: Any = os.path.join(self.root, self.folders[index], 'masks/') image_path: _path = os.path.join(image_folder, os.listdir(image_folder)[0]) mask_path: _path = os.path.join(mask_folder, os.listdir(mask_folder)[0]) image: _img = Image.open(image_path).convert('RGB') mask: _img = Image.open(mask_path).convert('RGB') if self.image_transforms and self.mask_transforms: image: torch.Tensor = self.image_transforms(image) mask: torch.Tensor = self.mask_transforms(mask) return image, mask #@: Mask-RCNN class MaskRCNN: def __init__(self, num_classes: int) -> None: self.num_classes = num_classes def forward(self) -> _model: model: _model = torchvision.models.detection.maskrcnn_resnet50_fpn(pretrained= True) in_features: int = model.roi_heads.box_predictor.cls_score.in_features model.roi_heads.box_predictor = FastRCNNPredictor(in_features, self.num_classes) in_features_mask: int = model.roi_heads.mask_predictor.conv5_mask.in_channels hidden_layer: int = 256 model.roi_heads.mask_predictor = MaskRCNNPredictor(in_features_mask, hidden_layer, self.num_classes) return model def double_conv(in_channels, out_channels): return nn.Sequential( nn.Conv2d(in_channels, out_channels, 3, padding=1), nn.ReLU(inplace=True), nn.Conv2d(out_channels, out_channels, 3, padding=1), nn.ReLU(inplace=True) ) class Unet(nn.Module): def __init__(self): super().__init__() self.dblock1 = double_conv(3, 64) self.dblock2 = double_conv(64, 128) self.dblock3 = double_conv(128,256) self.dblock4 = double_conv(256,512) self.pool = nn.MaxPool2d(2) self.upsample = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True) self.dblock5 = double_conv(256 + 512, 256) self.dblock6 = double_conv(128 + 256, 128) self.dblock7 = double_conv(128 + 64, 64) self.last_layer = nn.Conv2d(64,1,1) self.sigmoid = nn.Sigmoid() def forward(self,x): conv1 = self.dblock1(x) x = self.pool(conv1) conv2 = self.dblock2(x) x = self.pool(conv2) conv3 = self.dblock3(x) x = self.pool(conv3) conv4 = self.dblock4(x) x = self.upsample(conv4) x = torch.cat([x, conv3], dim=1) x = self.dblock5(x) x = self.upsample(x) x = torch.cat([x, conv2], dim=1) x = self.dblock6(x) x = self.upsample(x) x = torch.cat([x, conv1], dim=1) x = self.dblock7(x) out = self.last_layer(x) # out = self.sigmoid(x) return out #@: Criterion function -> DiceBCELoss() class DiceBCELoss(nn.Module): def __init__(self, weight: Optional[Any] = None, size_average: Optional[bool] = True) -> None: super(DiceBCELoss, self).__init__() def forward(self, inputs: torch.Tensor, targets: torch.Tensor, smooth: Optional[int] = 1) -> torch.Tensor: inputs: torch.Tensor = F.sigmoid(inputs) bce_weight: float = 0.5 inputs: torch.Tensor = inputs.view(-1) targets: torch.Tensor = targets.view(-1) intersection: Any = (inputs * targets).sum() dice_loss: Any = 1 - (2. * intersection + smooth)/ (inputs.sum() + targets.sum() + smooth) BCE: Any = F.binary_cross_entropy(inputs, targets, reduction= 'mean') loss_final: Any = BCE * bce_weight + dice_loss * (1 - bce_weight) return loss_final #@ -> Accuracy functions def iou_(y_pred: torch.Tensor, y: torch.Tensor) -> torch.Tensor|Any: inputs: torch.Tensor = y_pred.reshape(-1) targets: torch.Tensor = y.reshape(-1) intersection: Any = (inputs * targets).sum() total: Any = (inputs + targets).sum() union: Any = total - intersection smooth: int = 1 iou: Any = (intersection + smooth)/ (union + smooth) return iou def iou_batch(y_pred: torch.Tensor, y: torch.Tensor) -> list[torch.Tensor] | Any: ious: list = [] y_pred: torch.Tensor = F.sigmoid(y_pred) y_pred: np.ndarray = y_pred.clone().cpu().detach().numpy() y: np.ndarray = y.clone().cpu().detach().numpy() for pred, label in zip(y_pred, y): ious.append(iou_(pred, label)) iou: Any = np.nanmean(ious) return iou #@: Model Adaptor class class Model(): def __init__(self, network: _model, criterion: _criterion, optimizer: _optimizer, schedular: _schedular, dataloaders: dict[str, _loader], device: torch.device, dataset_size: dict[str, int], num_epochs: Optional[int] = 20) -> None: super(Model, self).__init__() self.net = network self.criterion = criterion self.optimizer = optimizer self.schedular = schedular self.dataloaders = dataloaders self.device = device self.dataset_size = dataset_size self.num_epochs = num_epochs def train(self) -> dict[str, float|int]: self.history: dict[str, float|int] = { x: [] for x in ['train_loss', 'val_loss', 'train_iou', 'val_iou'] } for epoch in range(self.num_epochs): print(f'Epoch {epoch + 1}/ {self.num_epochs}') running_train_loss: list = [] running_train_iou: list = [] for image, mask in self.dataloaders['Training']: image = image.to(self.device, dtype= torch.float) mask = mask.to(self.device, dtype= torch.float) pred_mask = self.net(image) loss = self.criterion(pred_mask, mask) score = iou_batch(pred_mask, mask) self.optimizer.zero_grad() loss.backward() self.optimizer.step() running_train_loss.append(loss.item()) running_train_iou.append(score) else: running_val_loss: list = [] running_val_iou: list = [] with torch.no_grad(): for image, mask in self.dataloaders['Validation']: image = image.to(self.device, dtype= torch.float) mask = mask.to(self.device, dtype= torch.float) pred_mask = self.net(image) loss = self.criterion(pred_mask, mask) score = iou_batch(pred_mask, mask) running_val_loss.append(loss.item()) running_val_iou.append(score) epoch_train_loss, epoch_train_iou = np.mean(running_train_loss), np.mean(running_train_iou) print(f'Train Loss: {epoch_train_loss}, IOU: {epoch_train_iou}') self.history['train_loss'].append(epoch_train_loss) self.history['train_iou'].append(epoch_train_iou) epoch_val_loss, epoch_val_iou = np.mean(running_val_loss), np.mean(running_val_iou) print(f'Val Loss: {epoch_val_loss}, IOU: {epoch_val_iou}') self.history['val_loss'].append(epoch_val_loss) self.history['val_iou'].append(epoch_val_iou) def iou_plot(self) -> _plot: train_iou_list: list[float] = [ float(x.cpu().numpy()) for x in self.history['train_iou'] ] val_iou_list: list[float] = [ float(x.cpu().numpy()) for x in self.history['val_iou'] ] plt.plot(train_iou_list, '-bx') plt.plot(test_iou_list, '-rx') plt.title('Model Accuracy/IOU Plot') plt.xlabel('No. of Epochs') plt.ylabel('Accuracy') plt.legend(['train', 'validation'], loc= 'best') plt.show() #@: Driver code if __name__.__contains__('__main__'): if not os.path.exists('/kaggle/working/data'): os.mkdir('/kaggle/working/data') zip_path: _path = '/kaggle/input/data-science-bowl-2018/stage1_train.zip' with zipfile.ZipFile(zip_path, 'r') as zip_ref: zip_ref.extractall('/kaggle/working/data') data_dir: _path = '/kaggle/working/data/' image_transforms: Container[_transform] = transforms.Compose([ transforms.Resize((128, 128)), transforms.RandomHorizontalFlip(), transforms.RandomVerticalFlip(), transforms.ToTensor(), transforms.Normalize((0.5), (0.5)) ]) mask_transforms: Container[_transforms] = transforms.Compose([ #transforms.Grayscale(num_output_channels=1), transforms.Resize((128, 128)), transforms.RandomHorizontalFlip(), transforms.RandomVerticalFlip(), transforms.ToTensor(), transforms.Normalize((0.5), (0.5)) ]) data: object = NuclieDataset(data_dir, image_transforms, mask_transforms) training_data, val_data = torch.utils.data.random_split(data, [580, 90]) dataset_size: dict[str, int] = { 'Training': len(training_data), 'Validation': len(val_data) } train_loader: _loader = torch.utils.data.DataLoader(training_data, batch_size= 10, shuffle= True) val_loader: _loader = torch.utils.data.DataLoader(val_data, batch_size= 10) dataloaders: dict[str, _loader] = { 'Training': train_loader, 'Validation': val_loader } #print(iter(train_loader).__next__()) device: torch.device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') num_classes: int = 2 #model: _model = MaskRCNN(num_classes).forward() model: _model = Unet() model = model.to(device) #print(model) optimizer: _optimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum= 0.9) scheduler: _scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=7, gamma=0.1) criterion: _criterion = DiceBCELoss() num_epochs: int = 25 Neucli_segmentor: _model = Model(model, criterion, optimizer, scheduler, dataloaders, device, dataset_size, num_epochs) history = Neucli_segmentor.train()