wip: just detect objects in a specific area

This commit is contained in:
blakeblackshear 2019-02-01 06:35:48 -06:00
parent 7e3d2f6611
commit 98ce5a4a59

View File

@ -29,9 +29,9 @@ 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(image_np, sess, detection_graph):
def detect_objects(cropped_frame, full_frame, sess, detection_graph):
# Expand dimensions since the model expects images to have shape: [1, None, None, 3]
image_np_expanded = np.expand_dims(image_np, 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.
@ -52,22 +52,41 @@ def detect_objects(image_np, sess, detection_graph):
objects = []
for index, value in enumerate(classes[0]):
object_dict = {}
if scores[0, index] > 0.5:
if scores[0, index] > 0.1:
object_dict[(category_index.get(value)).get('name').encode('utf8')] = \
scores[0, index]
objects.append(object_dict)
squeezed_boxes = np.squeeze(boxes)
squeezed_scores = np.squeeze(scores)
if(len(objects)>0):
# reposition bounding box based on full frame
for i, box in enumerate(squeezed_boxes):
if squeezed_scores[i] > .1:
ymin = ((box[0] * 300) + 200)/1080 # ymin
xmin = ((box[1] * 300) + 1300)/1920 # xmin
xmax = ((box[2] * 300) + 200)/1080 # ymax
ymax = ((box[3] * 300) + 1300)/1920 # xmax
print("ymin", box[0] * 300, ymin)
print("xmin", box[1] * 300, xmin)
print("ymax", box[2] * 300, ymax)
print("xmax", box[3] * 300, xmax)
# draw boxes for detected objects on image
vis_util.visualize_boxes_and_labels_on_image_array(
image_np,
np.squeeze(boxes),
cropped_frame,
squeezed_boxes,
np.squeeze(classes).astype(np.int32),
np.squeeze(scores),
squeezed_scores,
category_index,
use_normalized_coordinates=True,
line_thickness=4)
line_thickness=4,
min_score_thresh=.1)
# cv2.rectangle(full_frame, (800, 100), (1250, 550), (255,0,0), 2)
return objects, image_np
return objects, cropped_frame
def main():
# capture a single frame and check the frame shape so the correct array
@ -88,18 +107,21 @@ def main():
shared_frame_time = mp.Value('d', 0.0)
# compute the flattened array length from the array shape
flat_array_length = frame_shape[0] * frame_shape[1] * frame_shape[2]
# create shared array for passing the image data from capture to detect_objects
# create shared array for storing the full frame image data
shared_arr = mp.Array(ctypes.c_uint16, flat_array_length)
# create shared array for storing the cropped frame image data
# TODO: make dynamic
shared_cropped_arr = mp.Array(ctypes.c_uint16, 300*300*3)
# create shared array for passing the image data from detect_objects to flask
shared_output_arr = mp.Array(ctypes.c_uint16, flat_array_length)
shared_output_arr = mp.Array(ctypes.c_uint16, 300*300*3)#flat_array_length)
# create a numpy array with the image shape from the shared memory array
# this is used by flask to output an mjpeg stream
frame_output_arr = tonumpyarray(shared_output_arr).reshape(frame_shape)
frame_output_arr = tonumpyarray(shared_output_arr).reshape(300,300,3)
capture_process = mp.Process(target=fetch_frames, args=(shared_arr, shared_frame_time, frame_shape))
capture_process = mp.Process(target=fetch_frames, args=(shared_arr, shared_cropped_arr, shared_frame_time, frame_shape))
capture_process.daemon = True
detection_process = mp.Process(target=process_frames, args=(shared_arr, shared_output_arr, shared_frame_time, frame_shape))
detection_process = mp.Process(target=process_frames, args=(shared_arr, shared_cropped_arr, shared_output_arr, shared_frame_time, frame_shape))
detection_process.daemon = True
capture_process.start()
@ -119,9 +141,9 @@ def main():
# max out at 5 FPS
time.sleep(0.2)
# convert back to BGR
frame_bgr = cv2.cvtColor(frame_output_arr, cv2.COLOR_RGB2BGR)
# frame_bgr = cv2.cvtColor(frame_output_arr, cv2.COLOR_RGB2BGR)
# encode the image into a jpg
ret, jpg = cv2.imencode('.jpg', frame_bgr)
ret, jpg = cv2.imencode('.jpg', frame_output_arr)
yield (b'--frame\r\n'
b'Content-Type: image/jpeg\r\n\r\n' + jpg.tobytes() + b'\r\n\r\n')
@ -136,9 +158,10 @@ 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_frame_time, frame_shape):
def fetch_frames(shared_arr, shared_cropped_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(300,300,3)
# start the video capture
video = cv2.VideoCapture(RTSP_URL)
@ -158,6 +181,12 @@ def fetch_frames(shared_arr, shared_frame_time, frame_shape):
ret, frame = video.retrieve()
if ret:
# copy the frame into the numpy array
# Position 1
# cropped_frame[:] = frame[270:720, 550:1000]
# Position 2
# frame_cropped = frame[270:720, 100:550]
# Car
cropped_frame[:] = frame[200:500, 1300:1600]
arr[:] = frame
# signal to the detection_process by setting the shared_frame_time
shared_frame_time.value = frame_time.timestamp()
@ -165,11 +194,12 @@ def fetch_frames(shared_arr, shared_frame_time, frame_shape):
video.release()
# do the actual object detection
def process_frames(shared_arr, shared_output_arr, shared_frame_time, frame_shape):
def process_frames(shared_arr, shared_cropped_arr, shared_output_arr, shared_frame_time, frame_shape):
# shape shared input array into frame for processing
arr = tonumpyarray(shared_arr).reshape(frame_shape)
shared_cropped_frame = tonumpyarray(shared_cropped_arr).reshape(300,300,3)
# shape shared output array into frame so it can be copied into
output_arr = tonumpyarray(shared_output_arr).reshape(frame_shape)
output_arr = tonumpyarray(shared_output_arr).reshape(300,300,3)
# Load a (frozen) Tensorflow model into memory before the processing loop
detection_graph = tf.Graph()
@ -211,14 +241,15 @@ def process_frames(shared_arr, shared_output_arr, shared_frame_time, frame_shape
# make a copy of the frame
frame = arr.copy()
cropped_frame = shared_cropped_frame.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
# convert to RGB
frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
cropped_frame_rgb = cv2.cvtColor(cropped_frame, cv2.COLOR_BGR2RGB)
# do the object detection
objects, frame_overlay = detect_objects(frame_rgb, sess, detection_graph)
objects, frame_overlay = detect_objects(cropped_frame_rgb, frame, sess, detection_graph)
# copy the output frame with the bounding boxes to the output array
output_arr[:] = frame_overlay
if(len(objects) > 0):