mirror of
https://github.com/blakeblackshear/frigate.git
synced 2024-11-21 19:07:46 +01:00
crop the frame and calculate the coordinates in the subprocess and add labels to the image
This commit is contained in:
blakeblackshear
parent
11af9bb953
commit
b91c24bf8f
@ -23,9 +23,9 @@ PATH_TO_LABELS = '/label_map.pbtext'
|
||||
# TODO: make dynamic?
|
||||
NUM_CLASSES = 90
|
||||
|
||||
REGION_SIZE = 700
|
||||
REGION_X_OFFSET = 950
|
||||
REGION_Y_OFFSET = 380
|
||||
REGION_SIZE = 300
|
||||
REGION_X_OFFSET = 1250
|
||||
REGION_Y_OFFSET = 180
|
||||
|
||||
# Loading label map
|
||||
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)
|
||||
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]
|
||||
image_np_expanded = np.expand_dims(cropped_frame, axis=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]):
|
||||
score = scores[0, index]
|
||||
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
|
||||
|
||||
@ -84,16 +92,13 @@ def main():
|
||||
shared_arr = mp.Array(ctypes.c_uint16, flat_array_length)
|
||||
# shape current frame so it can be treated as an image
|
||||
frame_arr = tonumpyarray(shared_arr).reshape(frame_shape)
|
||||
# create shared array for storing the cropped frame image data
|
||||
# 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
|
||||
# create shared array for storing 10 detected objects
|
||||
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
|
||||
|
||||
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
|
||||
|
||||
capture_process.start()
|
||||
@ -113,21 +118,33 @@ def main():
|
||||
# max out at 5 FPS
|
||||
time.sleep(0.2)
|
||||
frame = frame_arr.copy()
|
||||
# convert to RGB for drawing
|
||||
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
|
||||
# draw the bounding boxes on the screen
|
||||
object_index = 0
|
||||
while(object_index < 60 and shared_output_arr[object_index] > 0):
|
||||
object_class = shared_output_arr[object_index]
|
||||
object_name = str(category_index.get(object_class).get('name'))
|
||||
score = shared_output_arr[object_index+1]
|
||||
ymin = int(((shared_output_arr[object_index+2] * REGION_SIZE) + REGION_Y_OFFSET))
|
||||
xmin = int(((shared_output_arr[object_index+3] * REGION_SIZE) + REGION_X_OFFSET))
|
||||
ymax = int(((shared_output_arr[object_index+4] * REGION_SIZE) + REGION_Y_OFFSET))
|
||||
xmax = int(((shared_output_arr[object_index+5] * REGION_SIZE) + REGION_X_OFFSET))
|
||||
cv2.rectangle(frame, (xmin, ymin), (xmax, ymax), (255,0,0), 2)
|
||||
display_str = '{}: {}%'.format(object_name, int(100*score))
|
||||
ymin = int(shared_output_arr[object_index+2])
|
||||
xmin = int(shared_output_arr[object_index+3])
|
||||
ymax = int(shared_output_arr[object_index+4])
|
||||
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
|
||||
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)
|
||||
# convert back to BGR
|
||||
frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
|
||||
# encode the image into a jpg
|
||||
ret, jpg = cv2.imencode('.jpg', frame)
|
||||
yield (b'--frame\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
|
||||
# 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
|
||||
arr = tonumpyarray(shared_arr).reshape(frame_shape)
|
||||
cropped_frame = tonumpyarray(shared_cropped_arr).reshape(REGION_SIZE,REGION_SIZE,3)
|
||||
|
||||
# start the video capture
|
||||
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]
|
||||
# Position 2
|
||||
# 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
|
||||
# signal to the detection_process by setting the shared_frame_time
|
||||
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()
|
||||
|
||||
# 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
|
||||
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
|
||||
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)
|
||||
continue
|
||||
|
||||
# make a copy of the frame
|
||||
# frame = arr.copy()
|
||||
cropped_frame = shared_cropped_frame.copy()
|
||||
# make a copy of the cropped frame
|
||||
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
|
||||
# signal that the frame has been used so a new one will be ready
|
||||
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
|
||||
cropped_frame_rgb = cv2.cvtColor(cropped_frame, cv2.COLOR_BGR2RGB)
|
||||
# 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
|
||||
shared_output_arr[:] = objects + [0.0] * (60-len(objects))
|
||||
|
||||
|
||||
Loading…
Reference in New Issue
Block a user