The project focuses on implementing a U-Net-based deep learning model for multi-class image segmentation, a fundamental yet challenging problem in the field of computer vision. The primary objective of this project is to efficiently segment input images into distinct, well-defined classes by training the U-Net model on a carefully curated and annotated dataset. This technique holds significant importance in various applications such as medical imaging for disease detection, autonomous vehicles for environment perception, and numerous other domains requiring precise image analysis.
This report provides a comprehensive overview of the code implementation, step-by-step methodology, evaluation metrics, and results achieved during the project. Additionally, it delves into the architectural details of the U-Net model, the strategies employed to optimize performance, and potential areas for improvement in future iterations of the work.
To develop a U-Net model capable of performing multi-class image segmentation, with the aim of precisely identifying and distinguishing between various image classes. The performance of the model is evaluated through robust metrics like Dice Loss, ensuring a thorough assessment of its segmentation accuracy and efficiency.
The following tools and libraries were used:
The dataset was preprocessed and loaded using PyTorch's DataLoader class. The preprocessing steps included:
train_data_loader = DataLoader(train_dataset, batch_size=16, shuffle=True) test_data_loader = DataLoader(test_dataset, batch_size=16, shuffle=False)
The U-Net architecture was implemented with the following design:
class UNet(nn.Module): def __init__(self, num_classes): super(UNet, self).__init__() self.num_classes = num_classes self.down_conv_11 = conv_block(in_channels=3, out_channels=64) self.down_conv_12 = nn.MaxPool2d(kernel_size=2, stride=2) self.down_conv_21 = conv_block(in_channels=64, out_channels=128) self.down_conv_22 = nn.MaxPool2d(kernel_size=2, stride=2) self.down_conv_31 = conv_block(in_channels=128, out_channels=256) self.down_conv_32 = nn.MaxPool2d(kernel_size=2, stride=2) self.down_conv_41 = conv_block(in_channels=256, out_channels=512) self.down_conv_42 = nn.MaxPool2d(kernel_size=2, stride=2) self.middle = conv_block(in_channels=512, out_channels=1024) def forward(self, x): x1 = self.down_conv_11(X) # [-1, 64, 256, 256] x2 = self.down_conv_12(x1) # [-1, 64, 128, 128] x3 = self.down_conv_21(x2) # [-1, 128, 128, 128] x4 = self.down_conv_22(x3) # [-1, 128, 64, 64] x5 = self.down_conv_31(x4) # [-1, 256, 64, 64] x6 = self.down_conv_32(x5) # [-1, 256, 32, 32] x7 = self.down_conv_41(x6) # [-1, 512, 32, 32] x8 = self.down_conv_42(x7) # [-1, 512, 16, 16] middle_out = self.middle(x8) # [-1, 1024, 16, 16] self.up_conv_11 = nn.ConvTranspose2d(in_channels=1024, out_channels=512, kernel_size=3, stride=2, padding=1, output_padding=1) self.up_conv_12 = conv_block(in_channels=1024, out_channels=512) self.up_conv_21 = nn.ConvTranspose2d(in_channels=512, out_channels=256, kernel_size=3, stride=2, padding=1, output_padding=1) self.up_conv_22 = conv_block(in_channels=512, out_channels=256) self.up_conv_31 = nn.ConvTranspose2d(in_channels=256, out_channels=128, kernel_size=3, stride=2, padding=1, output_padding=1) self.up_conv_32 = conv_block(in_channels=256, out_channels=128) self.up_conv_41 = nn.ConvTranspose2d(in_channels=128, out_channels=64, kernel_size=3, stride=2, padding=1, output_padding=1) self.up_conv_42 = conv_block(in_channels=128, out_channels=64) self.output = nn.Conv2d(in_channels=64, out_channels=num_classes, kernel_size=3, stride=1, padding=1) return output
The model was trained using the Adam optimizer and Dice Loss as the evaluation metric, emphasizing the significance of utilizing a cutting-edge loss function like Dice Loss in multi-class segmentation tasks. This state-of-the-art approach ensures enhanced evaluation of overlap and segmentation accuracy, making it particularly effective for complex segmentation challenges. The training loop included:
learning_rate = 0.001 epochs = 5 device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') device Umodel = UNet(num_classes=13).to(device) sample = (next(iter(train_data_loader))) sample[1].shape out = Umodel(sample[0].to(device)) out.shape optimizer = torch.optim.Adam(Umodel.parameters()) total_steps = len(train_data_loader) print(f"{epochs} epochs, {total_steps} total_steps per epoch") criterion = DiceLoss() import time epoch_losses = [] for epoch in range(epochs): start_time = time.time() epoch_loss = [] for batch_idx, (data, labels) in enumerate(train_data_loader): data, labels = data.to(device), labels.to(device) optimizer.zero_grad() outputs = Umodel(data) #loss = nn.CrossEntropyLoss(outputs,labels)# - torch.log(DiceLoss(outputs, labels)) loss = criterion(outputs, labels) loss.backward() optimizer.step() epoch_loss.append(loss.item()) if batch_idx % 200 == 0: print(f'batch index : {batch_idx} | loss : {loss.item()}') print(f'Epoch {epoch+1}, loss: ', np.mean(epoch_loss)) end_time = time.time() print(f'Spend time for 1 epoch: {end_time - start_time} sec') epoch_losses.append(epoch_loss)
Dice Loss was implemented to evaluate the overlap between predicted and ground truth segmentation maps, significantly enhancing the precision and reliability of the segmentation task. Utilizing Dice Loss as a state-of-the-art loss function underscores its effectiveness in handling complex multi-class segmentation challenges with improved accuracy.
class DiceLoss(nn.Module): def forward(self, logits, targets): smooth = 1 num = targets.size(0) probs = logits m1 = probs.reshape(num, -1) m2 = targets.reshape(num, -1) intersection = (m1 * m2) score = (2. * intersection.sum(1) + smooth) / (m1.sum(1) + m2.sum(1) + smooth) score = 1 - (score.sum() / num) return score
The trained model was saved in two formats: PyTorch checkpoint and TorchScript for inference.
save_model_path = './Unet_Model_dice_loss.pth' torch.save(model.state_dict(), save_model_path) model_scripted = torch.jit.script(model) model_scripted.save('modelnew.pt')
The model's performance was evaluated on unseen test data. Predictions were compared against ground truth labels, and visualization was done to assess segmentation accuracy.
model_scripted = torch.jit.load('modelnew.pt') model_scripted.eval() for data, labels in test_data_loader: with torch.no_grad(): outputs = model_scripted(data) # Visualization and metrics calculation
The training loss showed a consistent decrease, indicating successful learning.
Evaluation was conducted using Dice Score:
Sample predictions and ground truth were visualized to verify segmentation accuracy. The model successfully segmented images into their respective classes.
The project successfully implemented a U-Net model for multi-class image segmentation. With further refinements, this model can be adapted for diverse applications in medical imaging, autonomous vehicles, and beyond.
