人体姿势识别：openpose与yolov8_pose使用心得

2024-11-13

learn

openpose的使用

CodeWithGPU里的镜像只支持3090及其以上的GPU部署运行，在终端输入下面命令即可运行查看项目情况：

1 2	cd openpose ./build/examples/openpose/openpose.bin --video ~/test.mp4 --write_json ~/output_json --write_video ~/output.mp4 -display 0 -face -hand

下面让我们继续微调，在根目录下面，创建两个文件夹，一个test，一个test_run，放入我们的测试图片和输出图片。

再在根目录创建一个python文件，detect_pose函数里有着我们的判断算法（不过我当时写的时候感觉有点奇怪，效果也一般般，需要你自己调整）

main函数里还有需要一个SimHei.ttf，也就是中文字体，需要自己下载，放在根目录里，记得名称要对喔，拉姆当时就是因为下载的是SimHei.ttf，但是代码里写的是font_path = "simhei.ttf"，导致拉姆修了半个小时的bug（悲）

完整代码如下：

import sys
import cv2
import numpy as np
import os
import math
from PIL import Image, ImageDraw, ImageFont

sys.path.append('/root/openpose/build/python')
from openpose import pyopenpose as op

def calculate_angle(point1, point2, point3):
    """
    计算三个点形成的角度
    """
    if 0 in (point1[2], point2[2], point3[2]):  # 如果置信度为0，表示点没有被检测到
        return None
        
    a = np.array([point1[0], point1[1]])
    b = np.array([point2[0], point2[1]])
    c = np.array([point3[0], point3[1]])
    
    ba = a - b
    bc = c - b
    
    cosine_angle = np.dot(ba, bc) / (np.linalg.norm(ba) * np.linalg.norm(bc))
    angle = np.arccos(cosine_angle)
    
    return np.degrees(angle)

def detect_pose(pose_keypoints):
    """
    基于关键点判断姿势
    返回姿势名称和置信度
    """
    if pose_keypoints is None or len(pose_keypoints) == 0:
        return "No person detected", 0.0
        
    # 提取关键点
    # 鼻子和颈部
    nose = pose_keypoints[0]
    neck = pose_keypoints[1]
    # 臀部中心点
    hip_center = pose_keypoints[8]
    # 左右臀部
    left_hip = pose_keypoints[12]
    right_hip = pose_keypoints[9]
    # 左右膝盖
    left_knee = pose_keypoints[13]
    right_knee = pose_keypoints[10]
    # 左右脚踝
    left_ankle = pose_keypoints[14]
    right_ankle = pose_keypoints[11]
    # 左右肩膀
    left_shoulder = pose_keypoints[5]
    right_shoulder = pose_keypoints[2]
    # 左右肘部
    left_elbow = pose_keypoints[6]
    right_elbow = pose_keypoints[3]
    # 左右手腕
    left_wrist = pose_keypoints[7]
    right_wrist = pose_keypoints[4]
    
    # 计算膝盖角度
    left_knee_angle = calculate_angle(left_hip, left_knee, left_ankle)
    right_knee_angle = calculate_angle(right_hip, right_knee, right_ankle)
    
    # 获取髋部高度（相对于脚踝）
    hip_height = abs(hip_center[1] - left_ankle[1] + hip_center[1] - right_ankle[1]) / 2
    
    # 计算鼻子和脚高度（身高）
    body_height = abs(nose[1]-left_ankle[1])
    bo_height = abs(left_shoulder [1]-right_hip[1])
  
    
    # 姿态判断逻辑
    if left_knee_angle is None or right_knee_angle is None or hip_height is None:
        return "关键点缺失"
    
    # 站立判断：膝盖角度接近180度
    if (left_knee_angle > 160 and right_knee_angle > 160): 
        if (body_height > 2 * bo_height):
            return "站立"
        else:
            return "平躺"
    
    # # 关键点判断
    # if nose is None and neck is None:
    #     return "关键点缺失", 0.0
    
    # # 姿态判断逻辑
    # if left_knee_angle is None or right_knee_angle is None or hip_height is None:
    #     return "关键点缺失", 0.0
    
    # 蹲下判断：膝盖角度小于120度且髋部降低
    if (left_knee_angle < 120 and right_knee_angle < 120):
        confidence = (120 - max(left_knee_angle, right_knee_angle)) / 60
        return "蹲下"

    
    # 抬腿判断：一条腿的膝盖角度明显小于另一条腿
    knee_angle_diff = abs(left_knee_angle - right_knee_angle)
    if (knee_angle_diff > 30):
        return "抬腿"
    
    # 双腿抬起判断：两条腿的膝盖角度都小于某个阈值（例如120度）
    if (left_knee_angle < 120 and right_knee_angle < 120):
        return "双腿抬起"
    
    # 举手判断：手腕和肩膀的高度
    if (left_wrist[1] < left_shoulder[1] or right_wrist[1] < right_shoulder[1]):
        return "手部举起"
        
    return "未检测到正确动作"

  
