forked from divamgupta/image-segmentation-keras
-
Notifications
You must be signed in to change notification settings - Fork 0
/
train.py
executable file
·83 lines (59 loc) · 2.85 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
import argparse
import Models , LoadBatches
parser = argparse.ArgumentParser()
parser.add_argument("--save_weights_path", type = str )
parser.add_argument("--train_images", type = str )
parser.add_argument("--train_annotations", type = str )
parser.add_argument("--n_classes", type=int )
parser.add_argument("--input_height", type=int , default = 224 )
parser.add_argument("--input_width", type=int , default = 224 )
parser.add_argument('--validate',action='store_false')
parser.add_argument("--val_images", type = str , default = "")
parser.add_argument("--val_annotations", type = str , default = "")
parser.add_argument("--epochs", type = int, default = 5 )
parser.add_argument("--batch_size", type = int, default = 2 )
parser.add_argument("--val_batch_size", type = int, default = 2 )
parser.add_argument("--load_weights", type = str , default = "")
parser.add_argument("--model_name", type = str , default = "")
parser.add_argument("--optimizer_name", type = str , default = "adadelta")
args = parser.parse_args()
train_images_path = args.train_images
train_segs_path = args.train_annotations
train_batch_size = args.batch_size
n_classes = args.n_classes
input_height = args.input_height
input_width = args.input_width
validate = args.validate
save_weights_path = args.save_weights_path
epochs = args.epochs
load_weights = args.load_weights
optimizer_name = args.optimizer_name
model_name = args.model_name
if validate:
val_images_path = args.val_images
val_segs_path = args.val_annotations
val_batch_size = args.val_batch_size
modelFns = { 'vgg_segnet':Models.VGGSegnet.VGGSegnet , 'vgg_unet':Models.VGGUnet.VGGUnet , 'vgg_unet2':Models.VGGUnet.VGGUnet2 , 'fcn8':Models.FCN8.FCN8 , 'fcn32':Models.FCN32.FCN32 }
modelFN = modelFns[ model_name ]
m = modelFN( n_classes , input_height=input_height, input_width=input_width )
m.compile(loss='categorical_crossentropy',
optimizer= optimizer_name ,
metrics=['accuracy'])
if len( load_weights ) > 0:
m.load_weights(load_weights)
print "Model output shape" , m.output_shape
output_height = m.outputHeight
output_width = m.outputWidth
G = LoadBatches.imageSegmentationGenerator( train_images_path , train_segs_path , train_batch_size, n_classes , input_height , input_width , output_height , output_width )
if validate:
G2 = LoadBatches.imageSegmentationGenerator( val_images_path , val_segs_path , val_batch_size, n_classes , input_height , input_width , output_height , output_width )
if not validate:
for ep in range( epochs ):
m.fit_generator( G , 512 , epochs=1 )
m.save_weights( save_weights_path + "." + str( ep ) )
m.save( save_weights_path + ".model." + str( ep ) )
else:
for ep in range( epochs ):
m.fit_generator( G , 512 , validation_data=G2 , validation_steps=200 , epochs=1 )
m.save_weights( save_weights_path + "." + str( ep ) )
m.save( save_weights_path + ".model." + str( ep ) )