|
- """Perform data augmentation and preprocessing."""
- import argparse
- import json
- import math
- import os
- import random
- import cv2 as cv
- import numpy as np
-
-
- def get_parser():
- """Return argument parser for generating dataset."""
- parser = argparse.ArgumentParser()
- parser.add_argument('--dataset', required=True,
- choices=['trainval', 'test'],
- help="Generate trainval or test dataset.")
- parser.add_argument('--val_prop', type=float, default=0.1,
- help="The proportion of val sample in trainval.")
- parser.add_argument('--label_directory', required=True,
- help="The location of label directory.")
- parser.add_argument('--image_directory', required=True,
- help="The location of image directory.")
- parser.add_argument('--output_directory', required=True,
- help="The location of output directory.")
- return parser
-
-
- def boundary_check(centralied_marks):
- """Check situation that marking point appears too near to border."""
- for mark in centralied_marks:
- if mark[0] < -260 or mark[0] > 260 or mark[1] < -260 or mark[1] > 260:
- return False
- return True
-
-
- def overlap_check(centralied_marks):
- """Check situation that multiple marking points appear in same cell."""
- for i in range(len(centralied_marks) - 1):
- i_x = centralied_marks[i, 0]
- i_y = centralied_marks[i, 1]
- for j in range(i + 1, len(centralied_marks)):
- j_x = centralied_marks[j, 0]
- j_y = centralied_marks[j, 1]
- if abs(j_x - i_x) < 600 / 16 and abs(j_y - i_y) < 600 / 16:
- return False
- return True
-
-
- def generalize_marks(centralied_marks):
- """Convert coordinate to [0, 1] and calculate direction label."""
- generalized_marks = []
- for mark in centralied_marks:
- xval = (mark[0] + 300) / 600
- yval = (mark[1] + 300) / 600
- direction = math.atan2(mark[3] - mark[1], mark[2] - mark[0])
- generalized_marks.append([xval, yval, direction, mark[4]])
- return generalized_marks
-
-
- def write_image_and_label(name, image, centralied_marks, name_list):
- """Write image and label with given name."""
- name_list.append(os.path.basename(name))
- print("Processing NO.%d samples: %s..." % (len(name_list), name_list[-1]))
- image = cv.resize(image, (512, 512))
- cv.imwrite(name + '.jpg', image, [int(cv.IMWRITE_JPEG_QUALITY), 100])
- with open(name + '.json', 'w') as file:
- json.dump(generalize_marks(centralied_marks), file)
-
-
- def rotate_vector(vector, angle_degree):
- """Rotate a vector with given angle in degree."""
- angle_rad = math.pi * angle_degree / 180
- xval = vector[0]*math.cos(angle_rad) + vector[1]*math.sin(angle_rad)
- yval = -vector[0]*math.sin(angle_rad) + vector[1]*math.cos(angle_rad)
- return xval, yval
-
-
- def rotate_centralized_marks(centralied_marks, angle_degree):
- """Rotate centralized marks with given angle in degree."""
- rotated_marks = centralied_marks.copy()
- for i in range(centralied_marks.shape[0]):
- mark = centralied_marks[i]
- rotated_marks[i, 0:2] = rotate_vector(mark[0:2], angle_degree)
- rotated_marks[i, 2:4] = rotate_vector(mark[2:4], angle_degree)
- return rotated_marks
-
-
- def rotate_image(image, angle_degree):
- """Rotate image with given angle in degree."""
- rows, cols, _ = image.shape
- rotation_matrix = cv.getRotationMatrix2D((rows/2, cols/2), angle_degree, 1)
- return cv.warpAffine(image, rotation_matrix, (rows, cols))
-
-
- def generate_dataset(args):
- """Generate dataset according to arguments."""
- if args.dataset == 'trainval':
- val_directory = os.path.join(args.output_directory, 'val')
- args.output_directory = os.path.join(args.output_directory, 'train')
- elif args.dataset == 'test':
- args.output_directory = os.path.join(args.output_directory, 'test')
- os.makedirs(args.output_directory, exist_ok=True)
- name_list = []
- for label_file in os.listdir(args.label_directory):
- name = os.path.splitext(label_file)[0]
- image = cv.imread(os.path.join(args.image_directory, name + '.jpg'))
- with open(os.path.join(args.label_directory, label_file), 'r') as file:
- label = json.load(file)
- centralied_marks = np.array(label['marks'])
- if len(centralied_marks.shape) < 2:
- centralied_marks = np.expand_dims(centralied_marks, axis=0)
- centralied_marks[:, 0:4] -= 300.5
- if boundary_check(centralied_marks) or args.dataset == 'test':
- output_name = os.path.join(args.output_directory, name)
- write_image_and_label(output_name, image,
- centralied_marks, name_list)
-
- if args.dataset == 'test':
- continue
- for angle in range(5, 360, 5):
- rotated_marks = rotate_centralized_marks(centralied_marks, angle)
- if boundary_check(rotated_marks) and overlap_check(rotated_marks):
- rotated_image = rotate_image(image, angle)
- output_name = os.path.join(
- args.output_directory, name + '_' + str(angle))
- write_image_and_label(
- output_name, rotated_image, rotated_marks, name_list)
-
- if args.dataset == 'trainval':
- print("Dividing training set and validation set...")
- val_idx = random.sample(list(range(len(name_list))),
- int(round(len(name_list)*args.val_prop)))
- val_samples = [name_list[idx] for idx in val_idx]
- os.makedirs(val_directory, exist_ok=True)
- for val_sample in val_samples:
- train_directory = args.output_directory
- image_src = os.path.join(train_directory, val_sample + '.jpg')
- label_src = os.path.join(train_directory, val_sample + '.json')
- image_dst = os.path.join(val_directory, val_sample + '.jpg')
- label_dst = os.path.join(val_directory, val_sample + '.json')
- os.rename(image_src, image_dst)
- os.rename(label_src, label_dst)
- print("Done.")
-
-
- if __name__ == '__main__':
- generate_dataset(get_parser().parse_args())
|
