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

huangapple
go评论91阅读模式
英文:

Restricting Roboflow model inference to a user-defined region in OpenCV

问题

我目前正在进行一个项目，我在其中使用Roboflow模型来检测视频中的台球。视频信号经过OpenCV处理，模型在OpenAI Kit（OAK）设备上运行。使用roboflowoak包来接口模型和设备。

我想添加一个功能，使用户可以在视频中绘制一个矩形，模型只会检测这个矩形内的台球。基本上，矩形之外的任何东西都应该被模型忽略。

以下是我迄今为止编写的代码：

  1. import cv2
  2. import time
  3. import numpy as np
  4. from roboflowoak import RoboflowOak
  6. # 全局变量
  7. ix, iy = -1, -1
  8. drawing = False
  9. rect_over = False
  10. x_start, y_start, x_end, y_end = -1, -1, -1, -1
  12. # 鼠标回调函数
  13. def draw_rectangle(event, x, y, flags, param):
  14.     global ix, iy, drawing, rect_over, x_start, y_start, x_end, y_end
  16.     if event == cv2.EVENT_LBUTTONDOWN:
  17.         drawing = True
  18.         ix, iy = x, y
  19.         x_start, y_start = x, y
  21.     elif event == cv2.EVENT_MOUSEMOVE:
  22.         if drawing == True:
  23.             x_end, y_end = x, y
  25.     elif event == cv2.EVENT_LBUTTONUP:
  26.         drawing = False
  27.         rect_over = True
  28.         x_end, y_end = x, y
  30. def main():
  31.     # 实例化RoboflowOak对象
  32.     rf = RoboflowOak(
  33.         model="billiards-y0wwp",
  34.         confidence=0.05,
  35.         overlap=0.2,
  36.         version="3",
  37.         api_key="My API Key",
  38.         rgb=True,
  39.         depth=True,
  40.         device=None,
  41.         blocking=True
  42.     )
  44.     # 创建命名窗口以附加鼠标回调函数
  45.     cv2.namedWindow('frame')
  46.     cv2.setMouseCallback('frame', draw_rectangle)
  48.     # 循环以连续运行模型推断并显示输出
  49.     while True:
  50.         t0 = time.time()
  52.         # 运行模型推断
  53.         result, frame, raw_frame, depth = rf.detect()
  55.         # 在帧上绘制矩形
  56.         if rect_over:
  57.             cv2.rectangle(frame, pt1=(x_start, y_start), pt2=(x_end, y_end), color=(0, 255, 0), thickness=2)
  59.         # 过滤掉矩形之外的检测结果
  60.         if rect_over:
  61.             filtered_predictions = []
  62.             for p in result['predictions']:
  63.                 bbox = p['bbox']  # 假设预测有'bbox'键
  64.                 x_center = bbox['left'] + bbox['width'] / 2
  65.                 y_center = bbox['top'] + bbox['height'] / 2
  66.                 if x_start < x_center < x_end and y_start < y_center < y_end:
  67.                     filtered_predictions.append(p)
  69.             # 用过滤后的预测结果替换预测
  70.             result['predictions'] = filtered_predictions
  72.         # 打印推断时间
  73.         print(f"推断时间（毫秒）: {1 / (time.time() - t0)}")
  75.         # 打印预测
  76.         print(f"预测: {
    ]}")
  78.         # 规范化深度以进行可视化
  79.         max_depth = np.amax(depth)
  80.         normalized_depth = depth / max_depth
  82.         # 显示深度和帧
  83.         cv2.imshow("深度", normalized_depth)
  84.         cv2.imshow("帧", frame)
  86.         # 按下'q'键退出循环
  87.         if cv2.waitKey(1) == ord('q'):
  88.             break
  90. if __name__ == '__main__':
  91.     main()

如您所见，我已经实现了鼠标回调来绘制矩形，并在推断后过滤掉矩形之外的检测结果。然而，我觉得这可能不是最有效的方式，因为模型仍然在处理整个帧。

英文:

I'm currently working on a project where I'm using a Roboflow model to detect pool balls in a video feed. The video feed is processed using OpenCV and the model is running on an OpenAI Kit (OAK) device. The model and device are interfaced using the roboflowoak package.

I want to add a feature where the user can draw a rectangle on the video feed, and the model will only detect pool balls within this rectangle. Essentially, anything outside this rectangle should be ignored by the model.

