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)
