Spaces:
Runtime error
Runtime error
| import numpy as np | |
| import os | |
| import cv2 | |
| import random | |
| import json | |
| import argparse | |
| ap= argparse.ArgumentParser() | |
| ap.add_argument('--input_dir', '-i', required=True, help='Path to input dir for images') | |
| ap.add_argument('--output_dir', '-o', required=True, help='Path to output dir to store files. Must be created') | |
| ap.add_argument('--max_imgs', '-m', default=20000, type=int, help='Max number of images to generate') | |
| args= vars(ap.parse_args()) | |
| def apply_motion_blur(image, size, angle): | |
| k = np.zeros((size, size), dtype=np.float32) | |
| k[ (size-1)// 2 , :] = np.ones(size, dtype=np.float32) | |
| k = cv2.warpAffine(k, cv2.getRotationMatrix2D( (size / 2 -0.5 , size / 2 -0.5 ) , angle, 1.0), (size, size) ) | |
| k = k * ( 1.0 / np.sum(k) ) | |
| return cv2.filter2D(image, -1, k) | |
| folder = args['input_dir'] | |
| folder_save = args['output_dir'] | |
| max_images = args['max_imgs'] | |
| print(max_images) | |
| labels_angle = {} | |
| labels_length= {} | |
| images_done = 0 | |
| for filename in os.listdir(folder): | |
| img = cv2.imread(os.path.join(folder,filename)) | |
| if img is not None and img.shape[1] > img.shape[0]: | |
| img_gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) | |
| img_resized = cv2.resize(img_gray, (640,480), interpolation = cv2.INTER_AREA) | |
| length = random.randint(20,40) | |
| angle = random.randint(0,359) | |
| blurred = apply_motion_blur(img_resized, length, angle) | |
| cv2.imwrite(os.path.join(folder_save,filename), blurred) | |
| if angle>=180: | |
| angle_a= angle - 180 | |
| else: | |
| angle_a= angle | |
| labels_angle[filename] = angle_a | |
| labels_length[filename]= length | |
| images_done += 1 | |
| print("%s done"%images_done) | |
| if(images_done == max_images): | |
| print('Done!!!') | |
| break | |
| with open('angle_labels.json', 'w') as file: | |
| json.dump(labels_angle, file) | |
| with open('length_labels.json', 'w') as file: | |
| json.dump(labels_length, file) | |