-
Notifications
You must be signed in to change notification settings - Fork 528
/
main.py
134 lines (95 loc) · 3.41 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
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
import os
import cv2
import random
import numpy as np
import torch
import argparse
from shutil import copyfile
from src.config import Config
from src.edge_connect import EdgeConnect
def main(mode=None):
r"""starts the model
Args:
mode (int): 1: train, 2: test, 3: eval, reads from config file if not specified
"""
config = load_config(mode)
# cuda visble devices
os.environ['CUDA_VISIBLE_DEVICES'] = ','.join(str(e) for e in config.GPU)
# init device
if torch.cuda.is_available():
config.DEVICE = torch.device("cuda")
torch.backends.cudnn.benchmark = True # cudnn auto-tuner
else:
config.DEVICE = torch.device("cpu")
# set cv2 running threads to 1 (prevents deadlocks with pytorch dataloader)
cv2.setNumThreads(0)
# initialize random seed
torch.manual_seed(config.SEED)
torch.cuda.manual_seed_all(config.SEED)
np.random.seed(config.SEED)
random.seed(config.SEED)
# build the model and initialize
model = EdgeConnect(config)
model.load()
# model training
if config.MODE == 1:
config.print()
print('\nstart training...\n')
model.train()
# model test
elif config.MODE == 2:
print('\nstart testing...\n')
model.test()
# eval mode
else:
print('\nstart eval...\n')
model.eval()
def load_config(mode=None):
r"""loads model config
Args:
mode (int): 1: train, 2: test, 3: eval, reads from config file if not specified
"""
parser = argparse.ArgumentParser()
parser.add_argument('--path', '--checkpoints', type=str, default='./checkpoints', help='model checkpoints path (default: ./checkpoints)')
parser.add_argument('--model', type=int, choices=[1, 2, 3, 4], help='1: edge model, 2: inpaint model, 3: edge-inpaint model, 4: joint model')
# test mode
if mode == 2:
parser.add_argument('--input', type=str, help='path to the input images directory or an input image')
parser.add_argument('--mask', type=str, help='path to the masks directory or a mask file')
parser.add_argument('--edge', type=str, help='path to the edges directory or an edge file')
parser.add_argument('--output', type=str, help='path to the output directory')
args = parser.parse_args()
config_path = os.path.join(args.path, 'config.yml')
# create checkpoints path if does't exist
if not os.path.exists(args.path):
os.makedirs(args.path)
# copy config template if does't exist
if not os.path.exists(config_path):
copyfile('./config.yml.example', config_path)
# load config file
config = Config(config_path)
# train mode
if mode == 1:
config.MODE = 1
if args.model:
config.MODEL = args.model
# test mode
elif mode == 2:
config.MODE = 2
config.MODEL = args.model if args.model is not None else 3
config.INPUT_SIZE = 0
if args.input is not None:
config.TEST_FLIST = args.input
if args.mask is not None:
config.TEST_MASK_FLIST = args.mask
if args.edge is not None:
config.TEST_EDGE_FLIST = args.edge
if args.output is not None:
config.RESULTS = args.output
# eval mode
elif mode == 3:
config.MODE = 3
config.MODEL = args.model if args.model is not None else 3
return config
if __name__ == "__main__":
main()