-
Notifications
You must be signed in to change notification settings - Fork 6
/
main.py
57 lines (42 loc) · 2.37 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 numpy as np
import faiss
import tensorflow as tf
from data import *
from utils import args
from architectures import *
def main():
"""
Reads data and cmd arguments and trains models
"""
if args.args.data == 'HERA':
data = load_hera(args.args)
elif args.args.data == 'LOFAR':
data = load_lofar(args.args)
elif args.args.data == 'HIDE':
data = load_hide(args.args)
(unet_train_dataset, train_data, train_labels, train_masks,
ae_train_dataset, ae_train_data, ae_train_labels,
test_data, test_labels, test_masks,test_masks_orig) = data
print(" __________________________________ \n Save name {}".format(
args.args.model_name))
print(" __________________________________ \n")
if args.args.model == 'UNET':
train_unet(unet_train_dataset, train_data, train_labels, train_masks, test_data, test_labels, test_masks, test_masks_orig, args.args)
if args.args.model == 'RNET':
train_rnet(unet_train_dataset, train_data, train_labels, train_masks, test_data, test_labels, test_masks, test_masks_orig, args.args)
if args.args.model == 'RFI_NET':
train_rfi_net(unet_train_dataset, train_data, train_labels, train_masks, test_data, test_labels, test_masks, test_masks_orig, args.args)
elif args.args.model == 'DKNN':
train_resnet(ae_train_dataset, ae_train_data, ae_train_labels, test_data, test_labels, test_masks, test_masks_orig, args.args)
elif args.args.model == 'AE':
train_ae(ae_train_dataset, ae_train_data, ae_train_labels, test_data, test_labels, test_masks, test_masks_orig, args.args)
elif args.args.model == 'AE-SSIM':
train_ae_ssim(ae_train_dataset, ae_train_data, ae_train_labels, test_data, test_labels, test_masks, test_masks_orig, args.args)
elif args.args.model == 'DAE':
train_dae(ae_train_dataset, ae_train_data, ae_train_labels, test_data, test_labels, test_masks, test_masks_orig,args.args)
elif args.args.model == 'AOFlagger':
end_routine(train_data, test_data, test_data, test_masks, test_masks_orig, None, 'AOFlagger', args.args)
elif args.args.model == 'fine_tune':
train_fine_tune(unet_train_dataset, ae_train_data, train_data, train_labels, train_masks, test_data, test_labels, test_masks, test_masks_orig, args.args)
if __name__ == '__main__':
main()