def process_frame(frame: VideoFrame, threshold: float = DETECT_THRESHOLD) -> bool:
width = frame.video_info().width
height = frame.video_info().height
for tensor in frame.tensors():
dims = tensor.dims()
data = tensor.data()
object_size = dims[-1]
for i in range(dims[-2]):
image_id = data[i * object_size + 0]
label_id = data[i * object_size + 1]
confidence = data[i * object_size + 2]
x_min = int(data[i * object_size + 3] * width + 0.5)
y_min = int(data[i * object_size + 4] * height + 0.5)
x_max = int(data[i * object_size + 5] * width + 0.5)
y_max = int(data[i * object_size + 6] * height + 0.5)
if image_id != 0:
break
if confidence < threshold:
continue
roi = frame.add_region(x_min, y_min, x_max - x_min, y_max - y_min, str(label_id), confidence)
return True
def process_frame(frame: VideoFrame,
threshold: float = DETECT_THRESHOLD) -> bool:
global input_height, input_width, prob_threshold, iou_threshold
width = frame.video_info().width
height = frame.video_info().height
objects = list()
for tensor in frame.tensors():
dims = tensor.dims()
data = tensor.data()
layer_params = YoloParams(dims[2])
data = data.reshape(dims[0], dims[1], dims[2], dims[3])
objects += parse_yolo_region(data, (input_height, input_width),
(height, width),
layer_params,
prob_threshold,
is_proportional=False)
objects = filter_objects(objects, iou_threshold, prob_threshold)
with frame.data() as frame:
for obj in objects:
# Validation bbox of detected object
obj['xmax'] = min(obj['xmax'], width)
obj['ymax'] = min(obj['ymax'], height)
obj['xmin'] = max(obj['xmin'], 0)
obj['ymin'] = max(obj['ymin'], 0)
color = (min(obj['class_id'] * 12.5,
255), min(obj['class_id'] * 7,
255), min(obj['class_id'] * 5, 255))
det_label = labels_map[obj['class_id']] if labels_map and len(
labels_map) >= obj['class_id'] else str(obj['class_id'])
cv2.rectangle(frame, (obj['xmin'], obj['ymin']),
(obj['xmax'], obj['ymax']), color, 2)
cv2.putText(
frame, "#" + det_label + ' ' +
str(round(obj['confidence'] * 100, 1)) + ' %',
(obj['xmin'], obj['ymin'] - 7), cv2.FONT_HERSHEY_COMPLEX, 0.6,
color, 1)
return True