forked from Github/frigate
first working version, single region and motion detection disabled
This commit is contained in:
@@ -34,7 +34,7 @@ def detect_motion(shared_arr, shared_frame_time, frame_lock, frame_ready, motion
|
||||
|
||||
# lock and make a copy of the cropped frame
|
||||
with frame_lock:
|
||||
cropped_frame = arr[region_y_offset:region_y_offset+region_size, region_x_offset:region_x_offset+region_size].copy().astype('uint8')
|
||||
cropped_frame = arr[region_y_offset:region_y_offset+region_size, region_x_offset:region_x_offset+region_size].copy()
|
||||
frame_time = shared_frame_time.value
|
||||
|
||||
# convert to grayscale
|
||||
|
||||
@@ -1,9 +1,8 @@
|
||||
import datetime
|
||||
import cv2
|
||||
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
|
||||
|
||||
# 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.
|
||||
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
|
||||
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]
|
||||
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.
|
||||
(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
|
||||
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({
|
||||
'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
|
||||
@@ -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
|
||||
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
|
||||
while True:
|
||||
@@ -105,7 +78,7 @@ def detect_objects(shared_arr, object_queue, shared_frame_time, frame_lock, fram
|
||||
# convert to RGB
|
||||
cropped_frame_rgb = cv2.cvtColor(cropped_frame, cv2.COLOR_BGR2RGB)
|
||||
# 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:
|
||||
# ignore persons below the size threshold
|
||||
if obj['name'] == 'person' and (obj['xmax']-obj['xmin'])*(obj['ymax']-obj['ymin']) < min_person_area:
|
||||
|
||||
@@ -2,4 +2,4 @@ import numpy as np
|
||||
|
||||
# convert shared memory array into numpy array
|
||||
def tonumpyarray(mp_arr):
|
||||
return np.frombuffer(mp_arr.get_obj(), dtype=np.uint16)
|
||||
return np.frombuffer(mp_arr.get_obj(), dtype=np.uint8)
|
||||
@@ -78,7 +78,7 @@ class FrameTracker(threading.Thread):
|
||||
|
||||
# lock and make a copy of the frame
|
||||
with self.frame_lock:
|
||||
frame = self.shared_frame.copy().astype('uint8')
|
||||
frame = self.shared_frame.copy()
|
||||
frame_time = self.frame_time.value
|
||||
|
||||
# add the frame to recent frames
|
||||
|
||||
Reference in New Issue
Block a user