devbox
一键难忘520
这个屌丝很懒，什么也没留下！
关注作者
热门标签
热门文章
当前位置:   article > 正文

55、记录使用标注目标检测+关键点检测的数据集，进行训练和部署香橙派芯片_目标检测网络中添加key point

作者：一键难忘520 | 2024-07-18 00:13:32

赞

踩

目标检测网络中添加key point

基本思想：进行是数据集的预处理和标注，然后进行分开目标检测训练和关键点训练，然后增加追踪逻辑

一、标注关键点数据，这里标注只是使用标注检测框和关键点即可，在本质上也可以不用标注检测框，因为个人的数据集是鱼类，可以通过全身的关键点拟合成检测框，提供标注数据集，然后进行提取两类数据集一类为关键点数据集（coco），另一类为检测类数据集（yolov5/7）

 标注完数据之后,进行提取目标检测数据集和关键点数据集，（等待第一版模型出现，做自动化标注，全部以python调用onnx完成开发）代码：

  1. import numpy as np
  2. import json
  3. import glob
  4. import codecs
  5. import os
  6. import cv2
  7. import json
  8. import shutil
  9. import os
  10. from xml.dom.minidom import Document
  11.  
  12.  
  13. class MyEncoder(json.JSONEncoder):
  14.     def default(self, obj):
  15.         if isinstance(obj, np.integer):
  16.             return int(obj)
  17.         elif isinstance(obj, np.floating):
  18.             return float(obj)
  19.         elif isinstance(obj, np.ndarray):
  20.             return obj.tolist()
  21.         else:
  22.             return super(MyEncoder, self).default(obj)
  23.  
  24.  
  25. class Tococo(object):
  26.     def __init__(self, Path, save_path, labelimg_path, keypoints, skeleton, label, verify_flag, verify_path,
  27.                  inital_start_id, verify_labelimg_flag, supercategory):
  28.         self.labelimg_path = labelimg_path
  29.         self.verify_flag = verify_flag
  30.         self.label = label
  31.         categories = {}
  32.         categories.update({"supercategory": supercategory})
  33.         categories.update({"id": 0})
  34.         categories.update({"name": label})
  35.         categories.update({"keypoints": keypoints})
  36.         self.keypoints_length = len(keypoints)
  37.         self.inital_start_id = inital_start_id
  38.         categories.update({"skeleton": skeleton})
  39.         self.categories = [categories]
  40.         self.imagePath = None
  41.         self.imageData = None
  42.         self.imageHeight = None
  43.         self.imageWidth = None
  44.         self.imageName = None
  45.         self.imageName_Noext = None
  46.         self.image = []
  47.         self.annotations = []
  48.         self.jsonfile = glob.glob(os.path.join(Path, "*.json"))
  49.         self.ext = "png"
  50.         self.imgfile = glob.glob(os.path.join(Path, "*." + self.ext))
  51.         if len(self.imgfile) == 0:
  52.             self.imgfile = glob.glob(os.path.join(Path, "*.jpeg"))
  53.             self.ext = "jpeg"
  54.             if len(self.imgfile) == 0:
  55.                 self.imgfile = glob.glob(os.path.join(Path, "*.jpg"))
  56.                 self.ext = "jpg"
  57.         self.save_path = save_path  # 保存json的路径
  58.         self.class_id = label
  59.         self.coco = {}
  60.         self.path = Path
  61.         self.verify_path = verify_path
  62.         self.min_value_x = self.imageWidth
  63.         self.min_value_y = self.imageHeight
  64.         self.max_value_x = 0
  65.         self.max_value_y = 0
  66.         self.verify_labelimg = verify_labelimg_flag
  67.         self.full_img = None
  68.         self.coco_id=0;
  69.  
  70.     def save_label_img(self, all_object):
  71.  
  72.         m_filename = self.image[-1]['file_name']
  73.         print('m_filename=', m_filename)
  74.         m_width = self.image[-1]["width"]
  75.         print('m_width=', m_width)
  76.         m_height = self.image[-1]["height"]
  77.         print('m_height=', m_height)
  78.         object_name = os.path.splitext(m_filename)[0]
  79.         new_object_name = object_name + '.xml'
  80.         img_name = object_name + ".jpg"
  81.  
  82.         print(new_object_name)
  83.         doc = Document()  # 创建DOM文档对象
  84.         DOCUMENT = doc.createElement('annotation')  # 创建根元素
  85.  
  86.         folder = doc.createElement('folder')
  87.         folder_text = doc.createTextNode(m_filename)
  88.         folder.appendChild(folder_text)
  89.         DOCUMENT.appendChild(folder)
  90.         doc.appendChild(DOCUMENT)
  91.  
  92.         filename = doc.createElement('filename')
  93.         filename_text = doc.createTextNode(m_filename)
  94.         filename.appendChild(filename_text)
  95.         DOCUMENT.appendChild(filename)
  96.         doc.appendChild(DOCUMENT)
  97.  
  98.         path = doc.createElement('path')
  99.         path_text = doc.createTextNode(m_filename)
  100.         path.appendChild(path_text)
  101.         DOCUMENT.appendChild(path)
  102.         doc.appendChild(DOCUMENT)
  103.  
  104.         source = doc.createElement('source')
  105.         database = doc.createElement('database')
  106.         database_text = doc.createTextNode("Unknown")  # 元素内容写入
  107.         database.appendChild(database_text)
  108.         source.appendChild(database)
  109.         DOCUMENT.appendChild(source)
  110.         doc.appendChild(DOCUMENT)
  111.  
  112.         size = doc.createElement('size')
  113.         width = doc.createElement('width')
  114.         width_text = doc.createTextNode(str(m_width))  # 元素内容写入
  115.         width.appendChild(width_text)
  116.         size.appendChild(width)
  117.  
  118.         height = doc.createElement('height')
  119.         height_text = doc.createTextNode(str(m_height))
  120.         height.appendChild(height_text)
  121.         size.appendChild(height)
  122.  
  123.         depth = doc.createElement('depth')
  124.         depth_text = doc.createTextNode(str(3))
  125.         depth.appendChild(depth_text)
  126.         size.appendChild(depth)
  127.  
  128.         DOCUMENT.appendChild(size)
  129.  
  130.         segmented = doc.createElement('segmented')
  131.         segmented_text = doc.createTextNode(str(0))
  132.         segmented.appendChild(segmented_text)
  133.         DOCUMENT.appendChild(segmented)
  134.         doc.appendChild(DOCUMENT)
  135.         for item in all_object:
  136.             m_xmin_0 = int(item['bbox'][0])
  137.             print('m_xmin_0=', m_xmin_0)
  138.             m_ymin_0 = int(item['bbox'][1])
  139.             print('m_ymin_0=', m_ymin_0)
  140.             m_xmax_0 = int(item['bbox'][2])
  141.  
  142.             m_ymax_0 = int(item['bbox'][3])
  143.             print('m_ymax_0=', m_ymax_0)
  144.             m_name_0 = self.categories[0]["name"]
  145.             print('m_name_0=', m_name_0)
  146.             object = doc.createElement('object')
  147.             name = doc.createElement('name')
  148.             name_text = doc.createTextNode(m_name_0)
  149.             name.appendChild(name_text)
  150.             object.appendChild(name)
  151.  
  152.             pose = doc.createElement('pose')
  153.             pose_text = doc.createTextNode('Unspecified')
  154.             pose.appendChild(pose_text)
  155.             object.appendChild(pose)
  156.  
  157.             truncated = doc.createElement('truncated')
  158.             truncated_text = doc.createTextNode(str(0))
  159.             truncated.appendChild(truncated_text)
  160.             object.appendChild(truncated)
  161.  
  162.             difficult = doc.createElement('difficult')
  163.             difficult_text = doc.createTextNode(str(0))
  164.             difficult.appendChild(difficult_text)
  165.             object.appendChild(difficult)
  166.  
  167.             bndbox = doc.createElement('bndbox')
  168.             xmin = doc.createElement('xmin')
  169.             xmin_text = doc.createTextNode(str(int(m_xmin_0)))
  170.             xmin.appendChild(xmin_text)
  171.             bndbox.appendChild(xmin)
  172.  
  173.             ymin = doc.createElement('ymin')
  174.             ymin_text = doc.createTextNode(str(int(m_ymin_0)))
  175.             ymin.appendChild(ymin_text)
  176.             bndbox.appendChild(ymin)
  177.  
  178.             xmax = doc.createElement('xmax')
  179.             xmax_text = doc.createTextNode(str(int(m_xmax_0)))
  180.             xmax.appendChild(xmax_text)
  181.             bndbox.appendChild(xmax)
  182.  
  183.             ymax = doc.createElement('ymax')
  184.             ymax_text = doc.createTextNode(str(int(m_ymax_0)))
  185.             ymax.appendChild(ymax_text)
  186.             bndbox.appendChild(ymax)
  187.             object.appendChild(bndbox)
  188.  
  189.             DOCUMENT.appendChild(object)
  190.         new_path_filename = os.path.join(self.labelimg_path, new_object_name)
  191.         print('new_path_filename=', new_path_filename)
  192.         f = open(new_path_filename, 'w')
  193.  
  194.         doc.writexml(f, indent='\t', newl='\n', addindent='\t', encoding='utf-8')
  195.         f.close()
  196.  
  197.     def labelme_to_coco(self):
  198.         for num, json_file in enumerate(self.jsonfile):
  199.             print(json_file)
  200.             json_file = os.path.join(self.path, json_file)
  201.             full_path, json_file_name = os.path.split(json_file)
  202.             json_file_name_Noext, _ = os.path.splitext(json_file_name)
  203.             self.full_img = os.path.join(full_path, ".".join([json_file_name_Noext, self.ext]))
  204.             data = codecs.open(json_file, 'r')
  205.             data = json.load(data)
  206.             self.imageHeight, self.imageWidth = data["imageHeight"], data["imageWidth"]
  207.             _, self.imageName = os.path.split(data["imagePath"])
  208.             self.imageName_Noext, _ = os.path.splitext(self.imageName)
  209.             self.imagePath, self.imageData = data["imagePath"], None
  210.             self.image.append({"file_name": self.imageName, "width": self.imageWidth, "height": self.imageHeight,
  211.                                "id": int(json_file_name_Noext)})
  212.             shapes = data["shapes"]
  213.             segmentations = list(filter(lambda x: x["shape_type"] == "polygon", shapes))
  214.             rects = list(filter(lambda x: x["shape_type"] == "rectangle", shapes))
  215.             shapes = list(filter(lambda x: x["shape_type"] == "point", shapes))
  216.  
  217.             # print(len(shapes))
  218.             if self.verify_flag:
  219.                 _path, _img_name = os.path.split(self.full_img)
  220.                 img = cv2.imread(self.full_img)
  221.                 img_f = os.path.join(self.verify_path, _img_name)
  222.             temp_annotations = []
  223.             for num_id in range(0, int(len(shapes) / self.keypoints_length)):
  224.                 num_keypoints = 0
  225.                 keypoints = [0] * 3 * self.keypoints_length  # 这里是我们标注的关节点个数
  226.                 annotation = {}
  227.                 print(num_id)
  228.                 self.min_value_x = self.imageWidth
  229.                 self.min_value_y = self.imageHeight
  230.                 self.max_value_x = 0
  231.                 self.max_value_y = 0
  232.                 list_shape = shapes[
  233.                              num_id * self.keypoints_length:num_id * self.keypoints_length + self.keypoints_length]
  234.                 # print(shapes[num_id*self.keypoints_length:num_id*self.keypoints_length+self.keypoints_length])
  235.                 segmentation = []
  236.                 for m_id, shape in enumerate(list_shape):
  237.                     if shape["shape_type"] == "point":
  238.                         idx = int(shape['label']) - self.inital_start_id
  239.                         keypoints[idx * 3 + 0] = int(shape['points'][0][0])
  240.                         keypoints[idx * 3 + 1] = int(shape['points'][0][1])
  241.                         keypoints[idx * 3 + 2] = 2
  242.                         segmentation.extend([shape['points'][0][0], shape['points'][0][1]])
  243.                         num_keypoints = num_keypoints + 1
  244.                         self.min_value_x = min(self.min_value_x, shape['points'][0][0])
  245.                         self.min_value_y = min(self.min_value_y, shape['points'][0][1])
  246.                         self.max_value_x = max(self.max_value_x, shape['points'][0][0])
  247.                         self.max_value_y = max(self.max_value_y, shape['points'][0][1])
  248.                         annotation['bbox'] = [self.min_value_x, self.min_value_y, self.max_value_x, self.max_value_y]
  249.                     annotation['segmentation'] = [segmentation]
  250.                 annotation['num_keypoints'] = num_keypoints
  251.                 annotation['area'] = int(annotation['bbox'][2] - annotation['bbox'][0]) * int(
  252.                     annotation['bbox'][3] - annotation['bbox'][1])
  253.  
  254.                 annotation['iscrowd'] = 0
  255.                 annotation['keypoints'] = keypoints
  256.                 if len(rects):
  257.                     for inclde_item in rects:
  258.                         x1 = int(min(inclde_item["points"][0][0], inclde_item["points"][1][0]))
  259.                         y1 = int(min(inclde_item["points"][0][1], inclde_item["points"][1][1]))
  260.                         x2 = int(max(inclde_item["points"][0][0], inclde_item["points"][1][0]))
  261.                         y2 = int(max(inclde_item["points"][0][1], inclde_item["points"][1][1]))
  262.                         include_num = 0
  263.                         for ii in range(0, self.keypoints_length):
  264.                             x, y = int(keypoints[ii * 3]), int(keypoints[ii * 3 + 1])
  265.                             flag = (True if x >= x1 and x <= x2 and y >= y1 and y <= y2 else False)
  266.                             if flag == True:
  267.                                 include_num = include_num + 1
  268.                             if include_num == self.keypoints_length:
  269.                                 annotation['bbox'] = [x1, y1, x2, y2]
  270.  
  271.                 annotation['image_id'] = int(json_file_name_Noext)  # 对应的图片ID
  272.                 annotation['category_id'] = 0
  273.                 print("num=",num_id)
  274.                 annotation['id'] = self.coco_id  # 对象id
  275.                 self.coco_id =self.coco_id +1
  276.                 self.annotations.append(annotation)
  277.                 temp_annotations.append(annotation)
  278.                 self.image_id = int(json_file_name_Noext)
  279.  
  280.                 if self.verify_flag:
  281.                     for item in self.categories[0]["skeleton"]:
  282.                         idx = item[0]
  283.                         iddx = item[1]
  284.                         start_point = (annotation['keypoints'][idx * 3 + 0], annotation['keypoints'][idx * 3 + 1])
  285.                         end_point = (annotation['keypoints'][iddx * 3 + 0], annotation['keypoints'][iddx * 3 + 1])
  286.                         rect_start_point = (int(annotation['bbox'][0]), int(annotation['bbox'][1]))
  287.                         rect_end_point = (int(annotation['bbox'][2]), int(annotation['bbox'][3]))
  288.                         cv2.rectangle(img, rect_start_point, rect_end_point, (255, 0, 255), 1)
  289.                         cv2.putText(img, self.label, rect_start_point, cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 0), 2)
  290.  
  291.                         cv2.circle(img, start_point, 3, (255, 0, 0), 3)
  292.                         cv2.circle(img, end_point, 3, (255, 0, 0), 3)
  293.                         cv2.line(img, start_point, end_point, (0, 255, 0), 2)
  294.                         cv2.putText(img, "{}".format(idx), start_point, cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 28, 0), 2)
  295.                         cv2.imwrite(img_f, img)
  296.             if self.verify_labelimg:
  297.                 self.save_label_img(temp_annotations)
  298.                 shutil.copy(self.full_img, self.labelimg_path)
  299.         self.coco["images"] = self.image
  300.         self.coco["categories"] = self.categories
  301.         self.coco["annotations"] = self.annotations
  302.  
  303.     def get_images(self, filename, height, width):
  304.         image = {}
  305.         image["height"] = height
  306.         image['width'] = width
  307.         image["id"] = self.imagePath
  308.         image["file_name"] = filename
  309.         return image
  310.  
  311.     def get_categories(self, name, class_id):
  312.         category = {}
  313.         category["supercategory"] = self.class_id
  314.         category['id'] = class_id
  315.         category['name'] = name
  316.         return category
  317.  
  318.     def save_coco_json(self):
  319.         self.labelme_to_coco()
  320.         coco_data = self.coco
  321.         # 保存json文件
  322.         json.dump(coco_data, open(self.save_path, 'w'), indent=4, cls=MyEncoder)  # indent=4 更加美观显示
  323.  
  324.  
  325. all_img_json_path = r'/home/ubuntu/Downloads/fish_data/val'
  326. coco_path = r"/home/ubuntu/Downloads/fish_data/val.json"
  327. labelimg_path = r"/home/ubuntu/Downloads/labelimg"
  328. verify_path = r"/home/ubuntu/Downloads/verify"
  329. keypoints = ["0", "1", "2", "3", "4", "5", "6", "7", "8", "9", "10"]
  330. skeleton = [[0, 1], [1, 3], [3, 5], [5, 8], [8, 9], [9, 10], [10, 7], [7, 4], [4, 2], [2, 0], [5, 6], [6, 7], [6, 8],
  331.             [6, 9], [6, 10]]  # 故意写的关键点和关键点骨节数不一样
  332. supercategory = "fish"
  333. label = "fish_tilapia"
  334. verify_flag = False
  335. inital_start_id = 0  # 初始下标开始序号 ，根据标注数据集来确定
  336.  
  337. verify_labelimg_flag = False
  338.  
  339. full_path, json_file_name = os.path.split(coco_path)
  340. if not os.path.exists(full_path):
  341.     os.mkdir(full_path)
  342.  
  343. if not os.path.exists(verify_path):
  344.     os.mkdir(verify_path)
  345.  
  346. if not os.path.exists(labelimg_path):
  347.     os.mkdir(labelimg_path)
  348.  
  349. cocoData = Tococo(all_img_json_path, coco_path, labelimg_path, keypoints, skeleton, label, verify_flag, verify_path,
  350.                   inital_start_id, verify_labelimg_flag, supercategory)
  351. cocoData.save_coco_json()
  352.  
  353. print("finish")