# 图像进行测试    
def main():
    # OpenPose参数配置
    params = dict()
    params["face"] = True
    params["hand"] = True 
    params["model_folder"] = "/root/openpose/models/"
    # 初始化OpenPose
    opWrapper = op.WrapperPython()
    opWrapper.configure(params)
    opWrapper.start()
    
    # 图像文件夹路径
    img_folder = "./test"  # 替换为你的图像文件夹路径
    test_folder = "./test_run"  # 替换为你的输出文件夹路径
    
    # 确保输出文件夹存在
    if not os.path.exists(test_folder):
        os.makedirs(test_folder)
    
    # 获取图像文件列表
    image_paths = [os.path.join(img_folder, f) for f in os.listdir(img_folder) if f.endswith(('.png', '.jpg', '.jpeg'))]
    
    for image_path in image_paths:
        # 读取图像文件
        frame = cv2.imread(image_path)
        if frame is None:
            print(f"Error: Could not read image {image_path}")
            continue
            
        # 准备OpenPose输入
        datum = op.Datum()
        datum.cvInputData = frame
        
        # OpenPose处理
        opWrapper.emplaceAndPop(op.VectorDatum([datum]))
        
        # 获取姿态预测
        if datum.poseKeypoints is not None and len(datum.poseKeypoints) > 0:
            pose, confidence = detect_pose(datum.poseKeypoints[0])
            
            # 在图像上绘制关键点和框框
            for keypoints in datum.poseKeypoints:
                for i, keypoint in enumerate(keypoints):
                    x, y, confidence = keypoint
                    if confidence > 0.1:  # 只绘制置信度大于0.1的关键点
                        cv2.circle(frame, (int(x), int(y)), 5, (0, 255, 0), -1)
                        cv2.putText(frame, str(i), (int(x), int(y)), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 2)
            
            # 使用PIL库在图像上显示中文字符
            pil_image = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
            draw = ImageDraw.Draw(pil_image)
            font_path = "SimHei.ttf"  # 字体文件与脚本文件在同一目录下
            font = ImageFont.truetype(font_path, 24) 
            draw.text((50, 50), f"{pose} ({confidence:.2f})", font=font, fill=(0, 255, 0) if confidence > 0.5 else (0, 0, 255))
            frame = cv2.cvtColor(np.array(pil_image), cv2.COLOR_RGB2BGR)
        
        # 保存处理后的图像
        output_path = os.path.join(test_folder, os.path.basename(image_path))
        cv2.imwrite(output_path, frame)
        print(f"Processed image saved to {output_path}")
    
    # 释放资源
    cv2.destroyAllWindows()

if __name__ == "__main__":
    main()

祝使用愉快/记得要自己修判断算法逻辑喔

yolov8_pose的使用

安装指令：
首先克隆这个仓库到本地：

1 2	git clone https://github.com/Mrjianning/yolov8_pose_judge.git cd yolov8_pose_judge

安装所需的依赖：

1	pip install -r requirements.txt

快速开始:
运行以下命令以启动姿态检测：python main.py

