将Roboflow模型推理限制在OpenCV中的用户定义区域

s6fujrry  于 2023-08-06  发布在  其他
关注(0)|答案(1)|浏览(130)

我目前正在做一个项目,我正在使用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()

字符串
如您所见,我实现了一个鼠标回调函数来绘制矩形,并在推断之后过滤掉该矩形之外的检测。然而,我觉得这可能不是最有效的方法,因为模型仍然在处理整个帧。

7cjasjjr

7cjasjjr1#

import cv2
import time
import numpy as np

class test_class():

    def __init__(self):
# global variables # We dont like global variables
# Initiate class variables
        self.ix, self.iy = -1, -1
        self.drawing = False
        self.rect_over = False
        self.x_start, self.y_start, self.x_end, self.y_end = -1, -1, -1, -1
        
        self.toggle = False # INitialize our 'inference' toggle
        self.frame = np.zeros((250,250,3), np.uint8) #Dummy image

# mouse callback function
    def draw_rectangle(self,event, x, y, flags, param):
       
       

        if event == cv2.EVENT_LBUTTONDOWN:
            self.drawing = True
            self.ix, self.iy = x, y
            self.x_start, self.y_start = x, y

        elif event == cv2.EVENT_MOUSEMOVE:
            if self.drawing == True:
                self.x_end, self.y_end = x, y

        elif event == cv2.EVENT_LBUTTONUP:
            self.drawing = False
            self.rect_over = True
            self.x_end, self.y_end = x, y

    def inference(self): # Dummy Inferencing function. 
        #you can put all your AI /filtering here
        print("I am running AI stuff")
        time.sleep(5)
        return
        
    def main(self):

        cv2.namedWindow('frame')
        cv2.setMouseCallback('frame', self.draw_rectangle)

        # Loop to continuously run model inference and display output
        while True:
            t0 = time.time()
            f_copy = self.frame.copy() #Copying my blank image as I don't have your stream
            # Run model inference
            #result, frame, raw_frame, depth = rf.detect()
            if self.toggle is True: 
                self.inference()
            # Draw rectangle on frame
            if self.rect_over:
                print("im here")
                cv2.rectangle(f_copy, pt1=(self.x_start, self.y_start), pt2=(self.x_end, self.y_end), color=(0, 255, 0), thickness=2)

            cv2.imshow("frame", f_copy)
            if cv2.waitKey(1) == ord('f'): #probably a better way to do this. Cant exactly interrupt the inferencing step.....
                self.toggle = not self.toggle # Simple boolean toggle
            # Exit loop on 'q' key press
            if cv2.waitKey(1) == ord('q'):
                break

if __name__ == '__main__':
    my_program = test_class()
    my_program.main()

字符串
第1点:避免使用全局变量,因为不推荐使用全局变量
第2点:在调整框时,考虑切换推理。推理就像一个障碍物,所以由于延迟,你几乎看不到任何东西。
第3点。不确定你是否想有一个真实的的时间广场,因为你正在画盒子
我为AI推理创建了一个虚拟图像和一个虚拟函数。请注意打开时打印的延迟。这也会影响您的'框绘制'的回调。因此,一个可能的好主意是分开做,而不是同时做。如果你真的想在推理的同时做调整;你可能需要了解多线程/多处理。

相关问题