Skip to content
Snippets Groups Projects
Commit 9951cfe0 authored by hannandarryl's avatar hannandarryl
Browse files

added tflite code back

parent 6f56e879
No related branches found
No related tags found
No related merge requests found
......@@ -45,7 +45,7 @@ class YoloModel():
def get_bounding_boxes(self, original_image):
start = time.time()
print('image resizing\n')
# print('image resizing\n')
image_data = cv2.resize(original_image, (self.input_size, self.input_size))
image_data = image_data / 255.
images_data = []
......@@ -53,28 +53,42 @@ class YoloModel():
images_data.append(image_data)
images_data = np.asarray(images_data).astype(np.float32)
print('running as tensorflow\n')
# print('running as tensorflow\n')
#print('loading model\n')
# print(FLAGS.weights)
# saved_model_loaded = tf.saved_model.load(FLAGS.weights, tags=[tag_constants.SERVING])
#print('model loaded\n')
if FLAGS.framework == 'tflite':
# print('running as tflite\n')
interpreter = tf.lite.Interpreter(model_path=FLAGS.weights)
interpreter.allocate_tensors()
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()
interpreter.set_tensor(input_details[0]['index'], images_data)
interpreter.invoke()
pred = [interpreter.get_tensor(output_details[i]['index']) for i in range(len(output_details))]
if FLAGS.model == 'yolov3' and FLAGS.tiny == True:
boxes, pred_conf = filter_boxes(pred[1], pred[0], score_threshold=0.25, input_shape=tf.constant([input_size, input_size]))
else:
boxes, pred_conf = filter_boxes(pred[0], pred[1], score_threshold=0.25, input_shape=tf.constant([input_size, input_size]))
else:
infer = self.saved_model_loaded.signatures['serving_default']
print('batch data\n')
# print('batch data\n')
batch_data = tf.constant(images_data)
print('computing bounding box data\n')
# print('computing bounding box data\n')
yolo_start_time = time.time()
pred_bbox = infer(batch_data)
for key, value in pred_bbox.items():
boxes = value[:, :, 0:4]
pred_conf = value[:, :, 4:]
print("VALUE", value)
# print("VALUE", value)
yolo_end_time = time.time()
yolo_elapsed_time = yolo_end_time - yolo_start_time
print('Took %.2f seconds to run yolo\n' % (yolo_elapsed_time))
# print('Took %.2f seconds to run yolo\n' % (yolo_elapsed_time))
print('non max suppression\n')
# print('non max suppression\n')
boxes, scores, classes, valid_detections = tf.image.combined_non_max_suppression(
boxes=tf.reshape(boxes, (tf.shape(boxes)[0], -1, 1, 4)),
scores=tf.reshape(
......@@ -85,7 +99,7 @@ class YoloModel():
score_threshold=0.25
)
print('formatting bounding box data\n')
# print('formatting bounding box data\n')
boxes = boxes.numpy()
# remove bounding boxes with zero area
......@@ -103,5 +117,5 @@ class YoloModel():
end = time.time()
elapsed_time = end - start
print('Took %.2f seconds to run whole bounding box function\n' % (elapsed_time))
# print('Took %.2f seconds to run whole bounding box function\n' % (elapsed_time))
return boxes
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment