Skip to content

Commit 60ca5c3

Browse files
previprevi
committed
tflite completes
1 parent 0f52b69 commit 60ca5c3

File tree

15 files changed

+663
-1220
lines changed

15 files changed

+663
-1220
lines changed

‎cnn_classifier.py

Lines changed: 27 additions & 74 deletions
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,7 @@
2929
logger = logging.getLogger(__name__)
3030

3131
class CNNClassifier(object):
32-
def __init__(self, model_file, label_file, input_layer="input", output_layer="final_result", input_height=128, input_width=128, input_mean=127.5, input_std=127.5):
32+
def __init__(self, model_file, label_file):
3333
logger.info(model_file)
3434
self._interpreter = Interpreter(model_path=model_file)
3535
self._interpreter.set_num_threads(4)
@@ -40,73 +40,20 @@ def __init__(self, model_file, label_file, input_layer="input", output_layer="fi
4040
self._input_height=self._input_details[0]['shape'][1]
4141
self._input_width=self._input_details[0]['shape'][2]
4242
self._floating_model = (self._input_details[0]['dtype'] == np.float32)
43-
#input_name = "import/" + input_layer
44-
#output_name = "import/" + output_layer
45-
#self._input_operation = self._graph.get_operation_by_name(input_name)
46-
#self._output_operation = self._graph.get_operation_by_name(output_name)
47-
#self._session = tf.compat.v1.Session(graph=self._graph)
48-
#self._graph_norm = tf.Graph()
49-
#with self._graph_norm.as_default():
50-
# image_mat = tf.compat.v1.placeholder(tf.float32, None, name="image_rgb_in")
51-
# float_caster = tf.cast(image_mat, tf.float32)
52-
# dims_expander = tf.expand_dims(float_caster, 0)
53-
# resized = tf.compat.v1.image.resize_bilinear(dims_expander, [input_height, input_width])
54-
# normalized = tf.divide(tf.subtract(resized, [input_mean]), [input_std], name="image_norm_out")
55-
# self._input_operation_norm = self._graph_norm.get_operation_by_name("image_rgb_in")
56-
# self._output_operation_norm = self._graph_norm.get_operation_by_name("image_norm_out")
57-
#self._sess_norm = tf.Session(graph=self._graph_norm)
5843

5944
def close(self):
6045
pass
61-
#self._session.close()
62-
#self._sess_norm.close()
63-
64-
#def load_graph(self, model_file):
65-
# graph = tf.Graph()
66-
# graph_def = tf.compat.v1.GraphDef()
67-
#
68-
# with open(model_file, "rb") as f:
69-
# graph_def.ParseFromString(f.read())
70-
# with graph.as_default():
71-
# tf.import_graph_def(graph_def)
72-
#
73-
# return graph
74-
#
75-
#def read_tensor_from_image_file(self, file_name, input_height=299, input_width=299, input_mean=0, input_std=255):
76-
# input_name = "file_reader"
77-
# output_name = "normalized"
78-
#
79-
# file_reader = tf.read_file(file_name, input_name)
80-
#
81-
# if file_name.endswith(".png"):
82-
# image_reader = tf.image.decode_png(file_reader, channels=3, name='png_reader')
83-
# elif file_name.endswith(".gif"):
84-
# image_reader = tf.squeeze(tf.image.decode_gif(file_reader, name='gif_reader'))
85-
# elif file_name.endswith(".bmp"):
86-
# image_reader = tf.image.decode_bmp(file_reader, name='bmp_reader')
87-
# else:
88-
# image_reader = tf.image.decode_jpeg(file_reader, channels=3, name='jpeg_reader')
89-
#
90-
# float_caster = tf.cast(image_reader, tf.float32)
91-
# dims_expander = tf.expand_dims(float_caster, 0);
92-
# resized = tf.image.resize_bilinear(dims_expander, [self.input_height, self.input_width])
93-
# normalized = tf.divide(tf.subtract(resized, [input_mean]), [input_std])
94-
# sess = tf.Session()
95-
#
96-
# result = sess.run(normalized)
97-
# sess.close()
98-
#
99-
# return result
100-
#
101-
#def read_tensor_from_image_mat(self, image_mat, input_height=299, input_width=299, input_mean=0, input_std=255):
102-
# result = self._sess_norm.run(self._output_operation_norm.outputs[0], {self._input_operation_norm.outputs[0]: image_mat})
103-
# return result
104-
def read_tensor_from_image_mat(self, image_mat, input_height=299, input_width=299, input_mean=0, input_std=255):
46+
47+
def read_tensor_from_image_file(self, file_name):
48+
image = cv2.imread(file_name)
49+
return self.read_tensor_from_image_mat(image)
50+
51+
def read_tensor_from_image_mat(self, image_mat):
10552
frame_rgb = cv2.cvtColor(image_mat, cv2.COLOR_BGR2RGB)
10653
frame_resized = cv2.resize(frame_rgb, (self._input_width, self._input_height))
10754
input_data = np.expand_dims(frame_resized, axis=0)
10855

109-
# Normalize pixel values if using a floating model (i.e. if model is non-quantized)
56+
# Normalize pixel values if using a floating model (i.e. if model is non-quantized)
11057
if self._floating_model:
11158
input_mean = 127.5
11259
input_std = 127.5
@@ -124,33 +71,39 @@ def classify_image(self,
12471
image_file_or_mat,
12572
top_results=3):
12673
input_image = None
127-
#if isinstance(image_file_or_mat, str):
128-
# t = self.read_tensor_from_image_file(file_name=image_file_or_mat)
129-
#else:
130-
input_image = self.read_tensor_from_image_mat(image_file_or_mat)
74+
if isinstance(image_file_or_mat, str):
75+
input_image = self.read_tensor_from_image_file(file_name=image_file_or_mat)
76+
else:
77+
input_image = self.read_tensor_from_image_mat(image_file_or_mat)
13178

13279
self._interpreter.set_tensor(self._input_details[0]['index'], input_image)
13380
self._interpreter.invoke()
134-
scores = self._interpreter.get_tensor(self._output_details[0]['index'])[0] # Bounding box coordinates of detected objects
135-
81+
scores = self._interpreter.get_tensor(self._output_details[0]['index'])[0]
82+
83+
#print("scores: " + str(scores))
84+
confidence = 0.4
85+
base = 1
86+
# normalize to int8 for quantized models
87+
if len(scores)>0 and (scores[0] == int(scores[0])):
88+
confidence = 128
89+
base = 256
13690
pairs = []
13791
for i in range(0, len(scores)):
138-
if scores[i] > 128:
92+
if scores[i] > confidence:
13993
object_name = self._labels[i]
140-
pairs.append((object_name, int(100*scores[i]/256)))
94+
pairs.append((object_name, int(100*scores[i]/base)))
14195

14296
pairs = sorted(pairs, key=lambda x: x[1], reverse=True)[:top_results]
143-
logger.info(str(pairs))
14497
return pairs
14598

14699
def detect_objects(self,
147100
image_file_or_mat,
148101
top_results=3):
149102
input_image = None
150-
#if isinstance(image_file_or_mat, str):
151-
# t = self.read_tensor_from_image_file(file_name=image_file_or_mat)
152-
#else:
153-
input_image = self.read_tensor_from_image_mat(image_file_or_mat)
103+
if isinstance(image_file_or_mat, str):
104+
input_image = self.read_tensor_from_image_file(file_name=image_file_or_mat)
105+
else:
106+
input_image = self.read_tensor_from_image_mat(image_file_or_mat)
154107

155108
self._interpreter.set_tensor(self._input_details[0]['index'], input_image)
156109
self._interpreter.invoke()

‎cnn_manager.py

Lines changed: 14 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -63,6 +63,17 @@ def get_models(self):
6363
def get_model_status(self, model_name):
6464
return self._models[model_name]
6565

66+
@classmethod
67+
def get_model_info(cls, architecture):
68+
model_info = architecture.split("/")[1].split("_")
69+
return model_info
70+
71+
@classmethod
72+
def get_model_shape(cls, architecture):
73+
model_info = cls.get_model_info(architecture)
74+
size = int(model_info[3])
75+
return (size, size)
76+
6677
def _save_model_meta(self):
6778
f = open(MODEL_METADATA, "w")
6879
json.dump(self._models, f)
@@ -92,7 +103,7 @@ def train_new_model(self,
92103
self._trainers[model_name] = trainer
93104

94105
def save_model_status(self, model_name, architecture, status):
95-
model_info = architecture.split("_")
106+
model_info = self.get_model_info(architecture)
96107
self._models[model_name] = {"status": status, "image_height": model_info[3], "image_width": model_info[3], "output_layer": "final_result"}
97108
self._save_model_meta()
98109

@@ -104,10 +115,7 @@ def load_model(self, model_name):
104115
model_info = self._models.get(model_name)
105116
if model_info:
106117
return CNNClassifier(model_file=MODEL_PATH + "/" + model_name + ".tflite",
107-
label_file=MODEL_PATH + "/" + model_name + ".txt",
108-
output_layer=model_info["output_layer"],
109-
input_height=int(model_info["image_height"]),
110-
input_width=int(model_info["image_width"]))
118+
label_file=MODEL_PATH + "/" + model_name + ".txt")
111119
return None
112120
class TrainThread(threading.Thread):
113121

@@ -127,7 +135,7 @@ def update_train_status(self, model_name, status):
127135
model["status"] = status
128136

129137
def run(self):
130-
self.trainer = CNNTrainer(self.manager, self.architecture)
138+
self.trainer = CNNTrainer(self.manager, self.architecture, CNNManager.get_model_shape(self.architecture))
131139
self.manager.save_model_status(self.model_name, self.architecture, 0)
132140
image_dir = self.prepare_images()
133141
logging.info("retrain")

‎cnn_models/models.json

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1 +1 @@
1-
{"test_model_1": {"status": 0, "image_height": "128", "image_width": "128", "output_layer": "final_result"}, "mobilenet_v1_1_0_quant": {"status": 100, "image_height": 120, "output_layer": "final_result", "image_width": 160}, "v3-large_224_1.0_uint8": {"status": 1.0, "image_height": 224, "output_layer": "final_result", "image_width": 224}, "object_detection": {"status": 1.0, "image_height": 224, "output_layer": "final_result", "image_width": 224}}
1+
{"test_model_1": {"image_width": "160", "output_layer": "final_result", "status": 1, "image_height": "160"}, "base_high_slow": {"image_width": 224, "status": 1.0, "image_height": 224, "output_layer": "final_result"}, "object_detect": {"image_width": 224, "status": 1.0, "image_height": 224, "output_layer": "final_result"}, "base_low_fast": {"image_width": 224, "status": 100, "image_height": 224, "output_layer": "final_result"}}

0 commit comments

Comments
 (0)