没有问题的话，我们在根目录下创建一个python文件

python代码里有三个实现，分别是针对于图片的静态识别，针对于视频的动态识别，以及调用摄像头的实时动态识别

当然，determine_pose函数中姿态识别的算法逻辑代码，拉姆这边还是有一些问题的，请你自己修改喔~

完整代码如下：

# -*- coding: utf-8 -*-
import torch
import os
import cv2 as cv
import numpy as np
from ultralytics.data.augment import LetterBox
from ultralytics.utils import ops
from ultralytics.engine.results import Results
import copy
from PIL import Image, ImageDraw, ImageFont
import math
import sys

sys.stdout.reconfigure(encoding='utf-8')


class YOLOv8Pose:
    def __init__(self, model_path, device='cpu', conf=0.25, iou=0.7):
        self.device = device
        self.conf = conf
        self.iou = iou
        self.model = self.load_model(model_path)
        self.letterbox = LetterBox([640, 640], auto=True, stride=32)

    def load_model(self, model_path):
        ckpt = torch.load(model_path, map_location=self.device)
        model = ckpt['model'].to(self.device).float()
        model.eval()
        return model

    def preprocess(self, img_path):
        im = cv.imread(img_path)
        im = [im]
        orig_imgs = copy.deepcopy(im)
        im = [self.letterbox(image=x) for x in im]
        im = im[0][None]
        im = im[..., ::-1].transpose((0, 3, 1, 2))
        im = np.ascontiguousarray(im)
        im = torch.from_numpy(im)
        img = im.to(self.device)
        img = img.float()
        img /= 255
        return img, orig_imgs

    def infer(self, img):
        preds = self.model(img)
        prediction = ops.non_max_suppression(preds, self.conf, self.iou, agnostic=False, max_det=300, classes=None,
                                             nc=len(self.model.names))
        return prediction

    def postprocess(self, prediction, orig_imgs, img_shape, img_path):
        results = []
        for i, pred in enumerate(prediction):
            orig_img = orig_imgs[i] if isinstance(orig_imgs, list) else orig_imgs
            shape = orig_img.shape
            pred[:, :4] = ops.scale_boxes(img_shape, pred[:, :4], shape).round()
            pred_kpts = pred[:, 6:].view(len(pred), *self.model.kpt_shape) if len(pred) else pred[:, 6:]
            pred_kpts = ops.scale_coords(img_shape, pred_kpts, shape)

            results.append(
                Results(orig_img=orig_img,
                        path=img_path,
                        names=self.model.names,
                        boxes=pred[:, :6],
                        keypoints=pred_kpts))
        return results

    def plot_results(self, results, img):
        plot_args = {'line_width': None, 'boxes': True, 'conf': True, 'labels': True}
        plot_args['im_gpu'] = img[0]
        result = results[0]
        plotted_img = result.plot(**plot_args)
        return plotted_img

    def detect(self, img_path):
        img, orig_imgs = self.preprocess(img_path)
        prediction = self.infer(img)
        results = self.postprocess(prediction, orig_imgs, img.shape[2:], img_path)
        return results

    def show_results(self, results, img_path):
        img, _ = self.preprocess(img_path)
        plotted_img = self.plot_results(results, img)

        cv.imshow('plotted_img', plotted_img)
        cv.waitKey(0)
        cv.destroyAllWindows()

    def determine_pose(self, keypoints):
        """
        根据关键点确定人体姿态。
        :param keypoints: 包含每个关键点的坐标和置信度的元组列表，例如 [(x1, y1, conf1), ...]
        :return: 姿态类别：'站立','抬腿','平躺',"抬胳膊","蹲下","举左手","举右手","腿弯曲",'未检测到正确动作'
        """

        # 定义关键点索引
        left_ankle, right_ankle = keypoints[15], keypoints[16]
        left_knee, right_knee = keypoints[13], keypoints[14]
        left_hip, right_hip = keypoints[11], keypoints[12]
        left_shoulder, right_shoulder = keypoints[5], keypoints[6]
        left_elbow, right_elbow = keypoints[7], keypoints[8]
        left_wrist, right_wrist = keypoints[9], keypoints[10]
        nose = keypoints[0]

        # 检查置信度
        if any(conf < 0 for _, _, conf in
               [left_ankle, right_ankle, left_knee, right_knee, left_hip, right_hip, left_shoulder, right_shoulder]):
            return 'xxx未检测到正确动作'

        def calculate_angle(point1, point2, point3):
            """计算由三个点形成的角度。"""
            vector1 = (point2[0] - point1[0], point2[1] - point1[1])
            vector2 = (point3[0] - point2[0], point3[1] - point2[1])
            dot_product = vector1[0] * vector2[0] + vector1[1] * vector2[1]
            magnitude1 = math.sqrt(vector1[0] ** 2 + vector1[1] ** 2)
            magnitude2 = math.sqrt(vector2[0] ** 2 + vector2[1] ** 2)
            if magnitude1 * magnitude2 == 0:
                return 0
            cos_theta = max(-1, min(1, dot_product / (magnitude1 * magnitude2)))
            return math.degrees(math.acos(cos_theta))

        # 计算关键角度
        left_knee_angle = calculate_angle(left_hip, left_knee, left_ankle)
        right_knee_angle = calculate_angle(right_hip, right_knee, right_ankle)
        left_hip_angle = calculate_angle(left_shoulder, left_hip, left_knee)
        right_hip_angle = calculate_angle(right_shoulder, right_hip, right_knee)

        # 计算关键位置关系
        body_height = abs(nose[1] - left_ankle[1])
        bo_height = abs(left_ankle[1] - right_hip[1])
        ankle_distance = abs(left_ankle[0] - right_ankle[0])
        knee_distance = abs(left_knee[0] - right_knee[0])
        avg_hip_height = (left_hip[1] + right_hip[1]) / 2
        avg_shoulder_height = (left_shoulder[1] + right_shoulder[1]) / 2
        avg_knee_height = (left_knee[1] + right_knee[1]) / 2
        avg_ankle_height = (left_ankle[1] + right_ankle[1]) / 2
        shoulder_to_hip_dist = abs((left_shoulder[1] + right_shoulder[1]) / 2 -(left_hip[1] + right_hip[1]) / 2)
        hip_to_ankle_dist = abs((left_hip[1] + right_hip[1]) / 2 -(left_ankle[1] + right_ankle[1]) / 2)

        # if (left_shoulder[1] < left_hip[1] < left_knee[1] < left_ankle[1] and
        #         right_shoulder[1] < right_hip[1] < right_knee[1] < right_ankle[1]):
        #     # 确保身体没有抬起，适用于完全平躺的场景
        #     if (body_height > 2 * bo_height):
        #         return "1站立"
        #
        #     shoulder_to_hip_dist = abs((left_shoulder[1] + right_shoulder[1]) / 2 -
        #                                (left_hip[1] + right_hip[1]) / 2)
        #     hip_to_ankle_dist = abs((left_hip[1] + right_hip[1]) / 2 -
        #                             (left_ankle[1] + right_ankle[1]) / 2)
        #
        #     if shoulder_to_hip_dist > hip_to_ankle_dist * 0.5:
        #         if left_ankle[0]>left_knee[0]:
        #             return '2站立'
        #         else:
        #             return "平躺"
        #     else:
        #         if (left_knee[1] >= avg_hip_height and right_knee[1] >= avg_hip_height) or (
        #                 left_knee_angle >= 120 and right_knee_angle >= 120):
        #             return '平躺'
        #         else:
        #             return '可能抬腿'
        #
        # # 抬腿动作检测，要在其他条件下做进一步细化
        # knee_angle_diff = abs(left_knee_angle - right_knee_angle)
        # if (left_knee[1] < avg_hip_height or right_knee[1] < avg_hip_height) and (
        #         left_knee_angle < 120 or right_knee_angle < 120) or (knee_angle_diff > 30):
        #     return '抬腿'


        if (left_shoulder[1] < left_hip[1] < left_knee[1] < left_ankle[1] and
                right_shoulder[1] < right_hip[1] < right_knee[1] < right_ankle[1]):
            if (body_height > 2 * bo_height):
                return "1站立"

            # 如果肩膀到髋部的距离大于一定比例，判定为站立
            if shoulder_to_hip_dist > hip_to_ankle_dist * 0.5:
                return '2站立'
            else:
                return '平躺'

        knee_angle_diff = abs(left_knee_angle - right_knee_angle)
        # 2. 抬腿动作检测
        if (left_knee[1] < avg_hip_height or right_knee[1] < avg_hip_height) and (
                left_knee_angle < 120 or right_knee_angle < 120) or (knee_angle_diff > 30):
            return '抬腿'

        # 4. 蹲下检测
        if (avg_knee_height < avg_hip_height and left_knee_angle < 150
                and right_knee_angle < 150):
            return '蹲下'

        # 举手判断：手腕和肩膀的高度
        if (left_wrist[1] < left_shoulder[1] or right_wrist[1] < right_shoulder[1] or
                left_elbow[1] < left_shoulder[1] or right_elbow[1] < right_shoulder[1]):
            return "手部举起"

        return '未检测到正确动作'


    def get_results_as_dicts(self, results):
        results_list = []
        for result in results:
            for i, box in enumerate(result.boxes.data):
                cls_id = result.boxes.cls[i].item()  # 获取类别ID
                class_name = result.names.get(cls_id)
                confidence = result.boxes.conf[i].item()  # 获取置信度
                bbox = box[:4].tolist()  # 获取边框坐标

                # 准备关键点信息
                keypoints = []
                if result.keypoints.has_visible:
                    for kpt in result.keypoints.data:
                        for point in kpt:
                            x, y, conf = point.tolist()
                            keypoints.append((x, y, conf))

                # 判断姿态
                pose = self.determine_pose(keypoints)

                result_dict = {
                    "class_name": class_name,
                    "boxes": bbox,
                    "confidence": confidence,
                    "keypoints": keypoints,
                    "pose": pose  # 增加姿态信息
                }
                results_list.append(result_dict)

        return results_list


