def test_generator(batch_size):
directory = 'custom_dataset/'
test_image_files = load_test_images()
dataset_size = len(test_image_files)
indices = batch_indices(batch_size=batch_size, dataset_size=dataset_size)
print('Test Dataset Size: ', dataset_size)
while True:
# for i in range(0, 2):
# random.shuffle(valid_image_files)
for index in indices:
x_batch = []
y_batch = []
for n in range(index[0], index[1]):
image_name = test_image_files[n]
image, bbox = label_generator(directory, 'test', image_name)
yolo_out = bbox_to_grid(bbox)
x_batch.append(image)
y_batch.append(yolo_out)
# visualize(image, yolo_out, RGB2BGR=True)
x_batch = np.asarray(x_batch) / 255.0
y_batch = np.asarray(y_batch)
yield x_batch, y_batch
def valid_generator(batch_size):
directory = '../../../EgoGesture Dataset/'
valid_image_files = load_valid_images()
dataset_size = len(valid_image_files)
indices = batch_indices(batch_size=batch_size, dataset_size=dataset_size)
print('Validation Dataset Size: ', dataset_size)
while True:
for i in range(0, 10):
random.shuffle(valid_image_files)
for index in indices:
x_batch = []
y_batch = []
for n in range(index[0], index[1]):
image_name = valid_image_files[n]
image, bbox = label_generator(directory,
image_name,
type='Valid')
yolo_out = bbox_to_grid(bbox)
x_batch.append(image)
y_batch.append(yolo_out)
# visualize(image, yolo_out, RGB2BGR=True)
x_batch = np.asarray(x_batch) / 255.0
y_batch = np.asarray(y_batch)
yield x_batch, y_batch
def train_generator(batch_size, is_augment=True):
if is_augment:
batch_size = int(batch_size / 4)
directory = '../../../EgoGesture Dataset/'
train_image_files = load_train_images()
dataset_size = len(train_image_files)
indices = batch_indices(batch_size=batch_size, dataset_size=dataset_size)
print('Training Dataset Size: ', dataset_size)
while True:
for i in range(0, 10):
random.shuffle(train_image_files)
for index in indices:
x_batch = []
y_batch = []
for n in range(index[0], index[1]):
image_name = train_image_files[n]
image, bbox = label_generator(directory, image_name, type='')
yolo_out = bbox_to_grid(bbox)
x_batch.append(image)
y_batch.append(yolo_out)
# visualize(image, yolo_out, RGB2BGR=True)
# augment
image_aug, bbox_aug = augment(image, bbox)
yolo_out = bbox_to_grid(bbox_aug)
x_batch.append(image_aug)
y_batch.append(yolo_out)
# visualize(image_aug, yolo_out, RGB2BGR=True)
# horizontal flip
image_flip, bbox_flip = flip(image, bbox)
yolo_out = bbox_to_grid(bbox_flip)
x_batch.append(image_flip)
y_batch.append(yolo_out)
# visualize(image_flip, yolo_out, RGB2BGR=True)
# horizontal flip + augment
image_flip_aug, bbox_flip_aug = augment(image_flip, bbox_flip)
yolo_out = bbox_to_grid(bbox_flip_aug)
x_batch.append(image_flip_aug)
y_batch.append(yolo_out)
# visualize(image_flip_aug, yolo_out, RGB2BGR=True)
x_batch = np.asarray(x_batch) / 255.0
y_batch = np.asarray(y_batch)
yield x_batch, y_batch
aug = aug.to_deterministic()
image_aug = aug.augment_image(image)
image_aug = image_aug.copy()
bbs_aug = aug.augment_bounding_boxes([bbs])[0]
b = bbs_aug.bounding_boxes
bbs_aug = [b[0].x1, b[0].y1, b[0].x2, b[0].y2]
bbs_aug = np.asarray(bbs_aug)
bbs_aug[0] = bbs_aug[0] if bbs_aug[0] > 0 else 0
bbs_aug[1] = bbs_aug[1] if bbs_aug[1] > 0 else 0
bbs_aug[2] = bbs_aug[2] if bbs_aug[2] < size else size
bbs_aug[3] = bbs_aug[3] if bbs_aug[3] < size else size
return image_aug, bbs_aug
if __name__ == '__main__':
dir = '../../../../EgoGesture Dataset/'
img_name = 'BasketballField_Single_Eight_color_91.jpg'
img, box = label_generator(dir, img_name, type='')
# image_aug, bbox_aug = img, box
img_aug, bbox_aug = augment(image=img, bbox=box)
# image_aug, bbox_aug = flip(image=img, bbox=box)
# image_aug, bbox_aug = augment(image=image_aug, bbox=bbox_aug)
bbox_aug = [int(b) for b in bbox_aug]
x1, y1, x2, y2 = bbox_aug[0], bbox_aug[1], bbox_aug[2], bbox_aug[3]
img_aug = cv2.rectangle(img_aug, (x1, y1), (x2, y2), f.box_color, 3)
visualize(image=img_aug, title='visualize', RGB2BGR=True)