Yolov8n-pose

说明

Yolov8n-pose 是一款轻量级高效的人体姿态估计算法模型，基于广受欢迎的 YOLO（You Only Look Once）架构开发而成。该模型旨在实现快速且精准的姿态估计，非常适用于实时应用场景，例如：

健身追踪：监测运动过程并提供实时反馈。
虚拟试衣：通过追踪身体动作实现虚拟服装试穿。
安全监控：通过分析人体姿态来检测可疑活动。

修改说明

针对 NPU 的修改

安装

在 Python>=3.9 环境且 PyTorch==2.1.0 的条件下，通过 pip 安装包含所有依赖项的 ultralytics 软件包。

pip install -r examples/requirements.txt

将模型转换为支持NPU

python3 pth2onnx.py --pt models/Yolov8n-pose.pt

source /usr/local/Ascend/ascend-toolkit/set_env.sh
atc --framework=5 \
    --model=./models/Yolov8n-pose.onnx \
    --input_format=NCHW \
    --input_shape='images:1,3,640,640'\
    --output_type=FP32 \
    --output=./models/Yolov8n-pose \
    --log=error \
    --soc_version=Ascend910B2

推理

import argparse
import torch
import numpy as np
from ais_bench.infer.interface import InferSession

from ultralytics.models.yolo.detect import DetectionPredictor
from ultralytics import YOLO

### patch cpu lettebox begin #####
def patch_pre_transform(self, im):
    same_shapes = len({x.shape for x in im}) == 1
    self.model.pt = False 
    letterbox = LetterBox(self.imgsz, auto=same_shapes and self.model.pt, stride=self.model.stride)
    return [letterbox(image=x) for x in im]

from ultralytics.data.augment import LetterBox
from ultralytics.engine.predictor import BasePredictor
BasePredictor.pre_transform = patch_pre_transform
### patch cpu lettebox end #####

class OM_Model():
    def __init__(self):
        self.pt=False
        self.stride = 32
        self.fp16 = False
        self.names = {0: 'person', 1: 'bicycle', 2: 'car', 3: 'motorcycle', 4: 'airplane', 5: 'bus', 6: 'train', 7: 'truck', 8: 'boat', 9: 'traffic light', 10: 'fire hydrant', 11: 'stop sign', 12: 'parking meter', 13: 'bench', 14: 'bird', 15: 'cat', 16: 'dog', 17: 'horse', 18: 'sheep', 19: 'cow', 20: 'elephant', 21: 'bear', 22: 'zebra', 23: 'giraffe', 24: 'backpack', 25: 'umbrella', 26: 'handbag', 27: 'tie', 28: 'suitcase', 29: 'frisbee', 30: 'skis', 31: 'snowboard', 32: 'sports ball', 33: 'kite', 34: 'baseball bat', 35: 'baseball glove', 36: 'skateboard', 37: 'surfboard', 38: 'tennis racket', 39: 'bottle', 40: 'wine glass', 41: 'cup', 42: 'fork', 43: 'knife', 44: 'spoon', 45: 'bowl', 46: 'banana', 47: 'apple', 48: 'sandwich', 49: 'orange', 50: 'broccoli', 51: 'carrot', 52: 'hot dog', 53: 'pizza', 54: 'donut', 55: 'cake', 56: 'chair', 57: 'couch', 58: 'potted plant', 59: 'bed', 60: 'dining table', 61: 'toilet', 62: 'tv', 63: 'laptop', 64: 'mouse', 65: 'remote', 66: 'keyboard', 67: 'cell phone', 68: 'microwave', 69: 'oven', 70: 'toaster', 71: 'sink', 72: 'refrigerator', 73: 'book', 74: 'clock', 75: 'vase', 76: 'scissors', 77: 'teddy bear', 78: 'hair drier', 79: 'toothbrush'}

def pt_detect(input_args):
    model = YOLO(input_args.pt)
    output = model(source=input_args.data, save=True, save_txt=True)  # ndarray
    for i in output:
        print("pt result:", i.boxes.data)

def om_detect(input_args):
    om_model = InferSession(int(input_args.device_id), input_args.om)

    #data
    dp = DetectionPredictor()
    dp.model = OM_Model()
    dp.setup_source(input_args.data)
    for dp.batch in dp.dataset:
        paths, im0s, s = dp.batch

        # Preprocess
        im = dp.preprocess(im0s)

        # Inference
        im = np.ascontiguousarray(im).astype(np.float32)  # contiguous
        preds = om_model.infer([im])
        preds_tensor = torch.from_numpy(preds[0])
        #torch.save(preds_tensor,"om_tensor.pt")
        #print("om result:",preds_tensor, preds_tensor.shape)

        # Postprocess
        output = dp.postprocess(preds_tensor, im, im0s)
        print("om result:", output[0].boxes.data)

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--data', default='./data', help='data path')
    parser.add_argument('--pt', default='./models/yolov8n-pose.pt', help='pt model path')
    parser.add_argument('--om', default='./models/yolov8n-pose.om', help='om model path')
    parser.add_argument('--device_id', default='0', help='device id')
    parser.add_argument('--model_name_or_path', default='Jinan_AICC/Yolov8n-pose', help='Path to model')
    input_args = parser.parse_args()
    
    pt_detect(input_args)
    om_detect(input_args)