# 关键点的连接关系（COCO数据集的例子）
skeleton = [
    (0, 1), (0, 2), (1, 3), (2, 4), (3, 5), (4, 6),  # 头部到身体
    (5, 6), (5, 7), (7, 9), (6, 8), (8, 10),  # 身体到手臂
    (5, 11), (6, 12), (11, 12),  # 躯干
    (11, 13), (12, 14), (13, 15), (14, 16)  # 躯干到腿
]


# opencv绘制结果
def draw_results(image, results):
    height, width, _ = image.shape
    scale_factor = height / 500  # 假设基准高度为500像素

    line_thickness = max(1, int(2 * scale_factor))  # 线条粗细至少为1
    font_scale = max(0.5, 0.5 * scale_factor)  # 字体缩放至少为0.5
    circle_radius = max(1, int(3 * scale_factor))  # 圆点半径至少为1

    for res in results:
        # 绘制边框和类别名
        bbox = res['boxes']
        confidence = res['confidence']
        class_name = res['class_name']

        start_point = (int(bbox[0]), int(bbox[1]))
        end_point = (int(bbox[2]), int(bbox[3]))
        cv.rectangle(image, start_point, end_point, (255, 0, 0), line_thickness)

        label = f"{class_name}: {confidence:.2f}"
        cv.putText(image, label, (start_point[0], start_point[1] - 10), cv.FONT_HERSHEY_SIMPLEX, font_scale,
                   (0, 0, 255), line_thickness)

        # 绘制关键点
        keypoints = [(int(kpt[0]), int(kpt[1]), kpt[2]) for kpt in res['keypoints']]
        for idx, (x, y, conf) in enumerate(keypoints):
            if conf > 0.5:
                cv.circle(image, (x, y), circle_radius, (0, 255, 0), -1)
                cv.putText(image, str(idx), (x + 5, y + 5), cv.FONT_HERSHEY_SIMPLEX, font_scale, (0, 0, 255),
                           line_thickness)

        # 绘制骨架
        for start_idx, end_idx in skeleton:
            start_point = keypoints[start_idx]
            end_point = keypoints[end_idx]
            if start_point[2] > 0.5 and end_point[2] > 0.5:
                cv.line(image, (start_point[0], start_point[1]), (end_point[0], end_point[1]), (0, 255, 0),
                        line_thickness)

        # 显示姿态信息
        pose_label = f"姿态: {res['pose']}"
        # cv.putText(image, pose_label, (start_point[0], start_point[1] - 30), cv.FONT_HERSHEY_SIMPLEX, font_scale, (255, 255, 255), line_thickness)

        image = draw_chinese_text(image, pose_label, (50, 50), "./STSONG.TTF", 48, (0, 255, 0))

    return image


