Mirrors/blakeblackshear.frigate
1
0
Fork 0
mirror of https://github.com/blakeblackshear/frigate.git synced 2024-11-21 19:07:46 +01:00
Code Issues Projects Releases Wiki Activity

crop the frame and calculate the coordinates in the subprocess and add labels to the image

Browse Source
This commit is contained in:
blakeblackshear 2019-02-02 08:16:35 -06:00
parent 11af9bb953
commit b91c24bf8f
1 changed files with 41 additions and 29 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

70
detect_objects.py
View File

@ -23,9 +23,9 @@ PATH_TO_LABELS = '/label_map.pbtext'
# TODO: make dynamic? # TODO: make dynamic?
NUM_CLASSES = 90 NUM_CLASSES = 90
REGION_SIZE = 700 REGION_SIZE = 300
REGION_X_OFFSET = 950 REGION_X_OFFSET = 1250
REGION_Y_OFFSET = 380 REGION_Y_OFFSET = 180
# Loading label map # Loading label map
label_map = label_map_util.load_labelmap(PATH_TO_LABELS) label_map = label_map_util.load_labelmap(PATH_TO_LABELS)
@ -33,7 +33,7 @@ categories = label_map_util.convert_label_map_to_categories(label_map, max_num_c
use_display_name=True) use_display_name=True)
category_index = label_map_util.create_category_index(categories) category_index = label_map_util.create_category_index(categories)
def detect_objects(cropped_frame, sess, detection_graph): def detect_objects(cropped_frame, sess, detection_graph, region_size, region_x_offset, region_y_offset):
# 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') image_tensor = detection_graph.get_tensor_by_name('image_tensor:0')
@ -57,7 +57,15 @@ def detect_objects(cropped_frame, sess, detection_graph):
for index, value in enumerate(classes[0]): for index, value in enumerate(classes[0]):
score = scores[0, index] score = scores[0, index]
if score > 0.1: if score > 0.1:
objects += [value, scores[0, index]] + boxes[0, index].tolist() box = boxes[0, index].tolist()
box[0] = (box[0] * region_size) + region_y_offset
box[1] = (box[1] * region_size) + region_x_offset
box[2] = (box[2] * region_size) + region_y_offset
box[3] = (box[3] * region_size) + region_x_offset
objects += [value, scores[0, index]] + box
# only get the first 10 objects
if len(objects) = 60:
break
return objects return objects
@ -84,16 +92,13 @@ def main():
shared_arr = mp.Array(ctypes.c_uint16, flat_array_length) shared_arr = mp.Array(ctypes.c_uint16, flat_array_length)
# shape current frame so it can be treated as an image # shape current frame so it can be treated as an image
frame_arr = tonumpyarray(shared_arr).reshape(frame_shape) frame_arr = tonumpyarray(shared_arr).reshape(frame_shape)
# create shared array for storing the cropped frame image data # create shared array for storing 10 detected objects
# TODO: make dynamic
shared_cropped_arr = mp.Array(ctypes.c_uint16, REGION_SIZE*REGION_SIZE*3)
# create shared array for passing the image data from detect_objects to flask
shared_output_arr = mp.Array(ctypes.c_double, 6*10) shared_output_arr = mp.Array(ctypes.c_double, 6*10)
capture_process = mp.Process(target=fetch_frames, args=(shared_arr, shared_cropped_arr, shared_frame_time, frame_shape)) capture_process = mp.Process(target=fetch_frames, args=(shared_arr, shared_frame_time, frame_shape))
capture_process.daemon = True capture_process.daemon = True
detection_process = mp.Process(target=process_frames, args=(shared_arr, shared_cropped_arr, shared_output_arr, shared_frame_time, frame_shape)) detection_process = mp.Process(target=process_frames, args=(shared_arr, shared_output_arr, shared_frame_time, frame_shape, REGION_SIZE, REGION_X_OFFSET, REGION_Y_OFFSET))
detection_process.daemon = True detection_process.daemon = True
capture_process.start() capture_process.start()
@ -113,21 +118,33 @@ def main():
# max out at 5 FPS # max out at 5 FPS
time.sleep(0.2) time.sleep(0.2)
frame = frame_arr.copy() frame = frame_arr.copy()
# convert to RGB for drawing
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
# draw the bounding boxes on the screen # draw the bounding boxes on the screen
object_index = 0 object_index = 0
while(object_index < 60 and shared_output_arr[object_index] > 0): while(object_index < 60 and shared_output_arr[object_index] > 0):
object_class = shared_output_arr[object_index] object_class = shared_output_arr[object_index]
object_name = str(category_index.get(object_class).get('name'))
score = shared_output_arr[object_index+1] score = shared_output_arr[object_index+1]
ymin = int(((shared_output_arr[object_index+2] * REGION_SIZE) + REGION_Y_OFFSET)) display_str = '{}: {}%'.format(object_name, int(100*score))
xmin = int(((shared_output_arr[object_index+3] * REGION_SIZE) + REGION_X_OFFSET)) ymin = int(shared_output_arr[object_index+2])
ymax = int(((shared_output_arr[object_index+4] * REGION_SIZE) + REGION_Y_OFFSET)) xmin = int(shared_output_arr[object_index+3])
xmax = int(((shared_output_arr[object_index+5] * REGION_SIZE) + REGION_X_OFFSET)) ymax = int(shared_output_arr[object_index+4])
cv2.rectangle(frame, (xmin, ymin), (xmax, ymax), (255,0,0), 2) xmax = int(shared_output_arr[object_index+5])
vis_util.draw_bounding_box_on_image_array(frame,
ymin,
xmin,
ymax,
xmax,
color='red',
thickness=2,
display_str_list=[display_str],
use_normalized_coordinates=False)
object_index += 6 object_index += 6
print(category_index.get(object_class).get('name').encode('utf8'), score)
# encode the image into a jpg
cv2.rectangle(frame, (REGION_X_OFFSET, REGION_Y_OFFSET), (REGION_X_OFFSET+REGION_SIZE, REGION_Y_OFFSET+REGION_SIZE), (255,255,255), 2) cv2.rectangle(frame, (REGION_X_OFFSET, REGION_Y_OFFSET), (REGION_X_OFFSET+REGION_SIZE, REGION_Y_OFFSET+REGION_SIZE), (255,255,255), 2)
# convert back to BGR
frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
# encode the image into a jpg
ret, jpg = cv2.imencode('.jpg', frame) ret, jpg = cv2.imencode('.jpg', frame)
yield (b'--frame\r\n' yield (b'--frame\r\n'
b'Content-Type: image/jpeg\r\n\r\n' + jpg.tobytes() + b'\r\n\r\n') b'Content-Type: image/jpeg\r\n\r\n' + jpg.tobytes() + b'\r\n\r\n')
@ -143,10 +160,9 @@ def tonumpyarray(mp_arr):
# fetch the frames as fast a possible, only decoding the frames when the # fetch the frames as fast a possible, only decoding the frames when the
# detection_process has consumed the current frame # detection_process has consumed the current frame
def fetch_frames(shared_arr, shared_cropped_arr, shared_frame_time, frame_shape): def fetch_frames(shared_arr, shared_frame_time, frame_shape):
# convert shared memory array into numpy and shape into image array # convert shared memory array into numpy and shape into image array
arr = tonumpyarray(shared_arr).reshape(frame_shape) arr = tonumpyarray(shared_arr).reshape(frame_shape)
cropped_frame = tonumpyarray(shared_cropped_arr).reshape(REGION_SIZE,REGION_SIZE,3)
# start the video capture # start the video capture
video = cv2.VideoCapture(RTSP_URL) video = cv2.VideoCapture(RTSP_URL)
@ -170,8 +186,6 @@ def fetch_frames(shared_arr, shared_cropped_arr, shared_frame_time, frame_shape)
# cropped_frame[:] = frame[270:720, 550:1000] # cropped_frame[:] = frame[270:720, 550:1000]
# Position 2 # Position 2
# frame_cropped = frame[270:720, 100:550] # frame_cropped = frame[270:720, 100:550]
# Car
cropped_frame[:] = frame[REGION_Y_OFFSET:REGION_Y_OFFSET+REGION_SIZE, REGION_X_OFFSET:REGION_X_OFFSET+REGION_SIZE]
arr[:] = frame arr[:] = frame
# signal to the detection_process by setting the shared_frame_time # signal to the detection_process by setting the shared_frame_time
shared_frame_time.value = frame_time.timestamp() shared_frame_time.value = frame_time.timestamp()
@ -179,10 +193,9 @@ def fetch_frames(shared_arr, shared_cropped_arr, shared_frame_time, frame_shape)
video.release() video.release()
# do the actual object detection # do the actual object detection
def process_frames(shared_arr, shared_cropped_arr, shared_output_arr, shared_frame_time, frame_shape): def process_frames(shared_arr, shared_output_arr, shared_frame_time, frame_shape, region_size, region_x_offset, region_y_offset):
# 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)
shared_cropped_frame = tonumpyarray(shared_cropped_arr).reshape(REGION_SIZE,REGION_SIZE,3)
# Load a (frozen) Tensorflow model into memory before the processing loop # Load a (frozen) Tensorflow model into memory before the processing loop
detection_graph = tf.Graph() detection_graph = tf.Graph()
@ -222,9 +235,8 @@ def process_frames(shared_arr, shared_cropped_arr, shared_output_arr, shared_fra
time.sleep(0.01) time.sleep(0.01)
continue continue
# make a copy of the frame # make a copy of the cropped frame
# frame = arr.copy() cropped_frame = arr[region_y_offset:region_y_offset+region_size, region_x_offset:region_x_offset+region_size].copy()
cropped_frame = shared_cropped_frame.copy()
frame_time = shared_frame_time.value frame_time = shared_frame_time.value
# signal that the frame has been used so a new one will be ready # signal that the frame has been used so a new one will be ready
shared_frame_time.value = 0.0 shared_frame_time.value = 0.0
@ -232,7 +244,7 @@ def process_frames(shared_arr, shared_cropped_arr, shared_output_arr, shared_fra
# 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 = detect_objects(cropped_frame_rgb, sess, detection_graph) objects = detect_objects(cropped_frame_rgb, sess, detection_graph, region_size, region_x_offset, region_y_offset)
# copy the detected objects to the output array, filling the array when needed # copy the detected objects to the output array, filling the array when needed
shared_output_arr[:] = objects + [0.0] * (60-len(objects)) shared_output_arr[:] = objects + [0.0] * (60-len(objects))