二、测试mmpose得数据集

修改配置文件一/home/ubuntu/mmpose/configs/_base_/datasets/animalpose.py

  1. dataset_info = dict(
  2.     dataset_name='animalpose',
  3.     paper_info=dict(
  4.         author='Cao, Jinkun and Tang, Hongyang and Fang, Hao-Shu and '
  5.         'Shen, Xiaoyong and Lu, Cewu and Tai, Yu-Wing',
  6.         title='Cross-Domain Adaptation for Animal Pose Estimation',
  7.         container='The IEEE International Conference on '
  8.         'Computer Vision (ICCV)',
  9.         year='2019',
  10.         homepage='https://sites.google.com/view/animal-pose/',
  11.     ),
  12.     keypoint_info={
  13.         0:
  14.         dict(
  15.             name='0', id=0, color=[0, 255, 0], type='upper', swap=''),
  16.         1:
  17.         dict(
  18.             name='1',
  19.             id=1,
  20.             color=[255, 128, 0],
  21.             type='upper',
  22.             swap=''),
  23.         2:
  24.         dict(
  25.             name='2',
  26.             id=2,
  27.             color=[0, 255, 0],
  28.             type='upper',
  29.             swap=''),
  30.         3:
  31.         dict(
  32.             name='3',
  33.             id=3,
  34.             color=[255, 128, 0],
  35.             type='upper',
  36.             swap=''),
  37.         4:
  38.         dict(name='4', id=4, color=[51, 153, 255], type='upper', swap=''),
  39.         5:
  40.         dict(name='5', id=5, color=[51, 153, 255], type='upper', swap=''),
  41.         6:
  42.         dict(
  43.             name='6', id=6, color=[51, 153, 255], type='lower',
  44.             swap=''),
  45.         7:
  46.         dict(
  47.             name='7', id=7, color=[51, 153, 255], type='upper', swap=''),
  48.         8:
  49.         dict(
  50.             name='8',
  51.             id=8,
  52.             color=[0, 255, 0],
  53.             type='upper',
  54.             swap=''),
  55.         9:
  56.         dict(
  57.             name='9',
  58.             id=9,
  59.             color=[255, 128, 0],
  60.             type='upper',
  61.             swap=''),
  62.         10:
  63.         dict(
  64.             name='10',
  65.             id=10,
  66.             color=[0, 255, 0],
  67.             type='lower',
  68.             swap='')
  69.     },
  70.     skeleton_info={
  71.         0: dict(link=('0', '1'), id=0, color=[51, 153, 255]),
  72.         1: dict(link=('1', '3'), id=1, color=[0, 255, 0]),
  73.         2: dict(link=('3', '5'), id=2, color=[255, 128, 0]),
  74.         3: dict(link=('5', '8'), id=3, color=[0, 255, 0]),
  75.         4: dict(link=('6', '9'), id=4, color=[255, 128, 0]),
  76.         5: dict(link=('9', '10'), id=5, color=[51, 153, 255]),
  77.         6: dict(link=('10', '7'), id=6, color=[51, 153, 255]),
  78.         7: dict(link=('7', '4'), id=7, color=[51, 153, 255]),
  79.         8: dict(link=('4', '2'), id=8, color=[0, 255, 0]),
  80.         9: dict(link=('2', '0'), id=9, color=[0, 255, 0]),
  81.         10: dict(link=('5', '6'), id=10, color=[0, 255, 0]),
  82.         11: dict(link=('6', '7'), id=11, color=[0, 255, 0]),
  83.         12: dict(link=('6', '10'), id=12, color=[0, 255, 0]),
  84.         13: dict(link=('6', '9'), id=13, color=[0, 255, 0]),
  85.         14: dict(link=('6', '8'), id=14, color=[0, 255, 0]),
  86.         15: dict(link=('8', '9'), id=11, color=[0, 255, 0]),
  87.  
  88.     },
  89.     joint_weights=[
  90.         1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.,1
  91.     ],
  92.  
  93.     # Note: The original paper did not provide enough information about
  94.     # the sigmas. We modified from 'https://github.com/cocodataset/'
  95.     # 'cocoapi/blob/master/PythonAPI/pycocotools/cocoeval.py#L523'
  96.     sigmas=[
  97.         1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.
  98.     ])