def draw_chinese_text(image, text, position, font_path, font_size, color):
    """
    在图像上绘制中文文本
    :param image: OpenCV 图像
    :param text: 要绘制的中文文本
    :param position: 文本绘制位置 (x, y)
    :param font_path: 字体文件路径
    :param font_size: 字体大小
    :param color: 文本颜色 (B, G, R)
    :return: 绘制中文文本后的图像
    """
    # 将 OpenCV 图像转换为 PIL 图像
    image_pil = Image.fromarray(cv.cvtColor(image, cv.COLOR_BGR2RGB))
    draw = ImageDraw.Draw(image_pil)
    font = ImageFont.truetype(font_path, font_size, encoding="utf-8")

    # 绘制中文文本
    draw.text(position, text, font=font, fill=color)

    # 将 PIL 图像转换回 OpenCV 图像
    image = cv.cvtColor(np.array(image_pil), cv.COLOR_RGB2BGR)
    return

# 视频检测
# 修复绘制中文文本函数
def draw_chinese_text(image, text, position, font_path, font_size, color):
    """
    在图像上绘制中文文本
    :param image: OpenCV 图像
    :param text: 要绘制的中文文本
    :param position: 文本绘制位置 (x, y)
    :param font_path: 字体文件路径
    :param font_size: 字体大小
    :param color: 文本颜色 (B, G, R)
    :return: 绘制中文文本后的图像
    """
    # 将 OpenCV 图像转换为 PIL 图像
    image_pil = Image.fromarray(cv.cvtColor(image, cv.COLOR_BGR2RGB))
    draw = ImageDraw.Draw(image_pil)

    try:
        font = ImageFont.truetype(font_path, font_size, encoding="utf-8")
    except Exception as e:
        print(f"字体加载错误: {e}")
        # 如果字体加载失败，返回原图
        return image

    # 绘制中文文本
    draw.text(position, text, font=font, fill=color[::-1])  # 颜色需要转换BGR到RGB

    # 将 PIL 图像转换回 OpenCV 图像并返回
    return cv.cvtColor(np.array(image_pil), cv.COLOR_RGB2BGR)


