Scene Segmentation
[Getting Started Notebook] Semantic Segmentation
A Getting Started notebook for Semantic Segmentation Puzzle of BlitzXI.
Shubhamaicrowd
Starter Code for Semantic Segmentarion
What we are going to Learn¶
- Using segmentation_model.pytorch and Unet for multi class semantic segmentation
Note : Create a copy of the notebook and use the copy for submission. Go to File > Save a Copy in Drive to create a new copy
Downloading Dataset¶
Installing aicrowd-cli
In [ ]:
!pip install aicrowd-cli
%load_ext aicrowd.magic
In [ ]:
%aicrowd login
In [ ]:
!rm -rf data
!mkdir data
%aicrowd ds dl -c scene-segmentation -o data
In [ ]:
!unzip data/train.zip -d data/train > /dev/null
!unzip data/test.zip -d data/test > /dev/null
Downloading & Importing Libraries¶
Here we are going to use segmentation_models.pytorch which is a really popular library providing a tons of different segmentation models for pytorch including basic unets to DeepLabV3!
In [ ]:
!pip install git+https://github.com/qubvel/segmentation_models.pytorch
In [ ]:
# Pytorch
import torch
from torch import nn
import segmentation_models_pytorch as smp
from torch.utils.data import Dataset, DataLoader
# Reading Dataset, vis and miscellaneous
from PIL import Image
import matplotlib.pyplot as plt
import os
import numpy as np
import torch.nn as nn
from natsort import natsorted
from tqdm.notebook import tqdm
import cv2
Creating the Dataloader¶
In this section, we will be creating the dataloader that our model will use for the loading batches of corrosponding features and labels for training & testing
In [ ]:
class SemanticSegmentationDataset(Dataset):
def __init__(self, img_directory=None, label_directory=None, train=True):
self.img_directory = img_directory
self.label_directory = label_directory
# If the image direcotry is valid ( using only 500 images for faster training )
if img_directory != None:
if train:
self.img_list = natsorted(os.listdir(img_directory))[:500]
else:
self.img_list = natsorted(os.listdir(img_directory))
if train:
self.label_list = natsorted(os.listdir(label_directory))[:500]
self.train = train
self.labels = list(range(0, 25))
def __len__(self):
return len(self.img_list)
def __getitem__(self, idx):
# Reading the image
img = Image.open(os.path.join(self.img_directory, self.img_list[idx]))
img = img.convert("L")
if self.train == True:
# Readiding the mak image
mask = Image.open(os.path.join(self.label_directory, self.label_list[idx]))
# mask.show()
img = np.array(img, dtype=np.float32)
mask = np.array(mask, dtype=np.float32)
# Change image channel ordering
img = img[np.newaxis, :, :]
# Normalizing images
img = torch.from_numpy(img)
img = img.float()/255
binary_mask = np.array([(mask == v) for v in list(self.labels)])
binary_mask = np.stack(binary_mask, axis=-1).astype('float')
mask_preprocessed = binary_mask.transpose(2, 0, 1)
mask_preprocessed = torch.from_numpy(mask_preprocessed)
return img, mask_preprocessed
# If reading test dataset, only return image
else:
img = np.array(img, dtype=np.float32)
img = img[np.newaxis, :, :]
# img = np.moveaxis(img, -1, 0)
# Normalizing images
img = torch.from_numpy(img)
img = img.float()/255
return img
In [ ]:
# Creating the training dataset
train_dataset = SemanticSegmentationDataset(img_directory="./data/train/image", label_directory="./data/train/segmentation")
train_loader = DataLoader(train_dataset, batch_size=4, num_workers=0, shuffle=False, drop_last=True)
In [ ]:
# Reading the image and corrosponding segmentation
image_batch, segmentation_batch = next(iter(train_loader))
image_batch.shape, segmentation_batch.shape
Out[ ]:
Visualizing Dataset¶
In [ ]:
plt.rcParams["figure.figsize"] = (30,5)
# Going through each image and segmentation
for image, segmentation in zip(image_batch, segmentation_batch):
# Change the channel ordering
image = np.moveaxis(image.numpy()*255, 0, -1)
# Showing the image
plt.figure()
plt.subplot(1,2,1)
plt.imshow(image[:, :, 0])
plt.subplot(1,2,2)
plt.imshow(segmentation[6]*255)
plt.show()
In [ ]:
segmentation.shape
Out[ ]:
Creating the Model¶
Here we will get setting up the model architecture, optimizer and loss.
In [ ]:
ENCODER = 'timm-efficientnet-b0'
ENCODER_WEIGHTS = 'imagenet'
ACTIVATION = "softmax2d"
DEVICE = 'cuda'
# create segmentation model with pretrained encoder
model = smp.Unet(
encoder_name=ENCODER,
encoder_weights=ENCODER_WEIGHTS,
classes=len(train_dataset.labels),
in_channels=1,
activation=ACTIVATION,
)
Hyperparameters & Metrics¶
In [ ]:
# using DiceLoss
loss = smp.utils.losses.DiceLoss()
# using multiple metrics to train the model
metrics = [
smp.utils.metrics.IoU(threshold=0.5),
smp.utils.metrics.Fscore(threshold=0.5),
smp.utils.metrics.Accuracy(threshold=0.5),
smp.utils.metrics.Recall(threshold=0.5),
smp.utils.metrics.Precision(threshold=0.5),
]
# Using Adam optimizer
optimizer = torch.optim.Adam(params=model.parameters(), lr=0.006)
Training the Model¶
In [ ]:
# Setting up training epoch to train the model
train_epoch = smp.utils.train.TrainEpoch(
model,
loss=loss,
metrics=metrics,
optimizer=optimizer,
device=DEVICE,
verbose=True,
)
In [ ]:
# Training the model
train_logs = train_epoch.run(train_loader)
Generating Predictions¶
In [ ]:
# Creating the testing dataset
test_dataset = SemanticSegmentationDataset(img_directory="data/test/image", train=False)
test_loader = DataLoader(test_dataset, batch_size=1, num_workers=2, shuffle=False, drop_last=False)
In [ ]:
# Generating Model Predictions
!rm -rf segmentation
!mkdir segmentation
for n, batch in enumerate(tqdm(test_loader)):
# Getting the predictions
predictions = model.predict(batch.to(DEVICE)).cpu()
# Converting the predictions to right format
prediction_mask = (predictions.squeeze().cpu().numpy())
prediction_mask = np.transpose(prediction_mask, (1, 2, 0))
# Getting individual channel and combining them into single image
prediction_mask_gray = np.zeros((prediction_mask.shape[0],prediction_mask.shape[1]))
for ii in range(prediction_mask.shape[2]):
prediction_mask_gray = prediction_mask_gray + ii*prediction_mask[:,:,ii].round()
# Saving the image
prediction_mask_gray = Image.fromarray(prediction_mask_gray.astype(np.uint8))
prediction_mask_gray.save(os.path.join("segmentation", f"{n}.png"))
Submitting our Predictions¶
Note : Please save the notebook before submitting it (Ctrl + S)
In [ ]:
!aicrowd notebook submit -c scene-segmentation -a segmentation --no-verify
In [ ]:
Content
Comments
You must login before you can post a comment.