Here's the code I've written so far:

  1. import cv2
  2. import time
  3. import numpy as np
  4. from roboflowoak import RoboflowOak
  5. # global variables
  6. ix, iy = -1, -1
  7. drawing = False
  8. rect_over = False
  9. x_start, y_start, x_end, y_end = -1, -1, -1, -1
  10. # mouse callback function
  11. def draw_rectangle(event, x, y, flags, param):
  12. global ix, iy, drawing, rect_over, x_start, y_start, x_end, y_end
  13. if event == cv2.EVENT_LBUTTONDOWN:
  14. drawing = True
  15. ix, iy = x, y
  16. x_start, y_start = x, y
  17. elif event == cv2.EVENT_MOUSEMOVE:
  18. if drawing == True:
  19. x_end, y_end = x, y
  20. elif event == cv2.EVENT_LBUTTONUP:
  21. drawing = False
  22. rect_over = True
  23. x_end, y_end = x, y
  24. def main():
  25. # Instantiate RoboflowOak object
  26. rf = RoboflowOak(
  27. model=&quot;billiards-y0wwp&quot;,
  28. confidence=0.05,
  29. overlap=0.2,
  30. version=&quot;3&quot;,
  31. api_key=&quot;My API Key&quot;,
  32. rgb=True,
  33. depth=True,
  34. device=None,
  35. blocking=True
  36. )
  37. # create named window to attach mouse callback function
  38. cv2.namedWindow(&#39;frame&#39;)
  39. cv2.setMouseCallback(&#39;frame&#39;, draw_rectangle)
  40. # Loop to continuously run model inference and display output
  41. while True:
  42. t0 = time.time()
  43. # Run model inference
  44. result, frame, raw_frame, depth = rf.detect()
  45. # Draw rectangle on frame
  46. if rect_over:
  47. cv2.rectangle(frame, pt1=(x_start, y_start), pt2=(x_end, y_end), color=(0, 255, 0), thickness=2)
  48. # Filter out detections that are outside of the rectangle
  49. if rect_over:
  50. filtered_predictions = []
  51. for p in result[&#39;predictions&#39;]:
  52. bbox = p[&#39;bbox&#39;]  # Assuming the prediction has a &#39;bbox&#39; key
  53. x_center = bbox[&#39;left&#39;] + bbox[&#39;width&#39;] / 2
  54. y_center = bbox[&#39;top&#39;] + bbox[&#39;height&#39;] / 2
  55. if x_start &lt; x_center &lt; x_end and y_start &lt; y_center &lt; y_end:
  56. filtered_predictions.append(p)
  57. # Replace predictions with filtered predictions
  58. result[&#39;predictions&#39;] = filtered_predictions
  59. # Print inference time
  60. print(f&quot;INFERENCE TIME IN MS: {1 / (time.time() - t0)}&quot;)
  61. # Print predictions
  62. print(f&quot;PREDICTIONS: {
    ]}&quot;)
  63. # Normalize depth for visualization
  64. max_depth = np.amax(depth)
  65. normalized_depth = depth / max_depth
  66. # Show depth and frame
  67. cv2.imshow(&quot;depth&quot;, normalized_depth)
  68. cv2.imshow(&quot;frame&quot;, frame)
  69. # Exit loop on &#39;q&#39; key press
  70. if cv2.waitKey(1) == ord(&#39;q&#39;):
  71. break
  72. if __name__ == &#39;__main__&#39;:
  73. main()

As you can see, I've implemented a mouse callback to draw the rectangle, and I'm filtering out detections outside this rectangle after the inference. However, I feel this might not be the most efficient way since the model is still processing the entire frame.

答案1

得分: 0

  1.     import cv2
  2.     import time
  3.     import numpy as np
  4.     
  5.     class test_class():
  6.     
  7.         def __init__(self):
  8.             # 初始化类变量
  9.             self.ix, self.iy = -1, -1
  10.             self.drawing = False
  11.             self.rect_over = False
  12.             self.x_start, self.y_start, self.x_end, self.y_end = -1, -1, -1, -1
  13.             
  14.             self.toggle = False # 初始化我们的推断开关
  15.             self.frame = np.zeros((250,250,3), np.uint8) # 虚拟图像
  16.     
  17.         # 鼠标回调函数
  18.         def draw_rectangle(self,event, x, y, flags, param):
  19.             if event == cv2.EVENT_LBUTTONDOWN:
  20.                 self.drawing = True
  21.                 self.ix, self.iy = x, y
  22.                 self.x_start, self.y_start = x, y
  23.     
  24.             elif event == cv2.EVENT_MOUSEMOVE:
  25.                 if self.drawing == True:
  26.                     self.x_end, self.y_end = x, y
  27.     
  28.             elif event == cv2.EVENT_LBUTTONUP:
  29.                 self.drawing = False
  30.                 self.rect_over = True
  31.                 self.x_end, self.y_end = x, y
  32.     
  33.         def inference(self): # 虚拟推断函数
  34.             # 在这里放置所有的AI/过滤器
  35.             print("我在运行AI程序")
  36.             time.sleep(5)
  37.             return
  38.             
  39.         def main(self):
  40.     
  41.             cv2.namedWindow('frame')
  42.             cv2.setMouseCallback('frame', self.draw_rectangle)
  43.     
  44.             # 循环持续运行模型推断并显示输出
  45.             while True:
  46.                 t0 = time.time()
  47.                 f_copy = self.frame.copy() # 将我的空白图像复制，因为我没有你的数据流
  48.                 # 运行模型推断
  49.                 if self.toggle is True: 
  50.                     self.inference()
  51.                 # 在帧上绘制矩形
  52.                 if self.rect_over:
  53.                     print("我在这里")
  54.                     cv2.rectangle(f_copy, pt1=(self.x_start, self.y_start), pt2=(self.x_end, self.y_end), color=(0, 255, 0), thickness=2)
  55.     
  56.                 cv2.imshow("frame", f_copy)
  57.                 if cv2.waitKey(1) == ord('f'): # 可能有更好的方法来做这个。无法中断推断步骤.....
  58.                     self.toggle = not self.toggle # 简单的布尔切换
  59.                 # 在按下'q'键时退出循环
  60.                 if cv2.waitKey(1) == ord('q'):
  61.                     break
  62.     
  63.     if __name__ == '__main__':
  64.         my_program = test_class()
  65.         my_program.main()
英文: 
  1. import cv2
  2. import time
  3. import numpy as np
  4. class test_class():
  5. def __init__(self):
  6. # global variables # We dont like global variables
  7. # Initiate class variables
  8. self.ix, self.iy = -1, -1
  9. self.drawing = False
  10. self.rect_over = False
  11. self.x_start, self.y_start, self.x_end, self.y_end = -1, -1, -1, -1
  12. self.toggle = False # INitialize our &#39;inference&#39; toggle
  13. self.frame = np.zeros((250,250,3), np.uint8) #Dummy image
  14. # mouse callback function
  15. def draw_rectangle(self,event, x, y, flags, param):
  16. if event == cv2.EVENT_LBUTTONDOWN:
  17. self.drawing = True
  18. self.ix, self.iy = x, y
  19. self.x_start, self.y_start = x, y
  20. elif event == cv2.EVENT_MOUSEMOVE:
  21. if self.drawing == True:
  22. self.x_end, self.y_end = x, y
  23. elif event == cv2.EVENT_LBUTTONUP:
  24. self.drawing = False
  25. self.rect_over = True
  26. self.x_end, self.y_end = x, y
  27. def inference(self): # Dummy Inferencing function. 
  28. #you can put all your AI /filtering here
  29. print(&quot;I am running AI stuff&quot;)
  30. time.sleep(5)
  31. return
  32. def main(self):
  33. cv2.namedWindow(&#39;frame&#39;)
  34. cv2.setMouseCallback(&#39;frame&#39;, self.draw_rectangle)
  35. # Loop to continuously run model inference and display output
  36. while True:
  37. t0 = time.time()
  38. f_copy = self.frame.copy() #Copying my blank image as I don&#39;t have your stream
  39. # Run model inference
  40. #result, frame, raw_frame, depth = rf.detect()
  41. if self.toggle is True: 
  42. self.inference()
  43. # Draw rectangle on frame
  44. if self.rect_over:
  45. print(&quot;im here&quot;)
  46. cv2.rectangle(f_copy, pt1=(self.x_start, self.y_start), pt2=(self.x_end, self.y_end), color=(0, 255, 0), thickness=2)
  47. cv2.imshow(&quot;frame&quot;, f_copy)
  48. if cv2.waitKey(1) == ord(&#39;f&#39;): #probably a better way to do this. Cant exactly interrupt the inferencing step.....
  49. self.toggle = not self.toggle # Simple boolean toggle
  50. # Exit loop on &#39;q&#39; key press
  51. if cv2.waitKey(1) == ord(&#39;q&#39;):
  52. break
  53. if __name__ == &#39;__main__&#39;:
  54. my_program = test_class()
  55. my_program.main()

Point 1: Avoid globals as its not recommended

Point 2: Consider toggling your inferencing while you are adjusting your box. Inferencing is acting as a blocker so you can barely see anything due to the delay.

Point 3. Not sure if you wanted to have a real time square as you are drawing the box

I created a dummmy image and a dummy function for the AI inferencing. Note the delay in the prints when you toggle it on. This also affects the callback for your 'box drawing'. Hence a probably good idea to do it separately and not concurrently. If you really want to do the adjustment at the same time as the inferencing; you probably have to look at multithreading/multiprocessing.

huangapple
  • 本文由 发表于 2023年7月11日 10:20:05
  • 转载请务必保留本文链接：https://go.coder-hub.com/76658358.html
  • machine-learning
  • opencv
  • python
匿名

发表评论

匿名网友

:?: :razz: :sad: :evil: :!: :smile: :oops: :grin: :eek: :shock: :???: :cool: :lol: :mad: :twisted: :roll: :wink: :idea: :arrow: :neutral: :cry: :mrgreen:

确定