if __name__ == '__main__':
    main()

说明

健身追踪：监测运动过程并提供实时反馈。

虚拟试衣：通过追踪身体动作实现虚拟服装试穿。

安全监控：通过分析人体姿态来检测可疑活动。

将模型转换为支持NPU

python3 pth2onnx.py --pt models/Yolov8n-pose.pt

source /usr/local/Ascend/ascend-toolkit/set_env.sh
atc --framework=5 \
    --model=./models/Yolov8n-pose.onnx \
    --input_format=NCHW \
    --input_shape='images:1,3,640,640'\
    --output_type=FP32 \
    --output=./models/Yolov8n-pose \
    --log=error \
    --soc_version=Ascend910B2

推理

import argparse
import torch
import numpy as np
from ais_bench.infer.interface import InferSession

from ultralytics.models.yolo.detect import DetectionPredictor
from ultralytics import YOLO

### patch cpu lettebox begin #####
def patch_pre_transform(self, im):
    same_shapes = len({x.shape for x in im}) == 1
    self.model.pt = False 
    letterbox = LetterBox(self.imgsz, auto=same_shapes and self.model.pt, stride=self.model.stride)
    return [letterbox(image=x) for x in im]

from ultralytics.data.augment import LetterBox
from ultralytics.engine.predictor import BasePredictor
BasePredictor.pre_transform = patch_pre_transform
### patch cpu lettebox end #####

class OM_Model():
    def __init__(self):
        self.pt=False
        self.stride = 32
        self.fp16 = False
        self.names = {0: 'person', 1: 'bicycle', 2: 'car', 3: 'motorcycle', 4: 'airplane', 5: 'bus', 6: 'train', 7: 'truck', 8: 'boat', 9: 'traffic light', 10: 'fire hydrant', 11: 'stop sign', 12: 'parking meter', 13: 'bench', 14: 'bird', 15: 'cat', 16: 'dog', 17: 'horse', 18: 'sheep', 19: 'cow', 20: 'elephant', 21: 'bear', 22: 'zebra', 23: 'giraffe', 24: 'backpack', 25: 'umbrella', 26: 'handbag', 27: 'tie', 28: 'suitcase', 29: 'frisbee', 30: 'skis', 31: 'snowboard', 32: 'sports ball', 33: 'kite', 34: 'baseball bat', 35: 'baseball glove', 36: 'skateboard', 37: 'surfboard', 38: 'tennis racket', 39: 'bottle', 40: 'wine glass', 41: 'cup', 42: 'fork', 43: 'knife', 44: 'spoon', 45: 'bowl', 46: 'banana', 47: 'apple', 48: 'sandwich', 49: 'orange', 50: 'broccoli', 51: 'carrot', 52: 'hot dog', 53: 'pizza', 54: 'donut', 55: 'cake', 56: 'chair', 57: 'couch', 58: 'potted plant', 59: 'bed', 60: 'dining table', 61: 'toilet', 62: 'tv', 63: 'laptop', 64: 'mouse', 65: 'remote', 66: 'keyboard', 67: 'cell phone', 68: 'microwave', 69: 'oven', 70: 'toaster', 71: 'sink', 72: 'refrigerator', 73: 'book', 74: 'clock', 75: 'vase', 76: 'scissors', 77: 'teddy bear', 78: 'hair drier', 79: 'toothbrush'}

def pt_detect(input_args):
    model = YOLO(input_args.pt)
    output = model(source=input_args.data, save=True, save_txt=True)  # ndarray
    for i in output:
        print("pt result:", i.boxes.data)

def om_detect(input_args):
    om_model = InferSession(int(input_args.device_id), input_args.om)

    #data
    dp = DetectionPredictor()
    dp.model = OM_Model()
    dp.setup_source(input_args.data)
    for dp.batch in dp.dataset:
        paths, im0s, s = dp.batch

        # Preprocess
        im = dp.preprocess(im0s)

        # Inference
        im = np.ascontiguousarray(im).astype(np.float32)  # contiguous
        preds = om_model.infer([im])
        preds_tensor = torch.from_numpy(preds[0])
        #torch.save(preds_tensor,"om_tensor.pt")
        #print("om result:",preds_tensor, preds_tensor.shape)

        # Postprocess
        output = dp.postprocess(preds_tensor, im, im0s)
        print("om result:", output[0].boxes.data)

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--data', default='./data', help='data path')
    parser.add_argument('--pt', default='./models/yolov8n-pose.pt', help='pt model path')
    parser.add_argument('--om', default='./models/yolov8n-pose.om', help='om model path')
    parser.add_argument('--device_id', default='0', help='device id')
    parser.add_argument('--model_name_or_path', default='Jinan_AICC/Yolov8n-pose', help='Path to model')
    input_args = parser.parse_args()
    
    pt_detect(input_args)
    om_detect(input_args)

if __name__ == '__main__':
    main()