devbox
木道寻08
这个屌丝很懒，什么也没留下！
关注作者
热门标签
热门文章
当前位置:   article > 正文

基于YOLOV8-pose的自动标注_yolov8 自动标注

作者：木道寻08 | 2024-06-26 08:12:23

赞

踩

yolov8 自动标注

解析JSON文件

  1. {
  2.     "version": "5.1.1",  # 字典
  3.     "flags": {},    # 字典
  4.     "shapes": [    # 列表嵌套字典
  5.         {
  6.             "label": "person",    # 字典
  7.             "points": [    # 字典
  8.                 [
  9.                     1184.0,
  10.                     145.0
  11.                 ],
  12.                 [
  13.                     1563.0,
  14.                     743.0
  15.                 ]
  16.             ],
  17.             "group_id": 1,    # 字典
  18.             "shape_type": "rectangle",    # 字典
  19.             "flags": {}    # 字典
  20.         },
  21.         {
  22.             "label": "0",
  23.             "points": [
  24.                 [
  25.                     1519.164794921875,
  26.                     211.12527465820312
  27.                 ]
  28.             ],
  29.             "group_id": 1,
  30.             "shape_type": "point",    #区别 框的话是rectangle 点的话就是point
  31.             "flags": {}
  32.         },
  33.         ......
  34.         {
  35.             "label": "16",
  36.             "points": [
  37.                 [
  38.                     1305.552734375,
  39.                     662.0692138671875
  40.                 ]
  41.             ],
  42.             "group_id": 1,
  43.             "shape_type": "point",
  44.             "flags": {}
  45.         }
  46.     ],
  47.     "imagePath": "22112622_003270.jpg",    # 字典 文件名
  48.     "imageData": "......."       # base64编码 最重要的一项
  49.     "imageHeight": 1080,       # 照片高 字典
  50.     "imageWidth": 1920    # 照片宽 字典
  51. }

现在解析了JSON，之后就要分析怎么在yolov8-pose中取出我想要的值

从YOLOV8-pose取出对应的信息

  1. def postprocess(self, preds, img, orig_imgs):
  2.     """Return detection results for a given input image or list of images."""
  3.  
  4.     preds = ops.non_max_suppression(preds,
  5.                                     self.args.conf,
  6.                                     self.args.iou,
  7.                                     agnostic=self.args.agnostic_nms,
  8.                                     max_det=self.args.max_det,
  9.                                     classes=self.args.classes,
  10.                                     nc=len(self.model.names))
  11.  
  12.     if not isinstance(orig_imgs, list):  # input images are a torch.Tensor, not a list
  13.         orig_imgs = ops.convert_torch2numpy_batch(orig_imgs)
  14.     torch.set_printoptions(sci_mode=False)
  15.  
  16.     results = []
  17.     for i, pred in enumerate(preds):
  18.         orig_img = orig_imgs[i]
  19.         # 对应的全部信息
  20.         pred[:, :4] = ops.scale_boxes(img.shape[2:], pred[:, :4], orig_img.shape).round()
  21.         # 对应的点坐标
  22.         pred_kpts = pred[:, 6:].view(len(pred), *self.model.kpt_shape) if len(pred) else pred[:, 6:]
  23.         pred_kpts = ops.scale_coords(img.shape[2:], pred_kpts, orig_img.shape)
  24.         img_path = self.batch[0][i]
  25.         results.append(
  26.             Results(orig_img, path=img_path, names=self.model.names, boxes=pred[:, :6], keypoints=pred_kpts))
  27.  
  28.     return results

上述代码在 ultralytics/ultralytics/models/yolo/pose/predict.py中

yolov8-pose骨架信息图

可以根据此骨架关键点去做修改，改成自己需要的骨架信息

实现标注自动化

一张图片对应一个JSON文件，后面会有所改变，但是前面都一样，所以可以先把大概的轮廓搭建起来

  1. # 创建一个字典
  2. # 把原始字节码编码成base64字节码
  3. base64_bytes = base64.b64encode(byte_content)
  4.  
  5. # 把base64字节码解码成utf-8格式的字符串
  6. base64_string = base64_bytes.decode('utf-8')
  7. data = {
  8.  
  9.        "version": "5.1.1",
  10.        "flags": {},
  11.        "shapes": [
  12.  
  13.         ]
  14. }
  15. json_data = json.dumps(data, indent=4)
  16. # img_path在源代码有 是图片的绝对路径
  17. #这里是为了imagePath做准备以及保证保存的JSON文件与原文件同名
  18. imgName = img_path.split("/")[-1]
  19. savePath = "/home/test/fall/" + imgName.split(".")[0] + ".json"
  20. #将上述的字典写入我们创建的json文件中
  21. with open(savePath, 'w') as file:
  22.     file.write(json_data)

生成base64_string是借鉴这篇博客https://blog.csdn.net/nodototao/article/details/123800645?spm=1001.2101.3001.6650.2&utm_medium=distribute.pc_relevant.none-task-blog-2~default~CTRLIST~default-2.pc_relevant_default&depth_1-utm_source=distribute.pc_relevant.none-task-blog-2~default~CTRLIST~default-2.pc_relevant_default&utm_relevant_index=5

