34 lines
1.3 KiB
Python
Executable File
34 lines
1.3 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
|
|
import argparse
|
|
import os
|
|
from subprocess import call
|
|
|
|
|
|
if __name__ == '__main__':
|
|
parser = argparse.ArgumentParser(description="Generates dataset")
|
|
parser.add_argument('dataset', help='Dataset directory path', type=str)
|
|
parser.add_argument('-d', help='Decals count', type=int, default=8, dest="decals")
|
|
parser.add_argument('-n', help='Noise level', type=float, default=.4, dest="noise")
|
|
args = parser.parse_args()
|
|
|
|
dataset_path = args.dataset
|
|
ok_path = "ok_" + dataset_path
|
|
damaged_path = "damaged_" + dataset_path
|
|
|
|
if not os.path.exists(dataset_path):
|
|
raise RuntimeError("fuk")
|
|
|
|
if not os.path.exists(damaged_path):
|
|
os.makedirs(damaged_path)
|
|
|
|
if not os.path.exists(ok_path):
|
|
os.makedirs(ok_path)
|
|
|
|
dataset_images_paths = [os.path.join(dataset_path, f) for f in os.listdir(dataset_path)]
|
|
ok_images_paths = [os.path.join(ok_path, f) for f in os.listdir(dataset_path)]
|
|
damaged_images_paths = [os.path.join(damaged_path, f) for f in os.listdir(dataset_path)]
|
|
|
|
for dataset_image_path, ok_image_path, damaged_image_path in zip(dataset_images_paths, ok_images_paths, damaged_images_paths):
|
|
call(['python', './uss-mariusz.py', dataset_image_path, ok_image_path, damaged_image_path, '-d', str(args.decals), '-n', str(args.noise)])
|