# 修复绘制结果函数
def draw_results(image, results):
    if image is None or len(image.shape) != 3:
        print("无效的输入图像")
        return None

    height, width, _ = image.shape
    scale_factor = height / 500  # 假设基准高度为500像素

    line_thickness = max(1, int(2 * scale_factor))  # 线条粗细至少为1
    font_scale = max(0.5, 0.5 * scale_factor)  # 字体缩放至少为0.5
    circle_radius = max(1, int(3 * scale_factor))  # 圆点半径至少为1

    # 创建图像副本以避免修改原图
    image_with_results = image.copy()

    for res in results:
        try:
            # 绘制边框和类别名
            bbox = res['boxes']
            confidence = res['confidence']
            class_name = res['class_name']

            start_point = (int(bbox[0]), int(bbox[1]))
            end_point = (int(bbox[2]), int(bbox[3]))
            cv.rectangle(image_with_results, start_point, end_point, (255, 0, 0), line_thickness)

            label = f"{class_name}: {confidence:.2f}"
            cv.putText(image_with_results, label, (start_point[0], start_point[1] - 10),
                       cv.FONT_HERSHEY_SIMPLEX, font_scale, (0, 0, 255), line_thickness)

            # 绘制关键点
            keypoints = [(int(kpt[0]), int(kpt[1]), kpt[2]) for kpt in res['keypoints']]
            for idx, (x, y, conf) in enumerate(keypoints):
                if conf > 0.2:
                    cv.circle(image_with_results, (x, y), circle_radius, (0, 255, 0), -1)
                    cv.putText(image_with_results, str(idx), (x + 5, y + 5),
                               cv.FONT_HERSHEY_SIMPLEX, font_scale, (0, 0, 255), line_thickness)

            # 绘制骨架
            for start_idx, end_idx in skeleton:
                start_point = keypoints[start_idx]
                end_point = keypoints[end_idx]
                if start_point[2] > 0.5 and end_point[2] > 0.5:
                    cv.line(image_with_results, (start_point[0], start_point[1]),
                            (end_point[0], end_point[1]), (0, 255, 0), line_thickness)

            # 显示姿态信息
            pose_label = f"姿态: {res['pose']}"
            # 使用修复后的draw_chinese_text函数
            image_with_results = draw_chinese_text(
                image_with_results,
                pose_label,
                (50, 50),
                "./STSONG.TTF",
                48,
                (0, 255, 0)
            )

        except Exception as e:
            print(f"绘制结果时发生错误: {e}")
            continue

    return image_with_results