那么接下来就是将点的信息、框的信息写入到json文件中

  1. for j in range(pred.shape[0]):
  2.     # 将字典转换为JSON格式
  3.     # 打开文件，并将JSON数据写入文件
  4.     with open(savePath, 'r') as f:
  5.         data = json.load(f)
  6.     # 在字典列表中插入新的字典
  7.     # 更新组信息
  8.     bbox = {"label": "person", 'points': [[float(pred[j][0].item()), float(pred[j][1].item())],
  9.                                           [float(pred[j][2].item()), float(pred[j][3].item())]], "group_id": (j + 1),
  10.             "shape_type": "rectangle", "flags": {}}
  11.     data['shapes'].append(bbox)
  12.     data.update({"imagePath": imgName})
  13.     data.update({"imageData": base64_string})
  14.     # data.update({"imageData": str(base64.b64encode(open(img_path, "rb").read()))})
  15.     data.update({"imageHeight": orig_img.shape[0]})
  16.     data.update({"imageWidth": orig_img.shape[1]})
  17.     for i in range(17):
  18.         if float(pred_kpts[j][i][2].item()) > 0.5:
  19.             # print(float(pred_kpts[:, i, 0]), float(pred_kpts[:, i, 1]))
  20.             keypoints = {'label': str(i),
  21.                          'points': [[float(pred_kpts[j][i][0].item()), float(pred_kpts[j][i][1].item())]],
  22.                          'group_id': (j + 1), 'shape_type': 'point', 'flags': {}}
  23.  
  24.  
  25.             data['shapes'].append(keypoints)
  26.     # 将更新后的数据写回到JSON文件
  27.     with open(savePath, 'w') as f:
  28.         json.dump(data, f, indent=4)

完整代码

  1. def postprocess(self, preds, img, orig_imgs):
  2.     """Return detection results for a given input image or list of images."""
  3.  
  4.     preds = ops.non_max_suppression(preds,
  5.                                     self.args.conf,
  6.                                     self.args.iou,
  7.                                     agnostic=self.args.agnostic_nms,
  8.                                     max_det=self.args.max_det,
  9.                                     classes=self.args.classes,
  10.                                     nc=len(self.model.names))
  11.  
  12.     if not isinstance(orig_imgs, list):  # input images are a torch.Tensor, not a list
  13.         orig_imgs = ops.convert_torch2numpy_batch(orig_imgs)
  14.     torch.set_printoptions(sci_mode=False)
  15.  
  16.     results = []
  17.     for i, pred in enumerate(preds):
  18.         orig_img = orig_imgs[i]
  19.         pred[:, :4] = ops.scale_boxes(img.shape[2:], pred[:, :4], orig_img.shape).round()
  20.  
  21.         pred_kpts = pred[:, 6:].view(len(pred), *self.model.kpt_shape) if len(pred) else pred[:, 6:]
  22.         pred_kpts = ops.scale_coords(img.shape[2:], pred_kpts, orig_img.shape)
  23.         img_path = self.batch[0][i]
  24.         with open(img_path, 'rb') as jpg_file:
  25.             byte_content = jpg_file.read()
  26.         # 把原始字节码编码成base64字节码
  27.         base64_bytes = base64.b64encode(byte_content)
  28.  
  29.         # 把base64字节码解码成utf-8格式的字符串
  30.         base64_string = base64_bytes.decode('utf-8')
  31.         import json
  32.  
  33.         # 创建一个字典
  34.         data = {
  35.  
  36.             "version": "5.1.1",
  37.             "flags": {},
  38.             "shapes": [
  39.  
  40.             ]
  41.  
  42.         }
  43.         json_data = json.dumps(data, indent=4)
  44.         imgName = img_path.split("/")[-1]
  45.         print(img_path)
  46.         savePath = "/home/ebo/test/high/_merge/" + imgName.split(".")[0] + ".json"
  47.         with open(savePath, 'w') as file:
  48.             file.write(json_data)
  49.         for j in range(pred.shape[0]):
  50.             # 将字典转换为JSON格式
  51.             # 打开文件，并将JSON数据写入文件
  52.             with open(savePath, 'r') as f:
  53.                 data = json.load(f)
  54.             # 在字典列表中插入新的字典
  55.             bbox = {"label": "person", 'points': [[float(pred[j][0].item()), float(pred[j][1].item())],
  56.                                                   [float(pred[j][2].item()), float(pred[j][3].item())]],
  57.                     "group_id": (j + 1),
  58.                     "shape_type": "rectangle", "flags": {}}
  59.             data['shapes'].append(bbox)
  60.             data.update({"imagePath": imgName})
  61.             data.update({"imageData": base64_string})
  62.             # data.update({"imageData": str(base64.b64encode(open(img_path, "rb").read()))})
  63.             data.update({"imageHeight": orig_img.shape[0]})
  64.             data.update({"imageWidth": orig_img.shape[1]})
  65.             for i in range(17):
  66.                 if float(pred_kpts[j][i][2].item()) > 0.5:
  67.                     # print(float(pred_kpts[:, i, 0]), float(pred_kpts[:, i, 1]))
  68.                     keypoints = {'label': str(i),
  69.                                  'points': [[float(pred_kpts[j][i][0].item()), float(pred_kpts[j][i][1].item())]],
  70.                                  'group_id': (j + 1), 'shape_type': 'point', 'flags': {}}
  71.                     # keypoints = {'label': i, 'points': 0, 'group_id': 1, 'shape_type': 'point', 'flags': {}}
  72.                     # print(keypoints)da
  73.  
  74.                     data['shapes'].append(keypoints)
  75.             # 将更新后的数据写回到JSON文件
  76.             with open(savePath, 'w') as f:
  77.                 json.dump(data, f, indent=4)
  78.  
  79.         # print(pred)
  80.         results.append(
  81.             Results(orig_img, path=img_path, names=self.model.names, boxes=pred[:, :6], keypoints=pred_kpts))
  82.  
  83.     return results

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

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

  

闽ICP备14008679号