修改配置文件二/home/ubuntu/mmpose/configs/animal_2d_keypoint/topdown_heatmap/animalpose/td-hm_hrnet-w32_8xb64-210e_animalpose-256x256.py

  1. _base_ = ['../../../_base_/default_runtime.py']
  2.  
  3. # runtime
  4. train_cfg = dict(max_epochs=1000, val_interval=50)
  5.  
  6. # optimizer
  7. optim_wrapper = dict(optimizer=dict(
  8.     type='Adam',
  9.     lr=5e-4,
  10. ))
  11.  
  12. # learning policy
  13. param_scheduler = [
  14.     dict(
  15.         type='LinearLR', begin=0, end=500, start_factor=0.001,
  16.         by_epoch=False),  # warm-up
  17.     dict(
  18.         type='MultiStepLR',
  19.         begin=0,
  20.         end=210,
  21.         milestones=[170, 200],
  22.         gamma=0.1,
  23.         by_epoch=True)
  24. ]
  25.  
  26. # automatically scaling LR based on the actual training batch size
  27. auto_scale_lr = dict(base_batch_size=512)
  28.  
  29. # hooks
  30. default_hooks = dict(checkpoint=dict(save_best='coco/AP', rule='greater'))
  31.  
  32. # codec settings
  33. codec = dict(
  34.     type='MSRAHeatmap', input_size=(256, 256), heatmap_size=(64, 64), sigma=2)
  35.  
  36. # model settings
  37. model = dict(
  38.     type='TopdownPoseEstimator',
  39.     data_preprocessor=dict(
  40.         type='PoseDataPreprocessor',
  41.         mean=[123.675, 116.28, 103.53],
  42.         std=[58.395, 57.12, 57.375],
  43.         bgr_to_rgb=True),
  44.     backbone=dict(
  45.         type='HRNet',
  46.         in_channels=3,
  47.         extra=dict(
  48.             stage1=dict(
  49.                 num_modules=1,
  50.                 num_branches=1,
  51.                 block='BOTTLENECK',
  52.                 num_blocks=(4, ),
  53.                 num_channels=(64, )),
  54.             stage2=dict(
  55.                 num_modules=1,
  56.                 num_branches=2,
  57.                 block='BASIC',
  58.                 num_blocks=(4, 4),
  59.                 num_channels=(32, 64)),
  60.             stage3=dict(
  61.                 num_modules=4,
  62.                 num_branches=3,
  63.                 block='BASIC',
  64.                 num_blocks=(4, 4, 4),
  65.                 num_channels=(32, 64, 128)),
  66.             stage4=dict(
  67.                 num_modules=3,
  68.                 num_branches=4,
  69.                 block='BASIC',
  70.                 num_blocks=(4, 4, 4, 4),
  71.                 num_channels=(32, 64, 128, 256))),
  72.         init_cfg=dict(
  73.             type='Pretrained',
  74.             checkpoint='https://download.openmmlab.com/mmpose/'
  75.             'pretrain_models/hrnet_w32-36af842e.pth'),
  76.     ),
  77.     head=dict(
  78.         type='HeatmapHead',
  79.         in_channels=32,
  80.         out_channels=11,
  81.         deconv_out_channels=None,
  82.         loss=dict(type='KeypointMSELoss', use_target_weight=True),
  83.         decoder=codec),
  84.     test_cfg=dict(
  85.         flip_test=True,
  86.         flip_mode='heatmap',
  87.         shift_heatmap=True,
  88.     ))
  89.  
  90. # base dataset settings
  91. dataset_type = 'AnimalPoseDataset'
  92. data_mode = 'topdown'
  93. data_root = '/home/ubuntu/mmpose/fish_data/'
  94.  
  95. # pipelines
  96. train_pipeline = [
  97.     dict(type='LoadImage'),
  98.     dict(type='GetBBoxCenterScale'),
  99.     dict(type='RandomFlip', direction='horizontal'),
  100.     dict(type='RandomHalfBody'),
  101.     dict(type='RandomBBoxTransform'),
  102.     dict(type='TopdownAffine', input_size=codec['input_size']),
  103.     dict(type='GenerateTarget', encoder=codec),
  104.     dict(type='PackPoseInputs')
  105. ]
  106. val_pipeline = [
  107.     dict(type='LoadImage'),
  108.     dict(type='GetBBoxCenterScale'),
  109.     dict(type='TopdownAffine', input_size=codec['input_size']),
  110.     dict(type='PackPoseInputs')
  111. ]
  112.  
  113. # data loaders
  114. train_dataloader = dict(
  115.     batch_size=32,
  116.     num_workers=1,
  117.     persistent_workers=True,
  118.     sampler=dict(type='DefaultSampler', shuffle=True),
  119.     dataset=dict(
  120.         type=dataset_type,
  121.         data_root=data_root,
  122.         data_mode=data_mode,
  123.         ann_file='train.json',
  124.         data_prefix=dict(img='train/'),
  125.         pipeline=train_pipeline,
  126.     ))
  127. val_dataloader = dict(
  128.     batch_size=32,
  129.     num_workers=1,
  130.     persistent_workers=True,
  131.     drop_last=False,
  132.     sampler=dict(type='DefaultSampler', shuffle=False, round_up=False),
  133.     dataset=dict(
  134.         type=dataset_type,
  135.         data_root=data_root,
  136.         data_mode=data_mode,
  137.         ann_file='val.json',
  138.         data_prefix=dict(img='val/'),
  139.         test_mode=True,
  140.         pipeline=val_pipeline,
  141.     ))
  142. test_dataloader = val_dataloader
  143.  
  144. # evaluators
  145. val_evaluator = dict(
  146.     type='CocoMetric', ann_file=data_root + 'val.json')
  147. test_evaluator = val_evaluator

