graduation-design/abandon_code/img_process/yolo.py
GitLab 559a94c55a 1.大体实现机械臂跟随物体移动
2.使用线程队列将机器人控制与视觉识别解耦,大大提高了系统的流畅性
2024-04-01 13:59:33 +08:00

105 lines
3.1 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

import cv2
from ultralytics import YOLO
from pyorbbecsdk import Pipeline, FrameSet
from pyorbbecsdk import Config
from pyorbbecsdk import OBSensorType, OBFormat
from pyorbbecsdk import OBError
from pyorbbecsdk import VideoStreamProfile
from img_process.orbbec_camera.utils import frame_to_bgr_image
def load_model():
# 载入 YOLOv8 模型
model = YOLO('yolov8n.pt')
return model
# 参数说明
# --img: 输入图片的路径
# --is_save_video: 是否保存识别后的视频
def process_img(img=None):
# 暂时性配置
video_path = "output_video.mp4"
# 加载模型
model = load_model()
if img is not None:
# 对图像运行 YOLOv8 推理
results = model(img)
# 在图像上可视化结果
annotated_img = results[0].plot()
# 显示带有标注的图像
cv2.imshow("YOLOv8 推理", annotated_img)
cv2.waitKey(0)
cv2.destroyAllWindows()
else:
# 获取本机摄像头内容
cap = cv2.VideoCapture(0)
# 奥比中光摄像头
# ESC键的键值
ESC_KEY = 27
# 创建配置对象
config = Config()
# 创建Pipeline对象
pipeline = Pipeline()
try:
# 获取颜色传感器的流配置列表
profile_list = pipeline.get_stream_profile_list(OBSensorType.COLOR_SENSOR)
try:
# 尝试获取指定分辨率、格式和帧率的颜色配置
color_profile: VideoStreamProfile = profile_list.get_video_stream_profile(700, 0, OBFormat.RGB, 30)
except OBError as e:
# 如果出错,打印错误信息并使用默认的颜色配置
print(e)
color_profile = profile_list.get_default_video_stream_profile()
# 启用颜色流
config.enable_stream(color_profile)
except Exception as e:
# 如果出现异常,打印错误信息并返回
print(e)
return
# 启动Pipeline
pipeline.start(config)
while True:
try:
# 等待获取帧集最多等待100毫秒
frames: FrameSet = pipeline.wait_for_frames(100)
if frames is None:
continue
# 获取颜色帧
color_frame = frames.get_color_frame()
if color_frame is None:
continue
# 将帧转换为BGR格式的图像
color_image = frame_to_bgr_image(color_frame)
model.predict(source=color_image, show=True)
if color_image is None:
print("failed to convert frame to image")
continue
key = cv2.waitKey(1)
if key == ord('q') or key == ESC_KEY:
break
except KeyboardInterrupt:
# 如果捕获到键盘中断信号,则退出循环
break
# 释放视频捕获对象并关闭显示窗口
cap.release()
cv2.destroyAllWindows()
# 停止Pipeline
pipeline.stop()
if __name__ == '__main__':
process_img()