switch to a thread for object detection

This commit is contained in:
blakeblackshear
2019-03-27 20:44:57 -05:00
parent a074945394
commit 0514eeac03
2 changed files with 29 additions and 109 deletions

View File

@@ -21,89 +21,40 @@ def ReadLabelFile(file_path):
ret[int(pair[0])] = pair[1].strip()
return ret
def detect_objects(prepped_frame_array, prepped_frame_time,
prepped_frame_ready, prepped_frame_grabbed,
prepped_frame_box, object_queue, debug):
prepped_frame_np = tonumpyarray(prepped_frame_array)
# Load the edgetpu engine and labels
engine = DetectionEngine(PATH_TO_CKPT)
labels = ReadLabelFile(PATH_TO_LABELS)
frame_time = 0.0
region_box = [0,0,0]
while True:
# wait until a frame is ready
prepped_frame_ready.wait()
prepped_frame_copy = prepped_frame_np.copy()
frame_time = prepped_frame_time.value
region_box[:] = prepped_frame_box
prepped_frame_grabbed.set()
# print("Grabbed " + str(region_box[1]) + "," + str(region_box[2]))
# Actual detection.
objects = engine.DetectWithInputTensor(prepped_frame_copy, threshold=0.5, top_k=3)
# time.sleep(0.1)
# objects = []
# print(engine.get_inference_time())
# put detected objects in the queue
if objects:
for obj in objects:
box = obj.bounding_box.flatten().tolist()
object_queue.put({
'frame_time': frame_time,
'name': str(labels[obj.label_id]),
'score': float(obj.score),
'xmin': int((box[0] * region_box[0]) + region_box[1]),
'ymin': int((box[1] * region_box[0]) + region_box[2]),
'xmax': int((box[2] * region_box[0]) + region_box[1]),
'ymax': int((box[3] * region_box[0]) + region_box[2])
})
# else:
# object_queue.put({
# 'frame_time': frame_time,
# 'name': 'dummy',
# 'score': 0.99,
# 'xmin': int(0 + region_box[1]),
# 'ymin': int(0 + region_box[2]),
# 'xmax': int(10 + region_box[1]),
# 'ymax': int(10 + region_box[2])
# })
class PreppedQueueProcessor(threading.Thread):
def __init__(self, prepped_frame_array,
prepped_frame_time,
prepped_frame_ready,
prepped_frame_grabbed,
prepped_frame_box,
prepped_frame_queue):
def __init__(self, prepped_frame_queue, object_queue):
threading.Thread.__init__(self)
self.prepped_frame_array = prepped_frame_array
self.prepped_frame_time = prepped_frame_time
self.prepped_frame_ready = prepped_frame_ready
self.prepped_frame_grabbed = prepped_frame_grabbed
self.prepped_frame_box = prepped_frame_box
self.prepped_frame_queue = prepped_frame_queue
self.object_queue = object_queue
# Load the edgetpu engine and labels
self.engine = DetectionEngine(PATH_TO_CKPT)
self.labels = ReadLabelFile(PATH_TO_LABELS)
def run(self):
prepped_frame_np = tonumpyarray(self.prepped_frame_array)
# process queue...
while True:
frame = self.prepped_frame_queue.get()
# print(self.prepped_frame_queue.qsize())
prepped_frame_np[:] = frame['frame']
self.prepped_frame_time.value = frame['frame_time']
self.prepped_frame_box[0] = frame['region_size']
self.prepped_frame_box[1] = frame['region_x_offset']
self.prepped_frame_box[2] = frame['region_y_offset']
# print("Passed " + str(frame['region_x_offset']) + "," + str(frame['region_x_offset']))
self.prepped_frame_ready.set()
self.prepped_frame_grabbed.wait()
self.prepped_frame_grabbed.clear()
self.prepped_frame_ready.clear()
# Actual detection.
objects = self.engine.DetectWithInputTensor(frame['frame'], threshold=0.5, top_k=3)
# time.sleep(0.1)
# objects = []
# print(engine.get_inference_time())
# put detected objects in the queue
if objects:
for obj in objects:
box = obj.bounding_box.flatten().tolist()
self.object_queue.put({
'frame_time': frame['frame_time'],
'name': str(self.labels[obj.label_id]),
'score': float(obj.score),
'xmin': int((box[0] * frame['region_size']) + frame['region_x_offset']),
'ymin': int((box[1] * frame['region_size']) + frame['region_y_offset']),
'xmax': int((box[2] * frame['region_size']) + frame['region_x_offset']),
'ymax': int((box[3] * frame['region_size']) + frame['region_y_offset'])
})
# should this be a region class?
@@ -156,5 +107,5 @@ class FramePrepper(threading.Thread):
'region_x_offset': self.region_x_offset,
'region_y_offset': self.region_y_offset
})
# else:
# print("queue full. moving on")
else:
print("queue full. moving on")