测试数据集是否正确

ubuntu@ubuntu:~/mmpose$ python3 tools/misc/browse_dataset.py /home/ubuntu/mmpose/configs/animal_2d_keypoint/topdown_heatmap/animalpose/td-hm_hrnet-w32_8xb64-210e_animalpose-256x256.py --phase val --mode original

三、训练数据集

ubuntu@ubuntu:~/mmpose$ python3 tools/train.py /home/ubuntu/mmpose/configs/animal_2d_keypoint/topdown_heatmap/animalpose/td-hm_hrnet-w32_8xb64-210e_animalpose-256x256.py

四、是用mmdeploy转模型

  1. ubuntu@ubuntu:~/mmdeploy$ python3 tools/torch2onnx.py configs/mmpose/pose-detection_onnxruntime_static.py ../mmpose/configs/animal_2d_keypoint/topdown_heatmap/animalpose/td-hm_hrnet-w32_8xb64-210e_animalpose-256x256.py ../mmpose/work_dirs/td-hm_hrnet-w32_8xb64-210e_animalpose-256x256/epoch_670.pth ../mmpose/tests/data/animalpose/ca110.jpeg
  2. 04/24 17:19:28 - mmengine - INFO - torch2onnx: 
  3. 	model_cfg: ../mmpose/configs/animal_2d_keypoint/topdown_heatmap/animalpose/td-hm_hrnet-w32_8xb64-210e_animalpose-256x256.py 
  4. 	deploy_cfg: configs/mmpose/pose-detection_onnxruntime_static.py
  5. 04/24 17:19:29 - mmengine - WARNING - Failed to search registry with scope "mmpose" in the "Codebases" registry tree. As a workaround, the current "Codebases" registry in "mmdeploy" is used to build instance. This may cause unexpected failure when running the built modules. Please check whether "mmpose" is a correct scope, or whether the registry is initialized.
  6. 04/24 17:19:29 - mmengine - WARNING - Failed to search registry with scope "mmpose" in the "mmpose_tasks" registry tree. As a workaround, the current "mmpose_tasks" registry in "mmdeploy" is used to build instance. This may cause unexpected failure when running the built modules. Please check whether "mmpose" is a correct scope, or whether the registry is initialized.
  7. Loads checkpoint by local backend from path: ../mmpose/work_dirs/td-hm_hrnet-w32_8xb64-210e_animalpose-256x256/epoch_670.pth
  8. /home/ubuntu/mmpose/mmpose/datasets/datasets/utils.py:102: UserWarning: The metainfo config file "configs/_base_/datasets/animalpose.py" does not exist. A matched config file "/home/ubuntu/mmpose/mmpose/.mim/configs/_base_/datasets/animalpose.py" will be used instead.
  9.   warnings.warn(
  10. 04/24 17:19:31 - mmengine - WARNING - DeprecationWarning: get_onnx_config will be deprecated in the future. 
  11. 04/24 17:19:31 - mmengine - INFO - Export PyTorch model to ONNX: ./work-dir/end2end.onnx.
  12. 04/24 17:19:31 - mmengine - WARNING - Can not find torch._C._jit_pass_onnx_autograd_function_process, function rewrite will not be applied
  13. 04/24 17:20:10 - mmengine - INFO - Execute onnx optimize passes.
  14. 04/24 17:20:10 - mmengine - WARNING - Can not optimize model, please build torchscipt extension.
  15. More details: https://github.com/open-mmlab/mmdeploy/tree/1.x/docs/en/experimental/onnx_optimizer.md
  16. 04/24 17:20:10 - mmengine - INFO - torch2onnx finished. Results saved to ./work-dir

转ncnn

ubuntu@ubuntu:~/ncnn/build/install/bin$ ./onnx2ncnn /home/ubuntu/mmdeploy/work-dir/end2end.onnx /home/ubuntu/mmdeploy/work-dir/end2end.param /home/ubuntu/mmdeploy/work-dir/end2end.bin

五、转rk3588,产生txt

ubuntu@ubuntu:~/mmpose$ find /home/ubuntu/mmpose/fish_data -name "*.jpg" > a.txt

转模型

  1. from rknn.api import RKNN
  2.  
  3. ONNX_MODEL = '/home/ubuntu/mmdeploy/work-dir/end2end.onnx'
  4. RKNN_MODEL = '/home/ubuntu/mmdeploy/work-dir/end2end.rknn'
  5.  
  6. if __name__ == '__main__':
  7.  
  8.     # Create RKNN object
  9.     rknn = RKNN(verbose=True)
  10.  
  11.     # pre-process config
  12.     print('--> config model')
  13.     rknn.config(mean_values=[[0, 0, 0]], std_values=[[255, 255, 255]],
  14.                 target_platform='rk3588',
  15.                 quantized_dtype='asymmetric_quantized-8',  optimization_level=3)
  16.     print('done')
  17.  
  18.     print('--> Loading model')
  19.     ret = rknn.load_onnx(model=ONNX_MODEL)
  20.     if ret != 0:
  21.         print('Load model  failed!')
  22.         exit(ret)
  23.     print('done')
  24.  
  25.     # Build model
  26.     print('--> Building model')
  27.     ret = rknn.build(do_quantization=True, dataset='dataset.txt')  # ,pre_compile=True
  28.     if ret != 0:
  29.         print('Build end2end failed!')
  30.         exit(ret)
  31.     print('done')
  32.  
  33.     # Export rknn model
  34.     print('--> Export RKNN model')
  35.     ret = rknn.export_rknn(RKNN_MODEL)
  36.     if ret != 0:
  37.         print('Export end2end.rknn failed!')
  38.         exit(ret)
  39.     print('done')
  40.  
  41.     rknn.release()

cmakelists.txt

  1. cmake_minimum_required(VERSION 3.16)
  2. project(untitled10)
  3. set(CMAKE_CXX_FLAGS "-std=c++11")
  4. set(CMAKE_CXX_STANDARD 11)
  5. set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS}  ")
  6. set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ")
  7.  
  8. include_directories(${CMAKE_SOURCE_DIR})
  9. include_directories(${CMAKE_SOURCE_DIR}/include)
  10. find_package(OpenCV REQUIRED)
  11. #message(STATUS ${OpenCV_INCLUDE_DIRS})
  12. #添加头文件
  13. include_directories(${OpenCV_INCLUDE_DIRS})
  14. #链接Opencv库
  15. add_library(librknn_api SHARED IMPORTED)
  16. set_target_properties(librknn_api PROPERTIES IMPORTED_LOCATION ${CMAKE_SOURCE_DIR}/lib/librknn_api.so)
  17. #官方的zoo社区或者参考前几篇博客寻找其so出处
  18.  
  19. add_executable(untitled10 main.cpp)
  20. target_link_libraries(untitled10 ${OpenCV_LIBS} librknn_api )

