-
Notifications
You must be signed in to change notification settings - Fork 0
/
main.py
57 lines (49 loc) · 1.57 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
import argparse
import os
import tifffile
import pytest
from src import train, networks, util
from config import Config
import os
os.environ["KMP_DUPLICATE_LIB_OK"]="TRUE"
def main(mode, offline, tag):
"""[summary]
:param mode: [description]
:type mode: [type]
:param offline: [description]
:type offline: [type]
:param tag: [description]
:type tag: [type]
:raises ValueError: [description]
"""
print("Running in {} mode, tagged {}, offline {}".format(mode, tag, offline))
# Initialise Config object
c = Config(tag)
if mode == 'train':
overwrite = util.check_existence(tag, overwrite=True)
util.initialise_folders(tag, overwrite)
netD, netG = networks.make_nets(c, overwrite)
train(c, netG, netD, offline=offline, overwrite=overwrite)
elif mode == 'generate':
netD, netG = networks.make_nets(c, training=0)
net_g, net_d = netG(), netD()
util.generate(c, net_g, 97)
print("Img generated")
elif mode == 'test':
print('Performing unit tests')
pt = pytest.main(["-x", "src/test.py"])
else:
raise ValueError("Mode not recognised")
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("mode")
parser.add_argument("-t", "--tag")
parser.add_argument('-o', '--offline', action='store_true',
help='disable wandb')
args = parser.parse_args()
if args.tag:
tag = args.tag
else:
tag = 'test'
main(args.mode, args.offline, tag)
# main('train', False, 'test')