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) | |