main.cpp

  1. #include <stdio.h>
  2. #include <stdint.h>
  3. #include <stdlib.h>
  4. #include <queue>
  5. #include "rknn_api.h"
  6. #include "opencv2/core/core.hpp"
  7. #include "opencv2/imgproc/imgproc.hpp"
  8. #include "opencv2/highgui/highgui.hpp"
  9. #include <chrono>
  10. #include <iostream>
  11.  
  12. using namespace std;
  13.  
  14. struct Keypoints {
  15.     float x;
  16.     float y;
  17.     float score;
  18.  
  19.     Keypoints() : x(0), y(0), score(0) {}
  20.  
  21.     Keypoints(float x, float y, float score) : x(x), y(y), score(score) {}
  22. };
  23.  
  24. struct Box {
  25.     float center_x;
  26.     float center_y;
  27.     float scale_x;
  28.     float scale_y;
  29.     float scale_prob;
  30.     float score;
  31.  
  32.     Box() : center_x(0), center_y(0), scale_x(0), scale_y(0), scale_prob(0), score(0) {}
  33.  
  34.     Box(float center_x, float center_y, float scale_x, float scale_y, float scale_prob, float score) :
  35.             center_x(center_x), center_y(center_y), scale_x(scale_x), scale_y(scale_y), scale_prob(scale_prob),
  36.             score(score) {}
  37. };
  38.  
  39. void bbox_xywh2cs(float bbox[], float aspect_ratio, float padding, float pixel_std, float *center, float *scale) {
  40.     float x = bbox[0];
  41.     float y = bbox[1];
  42.     float w = bbox[2];
  43.     float h = bbox[3];
  44.     *center = x + w * 0.5;
  45.     *(center + 1) = y + h * 0.5;
  46.     if (w > aspect_ratio * h)
  47.         h = w * 1.0 / aspect_ratio;
  48.     else if (w < aspect_ratio * h)
  49.         w = h * aspect_ratio;
  50.  
  51.  
  52.     *scale = (w / pixel_std) * padding;
  53.     *(scale + 1) = (h / pixel_std) * padding;
  54. }
  55.  
  56. void rotate_point(float *pt, float angle_rad, float *rotated_pt) {
  57.     float sn = sin(angle_rad);
  58.     float cs = cos(angle_rad);
  59.     float new_x = pt[0] * cs - pt[1] * sn;
  60.     float new_y = pt[0] * sn + pt[1] * cs;
  61.     rotated_pt[0] = new_x;
  62.     rotated_pt[1] = new_y;
  63.  
  64. }
  65.  
  66. void _get_3rd_point(cv::Point2f a, cv::Point2f b, float *direction) {
  67.  
  68.     float direction_0 = a.x - b.x;
  69.     float direction_1 = a.y - b.y;
  70.     direction[0] = b.x - direction_1;
  71.     direction[1] = b.y + direction_0;
  72.  
  73.  
  74. }
  75.  
  76. void get_affine_transform(float *center, float *scale, float rot, float *output_size, float *shift, bool inv,
  77.                           cv::Mat &trans) {
  78.     float scale_tmp[] = {0, 0};
  79.     scale_tmp[0] = scale[0] * 200.0;
  80.     scale_tmp[1] = scale[1] * 200.0;
  81.     float src_w = scale_tmp[0];
  82.     float dst_w = output_size[0];
  83.     float dst_h = output_size[1];
  84.     float rot_rad = M_PI * rot / 180;
  85.     float pt[] = {0, 0};
  86.     pt[0] = 0;
  87.     pt[1] = src_w * (-0.5);
  88.     float src_dir[] = {0, 0};
  89.     rotate_point(pt, rot_rad, src_dir);
  90.     float dst_dir[] = {0, 0};
  91.     dst_dir[0] = 0;
  92.     dst_dir[1] = dst_w * (-0.5);
  93.     cv::Point2f src[3] = {cv::Point2f(0, 0), cv::Point2f(0, 0), cv::Point2f(0, 0)};
  94.     src[0] = cv::Point2f(center[0] + scale_tmp[0] * shift[0], center[1] + scale_tmp[1] * shift[1]);
  95.     src[1] = cv::Point2f(center[0] + src_dir[0] + scale_tmp[0] * shift[0],
  96.                          center[1] + src_dir[1] + scale_tmp[1] * shift[1]);
  97.     float direction_src[] = {0, 0};
  98.     _get_3rd_point(src[0], src[1], direction_src);
  99.     src[2] = cv::Point2f(direction_src[0], direction_src[1]);
  100.     cv::Point2f dst[3] = {cv::Point2f(0, 0), cv::Point2f(0, 0), cv::Point2f(0, 0)};
  101.     dst[0] = cv::Point2f(dst_w * 0.5, dst_h * 0.5);
  102.     dst[1] = cv::Point2f(dst_w * 0.5 + dst_dir[0], dst_h * 0.5 + dst_dir[1]);
  103.     float direction_dst[] = {0, 0};
  104.     _get_3rd_point(dst[0], dst[1], direction_dst);
  105.     dst[2] = cv::Point2f(direction_dst[0], direction_dst[1]);
  106.  
  107.     if (inv) {
  108.         trans = cv::getAffineTransform(dst, src);
  109.     } else {
  110.         trans = cv::getAffineTransform(src, dst);
  111.     }
  112.  
  113.  
  114. }
  115.  
  116. void
  117. transform_preds(std::vector<cv::Point2f> coords, std::vector<Keypoints> &target_coords, float *center, float *scale,
  118.                 int w, int h, bool use_udp = false) {
  119.     float scale_x[] = {0, 0};
  120.     float temp_scale[] = {scale[0] * 200, scale[1] * 200};
  121.     if (use_udp) {
  122.         scale_x[0] = temp_scale[0] / (w - 1);
  123.         scale_x[1] = temp_scale[1] / (h - 1);
  124.     } else {
  125.         scale_x[0] = temp_scale[0] / w;
  126.         scale_x[1] = temp_scale[1] / h;
  127.     }
  128.     for (int i = 0; i < coords.size(); i++) {
  129.         target_coords[i].x = coords[i].x * scale_x[0] + center[0] - temp_scale[0] * 0.5;
  130.         target_coords[i].y = coords[i].y * scale_x[1] + center[1] - temp_scale[1] * 0.5;
  131.     }
  132.  
  133. }
  134.  
  135.  
  136. void letterbox(cv::Mat rgb, cv::Mat &img_resize, int target_width, int target_height) {
  137.  
  138.     float shape_0 = rgb.rows;
  139.     float shape_1 = rgb.cols;
  140.     float new_shape_0 = target_height;
  141.     float new_shape_1 = target_width;
  142.     float r = std::min(new_shape_0 / shape_0, new_shape_1 / shape_1);
  143.     float new_unpad_0 = int(round(shape_1 * r));
  144.     float new_unpad_1 = int(round(shape_0 * r));
  145.     float dw = new_shape_1 - new_unpad_0;
  146.     float dh = new_shape_0 - new_unpad_1;
  147.     dw = dw / 2;
  148.     dh = dh / 2;
  149.     cv::Mat copy_rgb = rgb.clone();
  150.     if (int(shape_0) != int(new_unpad_0) && int(shape_1) != int(new_unpad_1)) {
  151.         cv::resize(copy_rgb, img_resize, cv::Size(new_unpad_0, new_unpad_1));
  152.         copy_rgb = img_resize;
  153.     }
  154.     int top = int(round(dh - 0.1));
  155.     int bottom = int(round(dh + 0.1));
  156.     int left = int(round(dw - 0.1));
  157.     int right = int(round(dw + 0.1));
  158.     cv::copyMakeBorder(copy_rgb, img_resize, top, bottom, left, right, cv::BORDER_CONSTANT, cv::Scalar(0, 0, 0));
  159.  
  160. }
  161.  
  162. void printRKNNTensor(rknn_tensor_attr *attr) {
  163.     printf("index=%d name=%s n_dims=%d dims=[%d %d %d %d] n_elems=%d size=%d "
  164.            "fmt=%d type=%d qnt_type=%d fl=%d zp=%d scale=%f\n",
  165.            attr->index, attr->name, attr->n_dims, attr->dims[3], attr->dims[2],
  166.            attr->dims[1], attr->dims[0], attr->n_elems, attr->size, 0, attr->type,
  167.            attr->qnt_type, attr->fl, attr->zp, attr->scale);
  168. }
  169.  
  170. int post_process_u8(uint8_t *input0, int model_in_h, int model_in_w) {
  171.  
  172.  
  173.     return 0;
  174. }
  175.  
  176. int main(int argc, char **argv) {
  177.     float keypoint_score = 0.1f;
  178.     cv::Mat bgr = cv::imread("../7.jpg");
  179.     cv::Mat rgb;
  180.     cv::cvtColor(bgr, rgb, cv::COLOR_BGR2RGB);
  181.     float image_target_w = 256;
  182.     float image_target_h = 256;
  183.     float padding = 1.25;
  184.     float pixel_std = 200;
  185.     float aspect_ratio = image_target_h / image_target_w;
  186.     float bbox[] = {516,210,516+702,210+468,
  187.                     9.995332e-01};// 需要检测框架　这个矩形框来自检测框架的坐标　x y w h score　这里来自我自己标注的检测框
  188. //    bbox[2] = bbox[2] - bbox[0];
  189. //    bbox[3] = bbox[3] - bbox[1];
  190.     float center[2] = {0, 0};
  191.     float scale[2] = {0, 0};
  192.     bbox_xywh2cs(bbox, aspect_ratio, padding, pixel_std, center, scale);
  193.     float rot = 0;
  194.     float shift[] = {0, 0};
  195.     bool inv = false;
  196.     float output_size[] = {image_target_h, image_target_w};
  197.     cv::Mat trans;
  198.     get_affine_transform(center, scale, rot, output_size, shift, inv, trans);
  199.     std::cout << trans << std::endl;
  200.     cv::Mat detect_image;//= cv::Mat::zeros(image_target_w ,image_target_h, CV_8UC3);
  201.     cv::warpAffine(rgb, detect_image, trans, cv::Size(image_target_h, image_target_w), cv::INTER_LINEAR);
  202.     const char *model_path = "../end2end_3588.rknn";
  203.     // Load model
  204.     FILE *fp = fopen(model_path, "rb");
  205.     if (fp == NULL) {
  206.         printf("fopen %s fail!\n", model_path);
  207.         return -1;
  208.     }
  209.     fseek(fp, 0, SEEK_END);
  210.     int model_len = ftell(fp);
  211.     void *model = malloc(model_len);
  212.     fseek(fp, 0, SEEK_SET);
  213.     if (model_len != fread(model, 1, model_len, fp)) {
  214.         printf("fread %s fail!\n", model_path);
  215.         free(model);
  216.         return -1;
  217.     }
  218.  
  219.  
  220.     rknn_context ctx = 0;
  221.  
  222.     int ret = rknn_init(&ctx, model, model_len, 0,NULL);
  223.     if (ret < 0) {
  224.         printf("rknn_init fail! ret=%d\n", ret);
  225.         return -1;
  226.     }
  227.  
  228.     /* Query sdk version */
  229.     rknn_sdk_version version;
  230.     ret = rknn_query(ctx, RKNN_QUERY_SDK_VERSION, &version,
  231.                      sizeof(rknn_sdk_version));
  232.     if (ret < 0) {
  233.         printf("rknn_init error ret=%d\n", ret);
  234.         return -1;
  235.     }
  236.     printf("sdk version: %s driver version: %s\n", version.api_version,
  237.            version.drv_version);
  238.  
  239.  
  240.     /* Get input,output attr */
  241.     rknn_input_output_num io_num;
  242.     ret = rknn_query(ctx, RKNN_QUERY_IN_OUT_NUM, &io_num, sizeof(io_num));
  243.     if (ret < 0) {
  244.         printf("rknn_init error ret=%d\n", ret);
  245.         return -1;
  246.     }
  247.     printf("model input num: %d, output num: %d\n", io_num.n_input,
  248.            io_num.n_output);
  249.  
  250.     rknn_tensor_attr input_attrs[io_num.n_input];
  251.     memset(input_attrs, 0, sizeof(input_attrs));
  252.     for (int i = 0; i < io_num.n_input; i++) {
  253.         input_attrs[i].index = i;
  254.         ret = rknn_query(ctx, RKNN_QUERY_INPUT_ATTR, &(input_attrs[i]),
  255.                          sizeof(rknn_tensor_attr));
  256.         if (ret < 0) {
  257.             printf("rknn_init error ret=%d\n", ret);
  258.             return -1;
  259.         }
  260.         printRKNNTensor(&(input_attrs[i]));
  261.     }
  262.  
  263.     rknn_tensor_attr output_attrs[io_num.n_output];
  264.     memset(output_attrs, 0, sizeof(output_attrs));
  265.     for (int i = 0; i < io_num.n_output; i++) {
  266.         output_attrs[i].index = i;
  267.         ret = rknn_query(ctx, RKNN_QUERY_OUTPUT_ATTR, &(output_attrs[i]),
  268.                          sizeof(rknn_tensor_attr));
  269.         printRKNNTensor(&(output_attrs[i]));
  270.     }
  271.  
  272.     int input_channel = 3;
  273.     int input_width = 0;
  274.     int input_height = 0;
  275.     if (input_attrs[0].fmt == RKNN_TENSOR_NCHW) {
  276.         printf("model is NCHW input fmt\n");
  277.         input_width = input_attrs[0].dims[0];
  278.         input_height = input_attrs[0].dims[1];
  279.         printf("input_width=%d input_height=%d\n", input_width, input_height);
  280.     } else {
  281.         printf("model is NHWC input fmt\n");
  282.         input_width = input_attrs[0].dims[1];
  283.         input_height = input_attrs[0].dims[2];
  284.         printf("input_width=%d input_height=%d\n", input_width, input_height);
  285.     }
  286.  
  287.     printf("model input height=%d, width=%d, channel=%d\n", input_height, input_width,
  288.            input_channel);
  289.  
  290.  
  291. /* Init input tensor */
  292.     rknn_input inputs[1];
  293.     memset(inputs, 0, sizeof(inputs));
  294.     inputs[0].index = 0;
  295.     inputs[0].buf = detect_image.data;
  296.     inputs[0].type = RKNN_TENSOR_UINT8;
  297.     inputs[0].size = input_width * input_height * input_channel;
  298.     inputs[0].fmt = RKNN_TENSOR_NHWC;
  299.     inputs[0].pass_through = 0;
  300.  
  301.     /* Init output tensor */
  302.     rknn_output outputs[io_num.n_output];
  303.     memset(outputs, 0, sizeof(outputs));
  304.     for (int i = 0; i < io_num.n_output; i++) {
  305.         outputs[i].want_float = 1;
  306.     }
  307.     printf("img.cols: %d, img.rows: %d\n", detect_image.cols, detect_image.rows);
  308.     auto t1 = std::chrono::steady_clock::now();
  309.     rknn_inputs_set(ctx, io_num.n_input, inputs);
  310.     std::chrono::steady_clock::time_point  now = std::chrono::steady_clock::now();
  311.  
  312.     ret = rknn_run(ctx, NULL);
  313.     if (ret < 0) {
  314.         printf("ctx error ret=%d\n", ret);
  315.         return -1;
  316.     }
  317.     auto t2 = std::chrono::steady_clock::now();
  318.     std::chrono::duration<double> time_span = std::chrono::duration_cast<std::chrono::duration<double>>(t2 - now);
  319.     std::cout << "It took me " << time_span.count()*1000 << " ms."<<std::endl;
  320.     ret = rknn_outputs_get(ctx, io_num.n_output, outputs, NULL);
  321.     if (ret < 0) {
  322.         printf("outputs error ret=%d\n", ret);
  323.         return -1;
  324.     }
  325.     int shape_b = 0;
  326.     int shape_c = 0;
  327.     int shape_w = 0;;
  328.     int shape_h = 0;;
  329.     for (int i = 0; i < io_num.n_output; ++i) {
  330.         shape_b = output_attrs[i].dims[0];
  331.         shape_c = output_attrs[i].dims[1];
  332.         shape_h = output_attrs[i].dims[2];;
  333.         shape_w = output_attrs[i].dims[3];;
  334.     }
  335.     printf("batch=%d channel=%d width=%d height= %d\n", shape_b, shape_c, shape_w, shape_h);
  336.     std::vector<float> vec_result_heap;
  337.     float *output = (float *) outputs[0].buf;
  338.     for (int i = 0; i < shape_c; i++) {
  339.         for (int j = 0; j < shape_h; j++) {
  340.             for (int k = 0; k < shape_w; k++) {
  341.                 float elements = output[i * shape_w * shape_h + j * shape_w + k];
  342.                 vec_result_heap.emplace_back(elements);
  343.             }
  344.         }
  345.  
  346.     }
  347.  
  348.     std::vector<Keypoints> all_preds;
  349.     std::vector<int> idx;
  350.     for (int i = 0; i < shape_c; i++) {
  351.         auto begin = vec_result_heap.begin() + i * shape_w * shape_h;
  352.         auto end = vec_result_heap.begin() + (i + 1) * shape_w * shape_h;
  353.         float maxValue = *max_element(begin, end);
  354.         int maxPosition = max_element(begin, end) - begin;
  355.         all_preds.emplace_back(Keypoints(0, 0, maxValue));
  356.         idx.emplace_back(maxPosition);
  357.     }
  358.     std::vector<cv::Point2f> vec_point;
  359.     for (int i = 0; i < idx.size(); i++) {
  360.         int x = idx[i] % shape_w;
  361.         int y = idx[i] / shape_w;
  362.         vec_point.emplace_back(cv::Point2f(x, y));
  363.     }
  364.  
  365.     for (int i = 0; i < shape_c; i++) {
  366.         int px = vec_point[i].x;
  367.         int py = vec_point[i].y;
  368.         if (px > 1 && px < shape_w - 1 && py > 1 && py < shape_h - 1) {
  369.             float diff_0 = vec_result_heap[py * shape_w + px + 1] - vec_result_heap[py * shape_w + px - 1];
  370.             float diff_1 = vec_result_heap[(py + 1) * shape_w + px] - vec_result_heap[(py - 1) * shape_w + px];
  371.             vec_point[i].x += diff_0 == 0 ? 0 : (diff_0 > 0) ? 0.25 : -0.25;
  372.             vec_point[i].y += diff_1 == 0 ? 0 : (diff_1 > 0) ? 0.25 : -0.25;
  373.         }
  374.     }
  375.     std::vector<Box> all_boxes;
  376.     bool heap_map = false;
  377.     if (heap_map) {
  378.         all_boxes.emplace_back(Box(center[0], center[1], scale[0], scale[1], scale[0] * scale[1] * 400, bbox[4]));
  379.     }
  380.     transform_preds(vec_point, all_preds, center, scale, shape_w, shape_h);
  381.     int skeleton[15][2] = {{0, 1},{1, 3},{3, 5},{5, 8},
  382.                            {8, 9},{9,  10},{10,  7},{7,4},{4,  2},
  383.                            {2,  0},{5,  6},
  384.                            {6,  7},{6,  10},{6,  8},{6,  9}};
  385.  
  386.     cv::rectangle(bgr, cv::Point(bbox[0], bbox[1]), cv::Point(bbox[0] + bbox[2], bbox[1] + bbox[3]),
  387.                   cv::Scalar(255, 0, 0));
  388.     for (int i = 0; i < all_preds.size(); i++) {
  389.         if (all_preds[i].score > keypoint_score) {
  390.             cv::circle(bgr, cv::Point(all_preds[i].x, all_preds[i].y), 3, cv::Scalar(0, 255, 120), -1);//画点，其实就是实心圆
  391.         }
  392.     }
  393.     for (int i = 0; i < sizeof(skeleton) / sizeof(sizeof(skeleton[1])); i++) {
  394.         int x0 = all_preds[skeleton[i][0]].x;
  395.         int y0 = all_preds[skeleton[i][0]].y;
  396.         int x1 = all_preds[skeleton[i][1]].x;
  397.         int y1 = all_preds[skeleton[i][1]].y;
  398.  
  399.         cv::line(bgr, cv::Point(x0, y0), cv::Point(x1, y1),
  400.                  cv::Scalar(0, 255, 0), 1);
  401.     }
  402.     cv::imwrite("../image.jpg", bgr);
  403.     ret = rknn_outputs_release(ctx, io_num.n_output, outputs);
  404.  
  405.     if (ret < 0) {
  406.         printf("rknn_query fail! ret=%d\n", ret);
  407.         goto Error;
  408.     }
  409.  
  410.     Error:
  411.     if (ctx > 0)
  412.         ret=rknn_destroy(ctx);
  413.     printf("%d \n",ret);
  414.     //感觉官方固件有问题，销毁对象存在问题，或许是我的固件有问题　https://t.rock-chips.com/forum.php?mod=viewthread&tid=2365
  415.     if (model)
  416.         free(model);
  417.     if (fp)
  418.         fclose(fp);
  419.     return 0;
  420. }

测试结果

附录 mp4转h264(瑞芯微用第二种)

  1. ubuntu@ubuntu:~/aaa$ ffmpeg -i people.mp4 -codec copy -bsf: hevc_mp4toannexb -f h264 sample.h264
  2. ubuntu@ubuntu:~/aaa$ ffmpeg -i people.mp4 -vcodec h264 people.h264

参考

声明：本文内容由网友自发贡献，不代表【wpsshop博客】立场，版权归原作者所有，本站不承担相应法律责任。如您发现有侵权的内容，请联系我们。转载请注明出处：https://www.wpsshop.cn/w/一键难忘520/article/detail/842692
推荐阅读
相关标签

Copyright © 2003-2013 www.wpsshop.cn 版权所有，并保留所有权利。

  

闽ICP备14008679号