yolov8的trt模型关键点标注代码

m0_56238089

已于 2023-09-11 10:31:45 修改

阅读量131

点赞数

文章标签： python 深度学习 YOLO

于 2023-09-11 10:23:39 首次发布

本文链接：https://blog.csdn.net/m0_56238089/article/details/132801491

版权

本文详细介绍了如何使用Python和深度学习框架，特别是针对YOLOv8模型，在NVIDIA TensorRT（TRT）上进行关键点检测的标注代码实现过程。

摘要由CSDN通过智能技术生成

from models import TRTModule
from pathlib import Path
import cv2
import torch
import numpy as np
 COLORS = {
    cls: [random.randint(0, 255) for _ in range(3)]
    for i, cls in enumerate(CLASSES)
}
KPS_COLORS= [
    [random.randint(0, 255) for _ in range(3)]
    for n in range(17)
]

import json
def save_json(imagePath,imageHeight,imageWidth, bboxes, scores, kpts,dw, dh,ratio):
    # 创建一个包含标注信息的LabelMe格式字典
    bboxes = bboxes.cpu().numpy().tolist()
    scores = scores.cpu().numpy().tolist()
    kpts = kpts.cpu().numpy().tolist()
    labelme_data = {
        "version": "5.1.1",
        "flags": {},
        "shapes": [],
        "imagePath": "your_image.jpg",  # 替换为您的图像文件名
        "imageData": None,
        "imageHeight": 800,  # 替换为图像高度
        "imageWidth": 800,  # 替换为图像宽度
    }
    labelme_data["imageHeight"] =imageHeight
    labelme_data["imageWidth"] =imageWidth
    labelme_data["imagePath"] =str(imagePath).rsplit("\\")[-1]
    # 遍历每个检测结果，将其添加到LabelMe格式字典中
    for bbox, score, keypoints in zip(bboxes, scores, kpts):
        x1, y1, x2, y2 = bbox

        shape_data = {
            "label": "mouse",  # 对象的标签
            "points": [[int(x1), int(y1)], [int(x2), int(y2)]],  # 边界框的顶点坐标
            "group_id": None,
            "shape_type": "rectangle",
            "flags": {},
        }

        labelme_data["shapes"].append(shape_data)
        for i in range(len(keypoints)):
            points_Data=keypoints[i]
            # print(points_Data)
            probability=poi