|
@@ -1,9 +1,8 @@
|
|
import datetime
|
|
import datetime
|
|
import cv2
|
|
import cv2
|
|
import numpy as np
|
|
import numpy as np
|
|
-import tensorflow as tf
|
|
|
|
-from object_detection.utils import label_map_util
|
|
|
|
-from object_detection.utils import visualization_utils as vis_util
|
|
|
|
|
|
+from edgetpu.detection.engine import DetectionEngine
|
|
|
|
+from PIL import Image
|
|
from . util import tonumpyarray
|
|
from . util import tonumpyarray
|
|
|
|
|
|
# TODO: make dynamic?
|
|
# TODO: make dynamic?
|
|
@@ -13,58 +12,38 @@ PATH_TO_CKPT = '/frozen_inference_graph.pb'
|
|
# List of the strings that is used to add correct label for each box.
|
|
# List of the strings that is used to add correct label for each box.
|
|
PATH_TO_LABELS = '/label_map.pbtext'
|
|
PATH_TO_LABELS = '/label_map.pbtext'
|
|
|
|
|
|
-# Loading label map
|
|
|
|
-label_map = label_map_util.load_labelmap(PATH_TO_LABELS)
|
|
|
|
-categories = label_map_util.convert_label_map_to_categories(label_map, max_num_classes=NUM_CLASSES,
|
|
|
|
- use_display_name=True)
|
|
|
|
-category_index = label_map_util.create_category_index(categories)
|
|
|
|
|
|
+# Function to read labels from text files.
|
|
|
|
+def ReadLabelFile(file_path):
|
|
|
|
+ with open(file_path, 'r') as f:
|
|
|
|
+ lines = f.readlines()
|
|
|
|
+ ret = {}
|
|
|
|
+ for line in lines:
|
|
|
|
+ pair = line.strip().split(maxsplit=1)
|
|
|
|
+ ret[int(pair[0])] = pair[1].strip()
|
|
|
|
+ return ret
|
|
|
|
|
|
# do the actual object detection
|
|
# do the actual object detection
|
|
-def tf_detect_objects(cropped_frame, sess, detection_graph, region_size, region_x_offset, region_y_offset, debug):
|
|
|
|
|
|
+def tf_detect_objects(cropped_frame, engine, labels, region_size, region_x_offset, region_y_offset, debug):
|
|
|
|
+ # Resize to 300x300
|
|
|
|
+ cropped_frame = cv2.resize(cropped_frame, dsize=(300, 300), interpolation=cv2.INTER_LINEAR)
|
|
# Expand dimensions since the model expects images to have shape: [1, None, None, 3]
|
|
# Expand dimensions since the model expects images to have shape: [1, None, None, 3]
|
|
image_np_expanded = np.expand_dims(cropped_frame, axis=0)
|
|
image_np_expanded = np.expand_dims(cropped_frame, axis=0)
|
|
- image_tensor = detection_graph.get_tensor_by_name('image_tensor:0')
|
|
|
|
-
|
|
|
|
- # Each box represents a part of the image where a particular object was detected.
|
|
|
|
- boxes = detection_graph.get_tensor_by_name('detection_boxes:0')
|
|
|
|
-
|
|
|
|
- # Each score represent how level of confidence for each of the objects.
|
|
|
|
- # Score is shown on the result image, together with the class label.
|
|
|
|
- scores = detection_graph.get_tensor_by_name('detection_scores:0')
|
|
|
|
- classes = detection_graph.get_tensor_by_name('detection_classes:0')
|
|
|
|
- num_detections = detection_graph.get_tensor_by_name('num_detections:0')
|
|
|
|
|
|
|
|
# Actual detection.
|
|
# Actual detection.
|
|
- (boxes, scores, classes, num_detections) = sess.run(
|
|
|
|
- [boxes, scores, classes, num_detections],
|
|
|
|
- feed_dict={image_tensor: image_np_expanded})
|
|
|
|
-
|
|
|
|
- if debug:
|
|
|
|
- if len([value for index,value in enumerate(classes[0]) if str(category_index.get(value).get('name')) == 'person' and scores[0,index] > 0.5]) > 0:
|
|
|
|
- vis_util.visualize_boxes_and_labels_on_image_array(
|
|
|
|
- cropped_frame,
|
|
|
|
- np.squeeze(boxes),
|
|
|
|
- np.squeeze(classes).astype(np.int32),
|
|
|
|
- np.squeeze(scores),
|
|
|
|
- category_index,
|
|
|
|
- use_normalized_coordinates=True,
|
|
|
|
- line_thickness=4)
|
|
|
|
- cv2.imwrite("/lab/debug/obj-{}-{}-{}.jpg".format(region_x_offset, region_y_offset, datetime.datetime.now().timestamp()), cropped_frame)
|
|
|
|
-
|
|
|
|
|
|
+ ans = engine.DetectWithInputTensor(image_np_expanded.flatten(), threshold=0.5, top_k=3)
|
|
|
|
|
|
# build an array of detected objects
|
|
# build an array of detected objects
|
|
objects = []
|
|
objects = []
|
|
- for index, value in enumerate(classes[0]):
|
|
|
|
- score = scores[0, index]
|
|
|
|
- if score > 0.5:
|
|
|
|
- box = boxes[0, index].tolist()
|
|
|
|
|
|
+ if ans:
|
|
|
|
+ for obj in ans:
|
|
|
|
+ box = obj.bounding_box.flatten().tolist()
|
|
objects.append({
|
|
objects.append({
|
|
- 'name': str(category_index.get(value).get('name')),
|
|
|
|
- 'score': float(score),
|
|
|
|
- 'ymin': int((box[0] * region_size) + region_y_offset),
|
|
|
|
- 'xmin': int((box[1] * region_size) + region_x_offset),
|
|
|
|
- 'ymax': int((box[2] * region_size) + region_y_offset),
|
|
|
|
- 'xmax': int((box[3] * region_size) + region_x_offset)
|
|
|
|
|
|
+ 'name': str(labels[obj.label_id]),
|
|
|
|
+ 'score': float(obj.score),
|
|
|
|
+ 'xmin': int((box[0] * region_size) + region_x_offset),
|
|
|
|
+ 'ymin': int((box[1] * region_size) + region_y_offset),
|
|
|
|
+ 'xmax': int((box[2] * region_size) + region_x_offset),
|
|
|
|
+ 'ymax': int((box[3] * region_size) + region_y_offset)
|
|
})
|
|
})
|
|
|
|
|
|
return objects
|
|
return objects
|
|
@@ -75,15 +54,9 @@ def detect_objects(shared_arr, object_queue, shared_frame_time, frame_lock, fram
|
|
# shape shared input array into frame for processing
|
|
# shape shared input array into frame for processing
|
|
arr = tonumpyarray(shared_arr).reshape(frame_shape)
|
|
arr = tonumpyarray(shared_arr).reshape(frame_shape)
|
|
|
|
|
|
- # Load a (frozen) Tensorflow model into memory before the processing loop
|
|
|
|
- detection_graph = tf.Graph()
|
|
|
|
- with detection_graph.as_default():
|
|
|
|
- od_graph_def = tf.GraphDef()
|
|
|
|
- with tf.gfile.GFile(PATH_TO_CKPT, 'rb') as fid:
|
|
|
|
- serialized_graph = fid.read()
|
|
|
|
- od_graph_def.ParseFromString(serialized_graph)
|
|
|
|
- tf.import_graph_def(od_graph_def, name='')
|
|
|
|
- sess = tf.Session(graph=detection_graph)
|
|
|
|
|
|
+ # Load the edgetpu engine and labels
|
|
|
|
+ engine = DetectionEngine(PATH_TO_CKPT)
|
|
|
|
+ labels = ReadLabelFile(PATH_TO_LABELS)
|
|
|
|
|
|
frame_time = 0.0
|
|
frame_time = 0.0
|
|
while True:
|
|
while True:
|
|
@@ -105,7 +78,7 @@ def detect_objects(shared_arr, object_queue, shared_frame_time, frame_lock, fram
|
|
# convert to RGB
|
|
# convert to RGB
|
|
cropped_frame_rgb = cv2.cvtColor(cropped_frame, cv2.COLOR_BGR2RGB)
|
|
cropped_frame_rgb = cv2.cvtColor(cropped_frame, cv2.COLOR_BGR2RGB)
|
|
# do the object detection
|
|
# do the object detection
|
|
- objects = tf_detect_objects(cropped_frame_rgb, sess, detection_graph, region_size, region_x_offset, region_y_offset, debug)
|
|
|
|
|
|
+ objects = tf_detect_objects(cropped_frame_rgb, engine, labels, region_size, region_x_offset, region_y_offset, debug)
|
|
for obj in objects:
|
|
for obj in objects:
|
|
# ignore persons below the size threshold
|
|
# ignore persons below the size threshold
|
|
if obj['name'] == 'person' and (obj['xmax']-obj['xmin'])*(obj['ymax']-obj['ymin']) < min_person_area:
|
|
if obj['name'] == 'person' and (obj['xmax']-obj['xmin'])*(obj['ymax']-obj['ymin']) < min_person_area:
|