-
Notifications
You must be signed in to change notification settings - Fork 11
/
main.py
90 lines (64 loc) · 3.12 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
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
import os
import argparse
from train import network_train
from test import network_test
def build_parser():
def str2bool(v):
if isinstance(v, bool):
return v
if v.lower() in ('yes', 'true', 't', 'y', '1'):
return True
elif v.lower() in ('no', 'false', 'f', 'n', '0'):
return False
else:
raise argparse.ArgumentTypeError('Boolean value expected !!!')
parser = argparse.ArgumentParser()
parser.add_argument('--cuda-device-no', type=int,
help='cpu: -1, gpu: 0 ~ n ', default=0)
parser.add_argument('--train-flag', type=str2bool,
help='Train flag', required=True)
parser.add_argument('--max-iter', type=int,
help='Train iterations', default=15000)
parser.add_argument('--batchs', type=int,
help='Batch size', default=8)
parser.add_argument('--lr', type=float,
help='Learning rate to optimize network', default=0.1)
parser.add_argument('--check-iter', type=int,
help='Number of iteration to check training logs', default=100)
parser.add_argument('--imsize', type=int,
help='Size for resize image during training', default=256)
parser.add_argument('--cropsize', type=int,
help='Size for crop image durning training', default=None)
parser.add_argument('--vgg-flag', type=str,
help='VGG flag for calculating losses', default='vgg16')
parser.add_argument('--content-layers', type=int, nargs='+',
help='layer indices to extract content features', default=[15])
parser.add_argument('--style-layers', type=int, nargs='+',
help='layer indices to extract style features', default=[3, 8, 15, 22])
parser.add_argument('--content-weight', type=float,
help='content loss weight', default=1.0)
parser.add_argument('--style-weight', type=float,
help='style loss weight', default=30.0)
parser.add_argument('--tv-weight', type=float,
help='tv loss weight', default=1.0)
parser.add_argument('--train-content', type=str,
help='Content images path for training')
parser.add_argument('--train-style', type=str,
help='The taget syle image path for training')
parser.add_argument('--save-path', type=str,
help='Save path', default='./trained_models/')
parser.add_argument('--model-load-path', type=str,
help="Trained model load path")
parser.add_argument('--test-content', type=str,
help="test content image path")
parser.add_argument('--output', type=str,
help='output image path to save the stylized image', default='stylized.jpg')
return parser
if __name__ == '__main__':
parser = build_parser()
args= parser.parse_args()
os.environ['CUDA_VISIBLE_DEVICES'] = str(args.cuda_device_no)
if args.train_flag:
transform_network = network_train(args)
else:
network_test(args)