import cv2



import pyopenpose as op
# RTSP地址
rtsp_url = "rtsp://192.168.1.115"

# 打开RTSP流
cap = cv2.VideoCapture(rtsp_url)

# 检查摄像头是否成功打开
if not cap.isOpened():
    print("无法打开摄像头")
    exit()

# 创建视频写入对象，用于保存视频
out = cv2.VideoWriter('output.mp4', cv2.VideoWriter_fourcc(*'mp4v'), 30, (640, 480))

# OpenPose 参数
params = {
    "model_folder": "../../../models/",  # 替换为您的 OpenPose 模型文件夹路径
    "hand": False,  # 是否检测手部关键点
    "face": False,  # 是否检测面部关键点
    "display": 0,  # 显示窗口的级别，0为不显示
}

# 初始化 OpenPose
opWrapper = op.WrapperPython()
opWrapper.configure(params)
opWrapper.start()

# 读取并显示视频流
while True:
    ret, frame = cap.read()
    if not ret:
        print("无法读取视频流")
        break

    # 执行姿态检测
    datum = op.Datum()
    datum.cvInputData = frame
    opWrapper.emplaceAndPop([datum])

    # 获取姿态检测结果
    pose_keypoints = datum.poseKeypoints

    # 在帧上绘制姿态关键点
    if pose_keypoints is not None:
        for person in pose_keypoints:
            for point in person:
                if point[2] > 0.1:  # 确保置信度足够高
                    cv2.circle(frame, (int(point[0]), int(point[1])), 4, (0, 255, 255), thickness=-1, lineType=cv2.FILLED)

    # 显示视频流
    cv2.imshow('RTSP Stream', frame)

    # 写入视频帧到输出文件
    out.write(frame)

    # 按下'q'键退出循环
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

# 释放资源
cap.release()
out.release()
cv2.destroyAllWindows()