# 图片检测
if __name__ == "__main__":
    yolov8 = YOLOv8Pose(model_path='yolov8x-pose.pt', device='cpu', conf=0.25, iou=0.7)
    # imgs = ["img_2.png"]
    img_folder = "./test/"
    output_folder = "./test_run/"

    imgs = [os.path.join(img_folder, img) for img in os.listdir(img_folder) if img.endswith(('.png', '.jpg', '.jpeg'))]
    # 存储处理后的图片结果
    processed_images = []

    # 处理每张图片并存储结果
    for img_path in imgs:
        img = cv.imread(img_path)
        results = yolov8.detect(img_path)
        results_ = yolov8.get_results_as_dicts(results)
        img_with_results = draw_results(img, results_)
        processed_images.append(img_with_results)

        # 保存处理后的图片到输出文件夹
        output_path = os.path.join(output_folder, os.path.basename(img_path))
        cv.imwrite(output_path, img_with_results)

    # 计算每张图片的显示尺寸
    max_width = 600
    for idx, img_result in enumerate(processed_images):
        height = int(img_result.shape[0] * (max_width / img_result.shape[1]))
        if height > 900:
            height = 900
            max_width = int(img_result.shape[1] * (height / img_result.shape[0]))
        processed_images[idx] = cv.resize(img_result, (max_width, height))

    # 显示每张图片的处理结果在独立窗口中
    for idx, img_result in enumerate(processed_images):
        cv.imshow(f'Processed Image {idx + 1}', img_result)

    cv.waitKey(0)
    cv.destroyAllWindows()

