"""Set of functions used to perform data augmentation on images and their assosiated .mat files"""
import tensorflow
from tensorflow.keras.preprocessing.image import ImageDataGenerator
import cv2
import os
import glob
import scipy.io as io
import numpy as np
import random
[docs]def dir_to_list(img_dir,prefix =('IMG','GT_IMG'),img_format ='jpg'):
"""Create lists of paths to the ground truth and images, as to be used for image augmentation, ground truth generation and model training.
For the list to be created correctly, the directory should have the following form :
Project
+--images
+ all images
+--ground_truth
+ all .mat files
:param str img_dir: path to the directory where the images are stored
:param prefix: tuple containing two strings if the name of the images and ground_truth files are different (eg : GT_IMG_3.mat is the ground_truth file corresponding to IMG_3.png, then prefix = ('IMG','GT_IMG'))
:param img_format: the format of the images, can only take the values 'png' and 'jpg'
:returns: 2 lists: the first one contain the paths to the images, the second one contain the paths to the associated ground truths (the Ground truth at index i correspond to the image at index i in the ifrst list)
"""
img_paths = []
GT_paths = []
#img into list
if img_format == 'jpg':
for img_path in glob.glob(os.path.join(img_dir, '*.jpg')):
img_paths.append(img_path)
GT_path = img_path.replace('images','ground_truth').replace('.jpg','.mat').replace(prefix[0],prefix[1])
GT_paths.append(GT_path)
if img_format == 'png':
for img_path in glob.glob(os.path.join(img_dir, '*.png')):
img_paths.append(img_path)
GT_path = img_path.replace('images','ground_truth').replace('.png','.mat').replace(prefix[0],prefix[1])
GT_paths.append(GT_path)
return img_paths, GT_paths
[docs]def Horizontal_Flip(img_paths,GT_paths, img_format = 'jpg'):
"""Flip a list of images horizontally and change the .mat file assosiated with it, save the new images and ground_truth in the specified directory
:param list img_paths: list of paths to the images to be augmented, as generated by the dir_to_list function
:param list GT_paths: list of paths to the Grount truth corresponding to the images to be augmented, as generated by the dir_to_list function
:param img_format: the format of the images, can only take the values 'png' and 'jpg'
"""
assert img_format == 'jpg' or img_format == 'png', 'img_format can only take the values : jpg, png'
for img_path,GT_path in list(zip(img_paths, GT_paths)):
#load
img = cv2.imread(img_path)
mat = io.loadmat(GT_path)
#flip img
new_im = cv2.flip(img,1)
#flip coordinates in .mat
information = mat['image_info']
i=information[0][0]
x=i.item(0)
for j in range(len(x[0])):
x[0][j][0]= img.shape[1]-x[0][j][0]
if img_format == 'jpg':
filename = img_path.replace('.jpg','_reverse.jpg')
if img_format == 'png':
filename = img_path.replace('.png','_reverse.png')
matname = GT_path.replace('.mat','_reverse.mat')
cv2.imwrite(filename, new_im)
io.savemat(matname, mat)
[docs]def brightness_variation(img_paths,GT_paths, img_format = 'jpg'):
"""Change the brightness of list of images, save the new images and ground_truth in the specified directory
:param list img_paths: list of paths to the images to be augmented, as generated by the dir_to_list function
:param list GT_paths: list of paths to the Grount truth corresponding to the images to be augmented, as generated by the dir_to_list function
:param img_format: the format of the images, can only take the values 'png' and 'jpg'
"""
assert img_format == 'jpg' or img_format == 'png', 'img_format can only take the values : jpg, png'
datagen = ImageDataGenerator(brightness_range=[0,1])
dic = {}
for img_path,GT_path in list(zip(img_paths, GT_paths)):
img = cv2.imread(img_path)
mat = io.loadmat(GT_path)
dic['brightness'] = random.uniform(0.3,0.65)
new_im = datagen.apply_transform(img, dic)
#save image
if img_format == 'jpg':
filename = img_path.replace('.jpg','_bright.jpg')
if img_format == 'png':
filename = img_path.replace('.png','_bright.png')
matname = GT_path.replace('.mat','_bright.mat')
cv2.imwrite(filename, new_im)
io.savemat(matname, mat)
[docs]def reverse_and_bright(img_paths,GT_paths,img_format = 'jpg'):
"""Change the brightness of list of images and flip them horizontally, save the new images and ground_truth in the specified directory
:param list img_paths: list of paths to the images to be augmented, as generated by the dir_to_list function
:param list GT_paths: list of paths to the Grount truth corresponding to the images to be augmented, as generated by the dir_to_list function
:param img_format: the format of the images, can only take the values 'png' and 'jpg'
"""
assert img_format == 'jpg' or img_format == 'png', 'img_format can only take the values : jpg, png'
for img_path,GT_path in list(zip(img_paths, GT_paths)):
img = cv2.imread(img_path)
mat = io.loadmat(GT_path)
datagen = ImageDataGenerator(brightness_range=[0,1])
dic = {}
new_image = cv2.flip(img,1)
dic['brightness'] = random.uniform(0.3,0.65)
new_im = datagen.apply_transform(new_image, dic)
#flip coordinates in .mat
information = mat['image_info']
i=information[0][0]
x=i.item(0)
for j in range(len(x[0])):
x[0][j][0]= img.shape[1]-x[0][j][0]
if img_format == 'jpg':
filename = img_path.replace('.jpg','_combine.jpg')
if img_format == 'png':
filename = img_path.replace('.png','_combine.png')
matname = GT_path.replace('.mat','_combine.mat')
cv2.imwrite(filename, new_im)
io.savemat(matname, mat)
[docs]def full_augment(img_paths,GT_paths, img_format = 'jpg'):
"""Perform the three different transformations on a set of images, save the new images and ground_truth in the specified directory
:param list img_paths: list of paths to the images to be augmented, as generated by the dir_to_list function
:param list GT_paths: list of paths to the Grount truth corresponding to the images to be augmented, as generated by the dir_to_list function
:param img_format: the format of the images, can only take the values 'png' and 'jpg'
"""
Horizontal_Flip(img_paths,GT_paths,img_format = img_format)
brightness_variation(img_paths,GT_paths,img_format = img_format)
reverse_and_bright(img_paths,GT_paths,img_format = img_format)