我目前正在做一个项目,我正在使用Roboflow模型来检测视频中的台球。视频源使用OpenCV进行处理,模型在OpenAI Kit(OAK)设备上运行。使用roboflowoak包连接模型和器械。
我想添加一个功能,用户可以在视频提要上画一个矩形,模型将只检测这个矩形内的台球。本质上,模型应该忽略此矩形之外的任何内容。
以下是我到目前为止编写的代码:
import cv2
import time
import numpy as np
from roboflowoak import RoboflowOak
# global variables
ix, iy = -1, -1
drawing = False
rect_over = False
x_start, y_start, x_end, y_end = -1, -1, -1, -1
# mouse callback function
def draw_rectangle(event, x, y, flags, param):
global ix, iy, drawing, rect_over, x_start, y_start, x_end, y_end
if event == cv2.EVENT_LBUTTONDOWN:
drawing = True
ix, iy = x, y
x_start, y_start = x, y
elif event == cv2.EVENT_MOUSEMOVE:
if drawing == True:
x_end, y_end = x, y
elif event == cv2.EVENT_LBUTTONUP:
drawing = False
rect_over = True
x_end, y_end = x, y
def main():
# Instantiate RoboflowOak object
rf = RoboflowOak(
model="billiards-y0wwp",
confidence=0.05,
overlap=0.2,
version="3",
api_key="My API Key",
rgb=True,
depth=True,
device=None,
blocking=True
)
# create named window to attach mouse callback function
cv2.namedWindow('frame')
cv2.setMouseCallback('frame', draw_rectangle)
# Loop to continuously run model inference and display output
while True:
t0 = time.time()
# Run model inference
result, frame, raw_frame, depth = rf.detect()
# Draw rectangle on frame
if rect_over:
cv2.rectangle(frame, pt1=(x_start, y_start), pt2=(x_end, y_end), color=(0, 255, 0), thickness=2)
# Filter out detections that are outside of the rectangle
if rect_over:
filtered_predictions = []
for p in result['predictions']:
bbox = p['bbox'] # Assuming the prediction has a 'bbox' key
x_center = bbox['left'] + bbox['width'] / 2
y_center = bbox['top'] + bbox['height'] / 2
if x_start < x_center < x_end and y_start < y_center < y_end:
filtered_predictions.append(p)
# Replace predictions with filtered predictions
result['predictions'] = filtered_predictions
# Print inference time
print(f"INFERENCE TIME IN MS: {1 / (time.time() - t0)}")
# Print predictions
print(f"PREDICTIONS: {[p for p in result['predictions']]}")
# Normalize depth for visualization
max_depth = np.amax(depth)
normalized_depth = depth / max_depth
# Show depth and frame
cv2.imshow("depth", normalized_depth)
cv2.imshow("frame", frame)
# Exit loop on 'q' key press
if cv2.waitKey(1) == ord('q'):
break
if __name__ == '__main__':
main()
字符串
如您所见,我实现了一个鼠标回调函数来绘制矩形,并在推断之后过滤掉该矩形之外的检测。然而,我觉得这可能不是最有效的方法,因为模型仍然在处理整个帧。
1条答案
按热度按时间7cjasjjr1#
字符串
第1点:避免使用全局变量,因为不推荐使用全局变量
第2点:在调整框时,考虑切换推理。推理就像一个障碍物,所以由于延迟,你几乎看不到任何东西。
第3点。不确定你是否想有一个真实的的时间广场,因为你正在画盒子
我为AI推理创建了一个虚拟图像和一个虚拟函数。请注意打开时打印的延迟。这也会影响您的'框绘制'的回调。因此,一个可能的好主意是分开做,而不是同时做。如果你真的想在推理的同时做调整;你可能需要了解多线程/多处理。