-
Notifications
You must be signed in to change notification settings - Fork 24
/
Copy pathGOPRO_preprocess.py
90 lines (62 loc) · 3.12 KB
/
GOPRO_preprocess.py
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
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
# coding: utf-8
# In[1]:
import os
import subprocess
import argparse
import numpy as np
parser = argparse.ArgumentParser()
# In[ ]:
def str2bool(v):
return v.lower() in ('true')
parser.add_argument("--GOPRO_path", type = str, default = './GOPRO')
parser.add_argument("--output_path", type = str, default = './output')
parser.add_argument("--train_num", type = int, default = 1000)
parser.add_argument("--test_num", type = int, default = 10)
parser.add_argument("--is_gamma", type = str2bool, default = False)
args = parser.parse_args()
if not os.path.exists(os.path.join(args.output_path, 'train')):
os.mkdir(os.path.join(args.output_path, 'train'))
os.mkdir(os.path.join(args.output_path, 'train/sharp'))
os.mkdir(os.path.join(args.output_path, 'train/blur'))
if not os.path.exists(os.path.join(args.output_path, 'test')):
os.mkdir(os.path.join(args.output_path, 'test'))
os.mkdir(os.path.join(args.output_path, 'test/sharp'))
os.mkdir(os.path.join(args.output_path, 'test/blur'))
GOPRO_train_path = os.path.join(args.GOPRO_path, 'train')
GOPRO_test_path = os.path.join(args.GOPRO_path, 'test')
train_blur = []
train_sharp = []
for direc in sorted(os.listdir(GOPRO_train_path)):
if args.is_gamma:
blur = os.path.join(os.path.join(GOPRO_train_path, direc), 'blur_gamma')
else:
blur = os.path.join(os.path.join(GOPRO_train_path, direc), 'blur')
sharp = os.path.join(os.path.join(GOPRO_train_path, direc), 'sharp')
sharp_imgs = sorted(os.listdir(sharp))
for i, img in enumerate(sorted(os.listdir(blur))):
train_blur.append(os.path.join(blur, img))
train_sharp.append(os.path.join(sharp, sharp_imgs[i]))
train_blur = np.asarray(train_blur)
train_sharp = np.asarray(train_sharp)
random_index = np.random.permutation(len(train_blur))[:args.train_num]
for index in random_index:
subprocess.call(['cp', train_blur[index], os.path.join(args.output_path, 'train/blur/%s'%('_'.join(train_blur[index].split('/')[-3:])))])
subprocess.call(['cp', train_sharp[index], os.path.join(args.output_path, 'train/sharp/%s'%('_'.join(train_sharp[index].split('/')[-3:])))])
test_blur = []
test_sharp = []
for direc in sorted(os.listdir(GOPRO_test_path)):
if args.is_gamma:
blur = os.path.join(os.path.join(GOPRO_test_path, direc), 'blur_gamma')
else:
blur = os.path.join(os.path.join(GOPRO_test_path, direc), 'blur')
sharp = os.path.join(os.path.join(GOPRO_test_path, direc), 'sharp')
sharp_imgs = sorted(os.listdir(sharp))
for i, img in enumerate(sorted(os.listdir(blur))):
test_blur.append(os.path.join(blur, img))
test_sharp.append(os.path.join(sharp, sharp_imgs[i]))
test_blur = np.asarray(test_blur)
test_sharp = np.asarray(test_sharp)
random_index = np.random.permutation(len(test_blur))[:args.test_num]
for index in random_index:
subprocess.call(['cp', test_blur[index], os.path.join(args.output_path, 'test/blur/%s'%('_'.join(test_blur[index].split('/')[-3:])))])
subprocess.call(['cp', test_sharp[index], os.path.join(args.output_path, 'test/sharp/%s'%('_'.join(test_sharp[index].split('/')[-3:])))])