forked from Ben-Louis/Deep-Image-Analogy-PyTorch
-
Notifications
You must be signed in to change notification settings - Fork 0
/
main.py
71 lines (50 loc) · 1.85 KB
/
main.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
import os
from utils import load_image
import argparse
from DeepAnalogy import analogy
import cv2
def str2bool(v):
return v.lower() in ('true')
if __name__=="__main__":
parser = argparse.ArgumentParser()
parser.add_argument('--resize_ratio', type=float, default=0.5)
parser.add_argument('--weight', type=int, default=2, choices=[2,3])
parser.add_argument('--img_A_path', type=str, default='data/demo/ava.png')
parser.add_argument('--img_BP_path', type=str, default='data/demo/mona.png')
parser.add_argument('--use_cuda', type=str2bool, default=True)
args = parser.parse_args()
# load images
print('Loading images...', end='')
img_A = load_image(args.img_A_path, args.resize_ratio)
img_BP = load_image(args.img_BP_path, args.resize_ratio)
print('\rImages loaded successfully!')
# setting parameters
config = dict()
params = {
'layers': [29,20,11,6,1],
'iter': 10,
}
config['params'] = params
if args.weight == 2:
config['weights'] = [1.0, 0.8, 0.7, 0.6, 0.1, 0.0]
elif args.weight == 3:
config['weights'] = [1.0, 0.9, 0.8, 0.7, 0.2, 0.0]
config['sizes'] = [3,3,3,5,5,3]
config['rangee'] = [32,6,6,4,4,2]
config['use_cuda'] = args.use_cuda
config['lr'] = [0.1, 0.005, 0.005, 0.00005]
# Deep-Image-Analogy
print("\n##### Deep Image Analogy - start #####")
img_AP, img_B, elapse = analogy(img_A, img_BP, config)
print("##### Deep Image Analogy - end | Elapse:"+elapse+" #####")
# save result
content = os.listdir('results')
count = 1
for c in content:
if os.path.isdir('results/'+c):
count += 1
save_path = 'results/expr_{}'.format(count)
os.mkdir(save_path)
cv2.imwrite(save_path+'/img_AP.png', img_AP)
cv2.imwrite(save_path+'/img_B.png', img_B)
print('Image saved!')