Tree Segmentation
Exploring Augmentations For Beginners
A notebook for beginners to get started with data augmentation
jinoooooooooo
This notebook will walk-through the albumentations library for data augmentation and how to get started for improving the quality of models you can build
Have you ever been at a point where you have an amazing idea that you could implement with your model , so you put on that headphones and start browsing the internet for relevant images... But uhhh there's not much of data that fits your need to start with 😔¶
But but ... the start of art networks all used thousands and thousands of images... What am I gonna do now ? 😢¶
Data augmentations to the rescue 🥏¶
So what exactly is data augmentation and why is it so popular?
Data Augmentation is a very powerful way of increasing your data by NOT copy pasting it but generate images through some special operations to trick the neural network that they are different images .
Refer this link for more examples and explanations - https://albumentations.ai/docs/
In [1]:
!pip install aicrowd-cli
%load_ext aicrowd.magic
In [2]:
%aicrowd login
In [3]:
!rm -rf data
!mkdir data
%aicrowd ds dl -c tree-segmentation -o data
In [4]:
!unzip data/train.zip -d data/train > /dev/null
!unzip data/test.zip -d data/test > /dev/null
In [5]:
import os
os.environ['CUDA_VISIBLE_DEVICES'] = '0'
import numpy as np
import cv2
import matplotlib.pyplot as plt
from tqdm import tqdm
import shutil
import copy
import random
In [6]:
TRAIN_DIR = '/content/data/train'
os.listdir(TRAIN_DIR)
Out[6]:
In [7]:
len(os.listdir('/content/data/train/image')), len(os.listdir('/content/data/train/segmentation'))
Out[7]:
In [8]:
x_train_dir = os.path.join(TRAIN_DIR + '/image')
y_train_dir = os.path.join(TRAIN_DIR + '/segmentation')
THE VISUALIZE FUNCTION 📷¶
This function will help us visualize the before and after effects of using augmentations
In [9]:
def visualize(**images):
n = len(images)
plt.figure(figsize = (10,10))
for i, (name, image) in enumerate(images.items()):
plt.subplot(1, n, i + 1)
plt.title(' '.join(name.split('_')).title())
plt.imshow(image)
plt.show()
LOADING THE DATA 🖨️¶
In [10]:
import torch
from torch import nn
from torch.utils.data import DataLoader
from torch.utils.data import Dataset
import cv2
from natsort import natsorted
from PIL import Image
In [11]:
class TreeSegmentationDataset(Dataset):
def __init__(self, img_directory=None, label_directory=None, train=True, augmentation = None, preprocessing = None):
self.img_directory = img_directory
self.label_directory = label_directory
self.augmentation = augmentation
# If the image direcotry is valid
if img_directory != None:
self.img_list = natsorted(os.listdir(img_directory))
self.label_list = natsorted(os.listdir(label_directory))
self.train = train
def __len__(self):
return len(self.img_list)
def __getitem__(self, idx):
image = cv2.imread(os.path.join(self.img_directory, self.img_list[idx]))
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
if self.train == True:
mask = cv2.imread(os.path.join(self.label_directory, self.label_list[idx]))
if self.augmentation:
sample = self.augmentation(image = image, mask = mask)
image, mask = sample['image'], sample['mask']
return image, mask
else:
return image
This is how our current mask and images look like 😮¶
In [17]:
dataset = TreeSegmentationDataset(x_train_dir, y_train_dir)
image, mask = dataset[10]
visualize(
image = image,
tree_mask = mask,
)
DATA AUGMENTATION¶
In [13]:
! pip install albumentations==0.4.6
In [14]:
import albumentations as albu
This is where we start defining our augmentations 🤩¶
In [15]:
def get_training_augmentation():
train_transform = [
albu.HorizontalFlip(p=0.5), # horizontal flips with 50% probability
albu.ShiftScaleRotate(scale_limit=0.5, rotate_limit=0, shift_limit=0.1, p=1, border_mode=0), #scaling and rotation both
albu.PadIfNeeded(min_height=320, min_width=320, always_apply=True, border_mode=0), #padding the images
albu.RandomCrop(height=320, width=320, always_apply=True), #randomly cropping the images
albu.IAAAdditiveGaussianNoise(p=0.2), # adding some gaussian noise with 20% probability
#the one of block helps us to apply any one of the operations inside the block with a given probability
albu.OneOf(
[
#some color transformations
albu.CLAHE(p=1),
albu.RandomBrightness(p=1),
albu.RandomGamma(p=1),
],
p=0.9,
),
albu.OneOf(
[
# image blurring
albu.IAASharpen(p=1),
albu.Blur(blur_limit=3, p=1),
albu.MotionBlur(blur_limit=3, p=1),
],
p=0.9,
),
albu.OneOf(
[
#changing the hue and saturation values
albu.RandomContrast(p=1),
albu.HueSaturationValue(p=1),
],
p=0.9,
),
]
return albu.Compose(train_transform)
OUR NEW IMAGES 👻¶
In [18]:
augmented_dataset = TreeSegmentationDataset(x_train_dir, y_train_dir, augmentation = get_training_augmentation())
for i in range(5): #visualize any 5 images and their mask
image, mask = augmented_dataset[1]
visualize(image = image, mask = mask)
Insane right? I hope this will help you get started with using augmentations and I'll encourage you to tune those parameters and try out some more cool augmentations and let me know how it went 😋¶
Content
Comments
You must login before you can post a comment.
Nice, but when I use your code and try to create a dataset, my dataset size is (torch.Size([4, 512, 512, 3]), torch.Size([4, 512, 512, 3])), this but it should be (torch.Size([4, 3, 512, 512]), torch.Size([4, 512, 512])), How to fix this!?
Hi sean benhur, glad you liked it . Now to visualize the images , i have not changed the order of dimensions . But to continue building the model in Pytorch, you need to change the order of dimensions like this (channels, height, width) For this you can try this small block of code.
def to_tensor(x, **kwargs): return x.transpose(2, 0, 1).astype(‘float32’)