Reputation: 907
I want to create my own custom DataGenerator
on my own dataset. I have read all the images and stored the locations and their labels in two variables named images
and labels
. I have written this custom generator:
def data_gen(img_folder, y, batch_size):
c = 0
n_image = list(np.arange(0,len(img_folder),1)) #List of training images
random.shuffle(n_image)
while (True):
img = np.zeros((batch_size, 224, 224, 3)).astype('float') #Create zero arrays to store the batches of training images
label = np.zeros((batch_size)).astype('float') #Create zero arrays to store the batches of label images
for i in range(c, c+batch_size): #initially from 0 to 16, c = 0.
train_img = imread(img_folder[n_image[i]])
# row,col= train_img.shape
train_img = cv2.resize(train_img, (224,224), interpolation = cv2.INTER_LANCZOS4)
train_img = train_img.reshape(224, 224, 3)
# binary_img = binary_img[:,:128//2]
img[i-c] = train_img #add to array - img[0], img[1], and so on.
label[i-c] = y[n_image[i]]
c+=batch_size
if(c+batch_size>=len((img_folder))):
c=0
random.shuffle(n_image)
# print "randomizing again"
yield img, label
What I want to know is how can I add other augmentations like flip
, crop
, rotate
to this generator? Moreover, how should I yield
these augmentations so that they are linked with the correct label.
Please let me know.
Upvotes: 0
Views: 103
Reputation: 246
You can add flip
, crop
, rotate
on train_img
before putting it into the img
. That is,
# ....
While(True):
# ....
# add your data augmentation function here
train_img = data_augmentor(train_img)
img[i-c] = train_img
# ....
Upvotes: 1