-
Notifications
You must be signed in to change notification settings - Fork 4
/
train.py
162 lines (115 loc) · 6.04 KB
/
train.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
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
import os
import sys
import tensorflow as tf
from keras.backend.tensorflow_backend import set_session
import keras , warnings
from keras.optimizers import *
from sacred import Experiment
from models import *
from utils import *
def configuration():
keras.backend.clear_session()
tf.reset_default_graph()
warnings.filterwarnings('ignore')
os.environ['CUDA_DEVICE_ORDER']='PCI_BUS_ID'
config = tf.ConfigProto(log_device_placement=False, allow_soft_placement=True)
config.gpu_options.allow_growth = True
#config.gpu_options.per_process_gpu_memory_fraction=0.95#0.95
config.gpu_options.visible_device_list="0"
set_session(tf.Session(config=config))
def get_dirs_or_files(input_data):
if os.path.isdir(input_data):
image_input, labels_input = os.path.join(input_data, 'images/'), os.path.join(input_data, 'labels/')
# Check if training dir exists
assert os.path.isdir(image_input), "{} is not a directory".format(image_input)
assert os.path.isdir(labels_input), "{} is not a directory".format(labels_input)
return image_input, labels_input
ex = Experiment()
@ex.config
def config_params():
n_classes=None # Number of classes. If your case study is binary case the set it to 2 and otherwise give your number of cases.
n_epochs=1
input_height=224*1
input_width=224*1
weight_decay=1e-6 # Weight decay of l2 regularization of model layers.
n_patches=20 # Number of patches of each image which is randomly created.
n_batch=1 # Number of batches at each iteration.
learning_rate=1e-4
patches=False # Make patches of image in order to use all information of image. In the case of page
# extraction this should be set to false since model should see all image.
augmentation=False
flip_aug=False # Flip image (augmentation).
elastic_aug=False # Elastic transformation (augmentation).
blur_aug=False # Blur patches of image (augmentation).
scaling=False # Scaling of patches (augmentation) will be imposed if this set to true.
binarization=False # Otsu thresholding. Used for augmentation in the case of binary case like textline prediction. For multicases should not be applied.
dir_train=None # Directory of training dataset (sub-folders should be named images and labels).
dir_eval=None # Directory of validation dataset (sub-folders should be named images and labels).
dir_output=None # Directory of output where the model should be saved.
pretraining=False # Set true to load pretrained weights of resnet50 encoder.
blur_k=['blur','guass','median'] # Used in order to blur image. Used for augmentation.
scales=[0.8, 1.2] # Scale patches with these scales. Used for augmentation.
flip_index=[0,1,-1] # Flip image. Used for augmentation.
@ex.automain
def run(n_classes,n_epochs,input_height,
input_width,weight_decay,n_patches,
n_batch,patches,augmentation,flip_aug,
elastic_aug,blur_aug,scaling,
binarization,blur_k,scales,dir_train,
flip_index,dir_eval ,dir_output,pretraining,learning_rate):
dir_img,dir_seg=get_dirs_or_files(dir_train)
dir_img_val,dir_seg_val=get_dirs_or_files(dir_eval)
# make first a directory in output for both training and evaluations in order to flow data from these directories.
dir_train_flowing=os.path.join(dir_output,'train')
dir_eval_flowing=os.path.join(dir_output,'eval')
dir_flow_train_imgs=os.path.join(dir_train_flowing,'images/')
dir_flow_train_labels=os.path.join(dir_train_flowing,'labels/')
dir_flow_eval_imgs=os.path.join(dir_eval_flowing,'images/')
dir_flow_eval_labels=os.path.join(dir_eval_flowing,'labels/')
if os.path.isdir(dir_train_flowing):
os.system('rm -rf '+dir_train_flowing)
os.makedirs(dir_train_flowing)
else:
os.makedirs(dir_train_flowing)
if os.path.isdir(dir_eval_flowing):
os.system('rm -rf '+dir_eval_flowing)
os.makedirs(dir_eval_flowing)
else:
os.makedirs(dir_eval_flowing)
os.mkdir(dir_flow_train_imgs)
os.mkdir(dir_flow_train_labels)
os.mkdir(dir_flow_eval_imgs)
os.mkdir(dir_flow_eval_labels)
#set the gpu configuration
configuration()
#writing patches into a sub-folder in order to be flowed from directory.
provide_patches(dir_img,dir_seg,dir_flow_train_imgs,
dir_flow_train_labels,
input_height,input_width,blur_k,blur_aug,n_patches,
flip_aug,binarization,elastic_aug,scaling,scales,flip_index,
augmentation=augmentation,patches=patches)
provide_patches(dir_img_val,dir_seg_val,dir_flow_eval_imgs,
dir_flow_eval_labels,
input_height,input_width,blur_k,blur_aug,n_patches,
flip_aug,binarization,elastic_aug,scaling,scales,flip_index,
augmentation=False,patches=True)
#get our model.
model = resnet50_unet(n_classes, input_height, input_width,weight_decay,pretraining)
#if you want to see the model structure just uncomment model summary.
#model.summary()
model.compile(loss='categorical_crossentropy',
optimizer = Adam(lr=learning_rate),metrics=['accuracy'])
#generating train and evaluation data
train_gen = data_gen(dir_flow_train_imgs,dir_flow_train_labels, batch_size = n_batch,
input_height=input_height, input_width=input_width,n_classes=n_classes )
val_gen = data_gen(dir_flow_eval_imgs,dir_flow_eval_labels, batch_size = n_batch,
input_height=input_height, input_width=input_width,n_classes=n_classes )
model.fit_generator(
train_gen,
steps_per_epoch=int(len(os.listdir(dir_flow_train_imgs))/n_batch),
validation_data=val_gen,
validation_steps=1,
epochs=n_epochs)
os.system('rm -rf '+dir_train_flowing)
os.system('rm -rf '+dir_eval_flowing)
model.save(dir_output+'/'+'model'+'.h5')