-
Notifications
You must be signed in to change notification settings - Fork 0
/
main.py
70 lines (49 loc) · 2.16 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
import os
import read_paths
from my_net import MyNet
from restore import Restore
from image_helper import read_image, save_compare_figure
from keras.models import load_model
from nets_factory import select_model
from get_input import get_input
TO_REPLACE = '#'
TEST_PATH = TO_REPLACE + '_dataset\\test'
TRAIN_PATH = TO_REPLACE + '_dataset\\train'
RESULTS_PATH = TO_REPLACE + '_dataset\\results'
NAME_TO_SAVE = TO_REPLACE + '_model.h5'
FULL_MODE = 'full'
QUICK_MODE = 'quick'
MY_MODES = [QUICK_MODE, FULL_MODE]
def select_mode():
mode_options = list('to ' + mode_option + ' press ' + str(i) for i, mode_option in enumerate(MY_MODES))
message = 'Please select the requested mode:\n'
mode = get_input(message, mode_options, MY_MODES)
return mode == QUICK_MODE
def create_model(my_model, corrupter, net_name, name_to_save):
net = MyNet()
paths = read_paths.list_paths(TRAIN_PATH.replace(TO_REPLACE, net_name))
my_model.build(paths, corrupter)
mode = select_mode()
trained_model = my_model.learn_model(net, quick_mode=mode)
trained_model.save(name_to_save)
return trained_model
def restore_image(trained_model, restore, corrupter, im, i, net_name):
im_corrupted = corrupter.corruption(im)
im_restored = restore.restore_image(im_corrupted, trained_model)
save_compare_figure([im_corrupted, im_restored], RESULTS_PATH.replace(TO_REPLACE, net_name), i)
def restore_images(trained_model, corrupter, net_name):
paths = read_paths.list_paths(TEST_PATH.replace(TO_REPLACE, net_name))
restore = Restore()
for i, path in enumerate(paths):
restore_image(trained_model, restore, corrupter, read_image(path), i, net_name)
def main():
my_model, corrupter, net_name = select_model()
name_to_save = NAME_TO_SAVE.replace(TO_REPLACE, net_name)
if os.path.isfile(name_to_save):
trained_model = load_model(name_to_save)
else:
trained_model = create_model(my_model, corrupter, net_name, name_to_save)
restore_images(trained_model, corrupter, net_name)
if __name__ == "__main__":
# os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
main()