105 lines
3.1 KiB
Python
105 lines
3.1 KiB
Python
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()
|