# # 视频处理主函数
# if __name__ == "__main__":
#     # 初始化模型
#     yolov8 = YOLOv8Pose(model_path='yolov8x-pose.pt', device='cpu', conf=0.25, iou=0.7)
#
#     # 设置输入和输出视频路径
#     input_video_path = "test.mp4"  # 替换为你的输入视频路径
#     output_video_path = "output_video.mp4"  # 输出视频的保存路径
#
#     # 打开视频文件
#     cap = cv.VideoCapture(input_video_path)
#     if not cap.isOpened():
#         print(f"无法打开视频文件: {input_video_path}")
#         exit()
#
#     # 获取视频的基本信息
#     frame_width = int(cap.get(cv.CAP_PROP_FRAME_WIDTH))
#     frame_height = int(cap.get(cv.CAP_PROP_FRAME_HEIGHT))
#     fps = int(cap.get(cv.CAP_PROP_FPS))
#     total_frames = int(cap.get(cv.CAP_PROP_FRAME_COUNT))
#
#     print(f"视频信息: {frame_width}x{frame_height} @ {fps}fps, 总帧数: {total_frames}")
#
#     # 创建视频写入器
#     fourcc = cv.VideoWriter_fourcc(*'mp4v')
#     out = cv.VideoWriter(output_video_path, fourcc, fps, (frame_width, frame_height))
#
#     processed_frames = 0
#     temp_img_path = "temp_frame.jpg"
#
#     try:
#         while cap.isOpened():
#             ret, frame = cap.read()
#             if not ret:
#                 print("\n到达视频结尾")
#                 break
#
#             try:
#                 # 确保帧是有效的
#                 if frame is None or frame.size == 0:
#                     print(f"\n无效帧 {processed_frames}")
#                     processed_frames += 1
#                     continue
#
#                 # 保存和处理当前帧
#                 cv.imwrite(temp_img_path, frame)
#
#                 # 检查临时文件是否成功创建
#                 if not os.path.exists(temp_img_path):
#                     print(f"\n临时文件创建失败，跳过帧 {processed_frames}")
#                     processed_frames += 1
#                     continue
#
#                 # 进行检测
#                 results = yolov8.detect(temp_img_path)
#                 results_ = yolov8.get_results_as_dicts(results)
#
#                 # 绘制结果
#                 frame_with_results = draw_results(frame, results_)
#
#                 if frame_with_results is not None:
#                     # 写入处理后的帧
#                     out.write(frame_with_results)
#
#                     # 显示处理进度
#                     processed_frames += 1
#                     progress = (processed_frames / total_frames) * 100
#                     print(f"\r处理进度: {progress:.2f}%", end='')
#
#                     # 显示预览
#                     cv.imshow('Video Processing', frame_with_results)
#
#                     if cv.waitKey(1) & 0xFF == ord('q'):
#                         print("\n用户中断处理")
#                         break
#                 else:
#                     print(f"\n帧 {processed_frames} 处理失败")
#                     processed_frames += 1
#
#             except Exception as e:
#                 print(f"\n处理帧 {processed_frames} 时发生错误: {str(e)}")
#                 processed_frames += 1
#                 continue
#
#     except Exception as e:
#         print(f"\n处理视频时发生错误: {str(e)}")
#
#     finally:
#         # 清理资源
#         print(f"\n处理完成，共处理 {processed_frames} 帧")
#         cap.release()
#         out.release()
#         cv.destroyAllWindows()
#
#         # 删除临时文件
#         if os.path.exists(temp_img_path):
#             try:
#                 os.remove(temp_img_path)
#             except Exception as e:
#                 print(f"删除临时文件时发生错误: {str(e)}")



# 摄像头检测
# if __name__ == "__main__":
#     # 初始化模型
#     yolov8 = YOLOv8Pose(model_path='yolov8n-pose.pt', device='cpu', conf=0.25, iou=0.7)
#
#     # 打开摄像头
#     cap = cv.VideoCapture(0)  # 使用默认摄像头，如果有多个摄像头可以改变参数
#
#     while True:
#         # 读取摄像头帧
#         ret, frame = cap.read()
#         if not ret:
#             print("无法读取摄像头画面")
#             break
#
#         # 保存当前帧到临时文件
#         temp_img_path = "temp_frame.jpg"
#         cv.imwrite(temp_img_path, frame)
#
#         # 进行检测
#         results = yolov8.detect(temp_img_path)
#         results_ = yolov8.get_results_as_dicts(results)
#
#         # 绘制结果
#         frame_with_results = draw_results(frame, results_)
#
#         # 显示结果
#         cv.imshow('实时姿态检测', frame_with_results)
#
#         # 按'q'键退出
#         if cv.waitKey(1) & 0xFF == ord('q'):
#             break
#
#     # 释放资源
#     cap.release()
#     cv.destroyAllWindows()