Spaces:
Runtime error
Runtime error
File size: 1,802 Bytes
e22b55b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 |
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)
|