Convolutional Neural Networks (CNNs) have emerged as a powerful tool for tasks such as image classification, object detection, and segmentation. However, the performance of these models heavily depends on the quantity and diversity of the training data available. In scenarios where the dataset is limited, data augmentation becomes a crucial technique to enhance model generalization and improve robustness.
Data augmentation involves applying various transformations to the existing dataset, generating new training samples with different appearances while preserving the underlying class labels. This process not only enriches the dataset but also helps the model learn invariant features and improves its ability to generalize to unseen data.
Two popular deep learning frameworks, TensorFlow and PyTorch, offer comprehensive support for integrating data augmentation into CNN models.
TensorFlow:
In TensorFlow, the tf.image module provides a rich set of functions for image processing and augmentation. This includes operations such as rotation, flipping, zooming, and brightness adjustments. By incorporating these augmentations directly into the data pipeline using tf.data and tf.image, developers can seamlessly integrate data augmentation into the training process.
import tensorflow as tf
# Example: Applying data augmentation in TensorFlow
data_augmentation = tf.keras.Sequential([
tf.keras.layers.experimental.preprocessing.Rescaling(1./255),
tf.keras.layers.experimental.preprocessing.RandomFlip("horizontal"),
tf.keras.layers.experimental.preprocessing.RandomRotation(0.2),
tf.keras.layers.experimental.preprocessing.RandomZoom(0.2),
])
# Use this data augmentation layer in your model
model = tf.keras.Sequential([
data_augmentation,
# Add your CNN layers here
])
PyTorch:
PyTorch, on the other hand, leverages the torchvision.transforms module to achieve data augmentation. This module provides a variety of transformations that can be applied to images during the training phase. Developers can easily compose these transformations and integrate them into the data loading process.
import torchvision.transforms as transforms
# Example: Applying data augmentation in PyTorch
transform = transforms.Compose([
transforms.RandomHorizontalFlip(),
transforms.RandomRotation(20),
transforms.RandomResizedCrop(224),
transforms.ToTensor(),
])
# Use this transform in your dataset loader
train_dataset = torchvision.datasets.ImageFolder(root='path/to/training/data', transform=transform)
TensorFlow:
In this section, all the functions needed to be merged in an ML project are presented. It is important to mention that this Python file is focused on increasing the dataset before training the ML model. Results are in: https://github.com/puldavid87/Greenhouse_project/tree/main/augmented_images
import os
import random
import matplotlib.image as mpimg
import matplotlib.pyplot as plt
import tensorflow as tf
from PIL import Image
from tensorflow.keras import layers
from tensorflow.keras.preprocessing.image import (ImageDataGenerator,
array_to_img, img_to_array,
load_img)
import numpy as np
train_path = '/your_path/'
output_path = '/your_path/'
def get_labels(train_path):
labels=[]
tam = []
for cont, i in enumerate((os.listdir(train_path))):
labels.append(i)
for cont, i in enumerate (labels):
num = len(os.listdir(train_path+"/"+i))
tam.append(num)
print("Label ", labels[cont], ": ", num)
tam_labels = np.array(tam)
return labels, tam_labels
def rotation_per_image(train_path, labels, output_path):
"""Rotate images in a folder and save to output folders
Args:
train_path: Path to folder containing images
labels: List of labels for output folders
output_path: Path to save rotated images
Returns:
None: Function does not return anything
Rotate each image in train_path by 90, 180 and 270 degrees clockwise and save copies to folders in output_path matching each image's label:
- Read each image from train_path
- Rotate image by 90, 180 and 270 degrees
- Get label of image from labels list
- Save rotated copies to folders in output_path matching the label"""
for i in labels:
os.makedirs(output_path + "/" + i, exist_ok=True)
for i in labels:
file_names = os.listdir(train_path + "/" + i)
img = random.sample(file_names, 1)
image = load_img(train_path + "/" + i + "/" + str(img[0]))
x = img_to_array(image)
x = np.expand_dims(x, axis=0)
cont = 0
for batch in datagen.flow(x, batch_size=1):
augmented_image = array_to_img(batch[0])
save_path = os.path.join(
output_path + "/" + i,
f"augmented_image_{i}_{cont}.jpg")
augmented_image.save(save_path)
cont += 1
# Generate 20 augmented images (you can change this number)
if cont >= 20:
break
plot_data_augmentation(image, output_path + "/" + i)
datagen = ImageDataGenerator(
rotation_range=40, # Rotate images randomly up to 40 degrees
zoom_range=0.1, # Zoom in or out by up to 20%
width_shift_range=0.1,
height_shift_range=0.1,
brightness_range=[0.4, 1.5],
horizontal_flip=True, # Flip the image horizontally
vertical_flip=True,
# fill_mode='nearest' # Fill pixels with the nearest available value
)
def plot_data_augmentation(image, output_path):
"""
Plots original and augmented images side by side.
Args:
image: Original image to augment in one line
output_path: Path to directory containing augmented images in one line
Returns:
None: Does not return anything, just plots images in one line
Processing Logic:
- Display original image and title it "Original Image"
- Load all augmented images from output directory
- Display each augmented image with subplot
"""
# Display the original and augmented images
plt.figure(figsize=(12, 6))
plt.subplot(2, 1, 1)
plt.imshow(image)
plt.title("Original Image")
plt.subplot(2, 1, 2)
augmented_images = os.listdir(output_path)
for l, augmented_image_name in enumerate(augmented_images):
augmented_image_path = os.path.join(output_path, augmented_image_name)
augmented_image = load_img(augmented_image_path)
plt.subplot(4, 5, l + 1)
plt.imshow(augmented_image)
plt.title(f"Augmented Image {l + 1}")
plt.tight_layout()
plt.show()
labels, tam_labels = get_labels(train_path)
rotation_per_image(train_path, labels, output_path)
PyTorch:
In a PyTorch environment, the following Python script replicates the same functionality that the previous section. Results are: https://github.com/puldavid87/Greenhouse_project/tree/main/augmented_images_torch
from PIL import Image
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
import os
import random
from torchvision.io import read_image
import numpy as np
from torchvision.utils import save_image
# Define the path to your image"
train_path = '/your_path"
output_path = '/your_path/'
os.makedirs(output_path, exist_ok=True)
def get_labels(train_path):
labels=[]
tam = []
for cont, i in enumerate((os.listdir(train_path))):
labels.append(i)
for cont, i in enumerate (labels):
num = len(os.listdir(train_path+"/"+i))
tam.append(num)
print("Label ", labels[cont], ": ", num)
tam_labels = np.array(tam)
return labels, tam_labels
def rotation_per_image_torch(train_path, labels, output_path):
for i in labels:
os.makedirs(output_path + "/" + i, exist_ok=True)
for i in labels:
file_names = os.listdir(train_path + "/" + i)
img = random.sample(file_names, 1)
# Load the image and apply data augmentation
image = Image.open(train_path + "/" + i + "/" + img[0])
# Create 20 augmented versions of the image
augmented_images = [transform(image) for _ in range(20)]
cont = 0
for augmented_image in augmented_images:
augmented_image = transforms.ToPILImage()(augmented_image)
save_path = os.path.join(
output_path + "/" + i,
f"augmented_image_{i}_{cont}.jpg")
augmented_image.save(save_path)
cont += 1
return augmented_images
# Define a transform that includes data augmentation
transform = transforms.Compose([
transforms.RandomHorizontalFlip(), # Randomly flip the image horizontally
transforms.RandomVerticalFlip(),
# Randomly rotate the image by up to 30 degrees
transforms.RandomRotation(40),
# Adjust brightness, contrast, saturation, and hue
transforms.ColorJitter(brightness=0.2, contrast=0.2),
# Randomly crop and resize the image to 224x224 pixels
transforms.RandomResizedCrop(224),
transforms.ToTensor(), # Convert the image to a PyTorch tensor
])
def plot_data_augmentation(image, augmented_images):
# Display the original and augmented images
plt.figure(figsize=(12, 8))
plt.subplot(4, 5, 1)
plt.title("Original Image")
plt.imshow(image)
# plt.axis('off')
#plt.subplot(2, 1, 2)
for i, augmented_image in enumerate(augmented_images):
plt.subplot(4, 5, i + 2)
plt.title(f"Augmented Image {i + 1}")
# Convert tensor to (H, W, C) for display
plt.imshow(augmented_image.permute(1, 2, 0))
plt.axis('off')
plt.tight_layout()
plt.show()
labels, tam_labels = get_labels(train_path)
rotation_per_image_torch(train_path, labels, output_path)