forked from foamliu/Colorful-Image-Colorization
-
Notifications
You must be signed in to change notification settings - Fork 0
/
run_demo.py
82 lines (67 loc) · 2.56 KB
/
run_demo.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
import os, random, shutil
import glob
import string
import generate_demo_file_names
def remove_folder_content(path):
for filename in os.listdir(path):
file_path = os.path.join(path, filename)
try:
if os.path.isfile(file_path) or os.path.islink(file_path):
os.unlink(file_path)
elif os.path.isdir(file_path):
shutil.rmtree(file_path)
except Exception as e:
print('Failed to delete %s. Reason: %s' % (file_path, e))
files_to_pick = 10
datasets_folder = "../datasets/"
datasets = [
f"{datasets_folder}google_landscapes",
f"{datasets_folder}celebryci",
f"{datasets_folder}kwiatki",
f"{datasets_folder}landscape",
f"{datasets_folder}ptaki"
]
demo_input_path = './demo_input'
demo_output_path = './demo_output'
def run_demo(picked_datasets_index):
# First clear some folders
remove_folder_content(demo_input_path)
remove_folder_content(demo_output_path)
# Now pick some random images from dataset
random_files = []
folder_path = datasets[picked_datasets_index]
for _ in range(files_to_pick):
random_files.append(folder_path + "/" + random.choice(os.listdir(folder_path)))
for random_file in random_files:
shutil.copy2(random_file, demo_input_path)
# Now generate names for them
generate_demo_file_names.main()
import demo
demo.main()
# Now save results with renamed name to results folder
results_folder = '../results'
N = 10
file_index = 0
for file in os.listdir(demo_output_path):
# 0 - gt
# 1 - black and white
# 2 - out
file_extension = file.split('.')[-1]
if file_index % 3 == 0:
dataset_name = datasets[picked_datasets_index].split('/')[-1]
random_name = ''.join(random.choices(string.ascii_uppercase + string.digits, k=N))
gt = f"{random_name}_gt.{file_extension}"
bandw = f"{random_name}_black_and_white.{file_extension}"
out = f"{random_name}_out.{file_extension}"
shutil.copy2(f"{demo_output_path}/{file}", f"{results_folder}/{dataset_name}/{gt}")
elif file_index % 3 == 1:
# Black and white
shutil.copy2(f"{demo_output_path}/{file}", f"{results_folder}/{dataset_name}/{bandw}")
else:
# out
shutil.copy2(f"{demo_output_path}/{file}", f"{results_folder}/{dataset_name}/{out}")
file_index += 1
for i in range(len(datasets)):
run_demo(i)
remove_folder_content(demo_input_path)
remove_folder_content(demo_output_path)