-
Notifications
You must be signed in to change notification settings - Fork 10
/
dog_breed_detector_trainer.py
127 lines (99 loc) · 4.35 KB
/
dog_breed_detector_trainer.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
'''
Trains the dog breed detector and saves the final model for the serving
'''
import os
import shutil
from glob import glob
from data_extractor import DataExtractor
from dog_breed_detector_model import DogBreedDetectorModel
from final_model import FinalModel
from tensorflow.python.keras.applications.densenet import preprocess_input
from tensorflow.python.keras.estimator import model_to_estimator
import tensorflow as tf
def preprocess_image(image_buffer, image_shape=(224, 224, 3)):
'''
Preprocess JPEG encoded bytes to 3D float Tensor and rescales
it so that pixels are in a range of [-1, 1]
:param image_buffer: Buffer that contains JPEG image
:return: 4D image tensor (1, width, height,channels) with pixels scaled
to [-1, 1]. First dimension is a batch size (1 is our case)
'''
# Decode the string as an RGB JPEG.
# Note that the resulting image contains an unknown height and width
# that is set dynamically by decode_jpeg. In other words, the height
# and width of image is unknown at compile-time.
image_buffer = tf.reshape(image_buffer, [])
img = tf.image.decode_jpeg(image_buffer, channels=3)
# Networks accept images in batches.
# The first dimension usually represents the batch size.
img = tf.image.resize_images(img, size=(image_shape[1], image_shape[0]))
# Finally, preprocess image for the model
img = preprocess_input(img)
return img
def serving_input_receiver_fn():
'''
An input receiver that expects a serialized tf.Example
:return: a return type for a serving_input_receiver_fn
'''
serialized_tf_example = tf.placeholder(tf.string, name='input_image')
feature_configs = {
'image/encoded': tf.FixedLenFeature(
shape=[], dtype=tf.string),
}
tf_example = tf.parse_example(serialized_tf_example, feature_configs)
jpegs = tf_example['image/encoded']
receiver_tensors = {'examples': jpegs}
images = tf.map_fn(preprocess_image, jpegs, dtype=tf.float32)
features = {'input_2': images}
return tf.estimator.export.ServingInputReceiver(features, receiver_tensors)
def export_for_serving(model):
'''
Converts model to the TensorFlow estimator and saves it to the disk
:param model: keras model to prepare for serving
'''
export_dir = 'export_model/'
if os.path.exists(export_dir):
shutil.rmtree(export_dir)
tf_estimator = model_to_estimator(keras_model=model)
tf_estimator.export_savedmodel(
export_dir,
serving_input_receiver_fn,
strip_default_attrs=True)
def train_test_dog_breed_detector(do_model_test=True):
'''
Trains and tests the dog breed detector model
:param do_model_test: set to True for running prediction on the test data
:return: dog breed detector model
'''
# get train and validation data and labels
data_extractor = DataExtractor('dog_images/train', 'dog_images/valid', 'dog_images/test')
train_data = data_extractor.load_train_data()
train_labels, train_num_dog_breeds = data_extractor.load_train_labels()
valid_data = data_extractor.load_valid_data()
valid_labels, valid_num_dog_breeds = data_extractor.load_valid_labels()
assert train_num_dog_breeds == valid_num_dog_breeds
# create and init the model
model = DogBreedDetectorModel(num_dog_breeds=train_num_dog_breeds)
model.init_model(train_data)
# train the model
if not model.load_model():
model.train(train_data, train_labels, valid_data, valid_labels)
# load test data and predict
if do_model_test:
test_data = data_extractor.load_test_data()
test_labels, test_num_dog_breeds = data_extractor.load_test_labels()
assert train_num_dog_breeds == test_num_dog_breeds
model.predict(test_data, test_labels)
return model
def create_save_final_model(top_model):
'''
Creates and save the final model - it consists of pretrained DenseNet201 and
dog breed detector on top of it
:param top_model: dog breed model is expected
'''
dog_breed_names = [item[21:-1] for item in sorted(glob('{}/*/'.format('dog_images/train')))]
final_model = FinalModel(top_model.model, dog_breed_names)
export_for_serving(final_model.model)
if __name__ == '__main__':
dog_breed_detector_model = train_test_dog_breed_detector(do_model_test=False)
create_save_final_model(dog_breed_detector_model)