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

yolov7 mask训练笔记_yolo mask

作者：weixin_40725706 | 2024-03-20 18:57:53

赞

踩

yolo mask

目录

训练过程：

不要过早停止训练：

yaml文件解析

关于预训练：

dataloader

coco128-seg.yaml

segments 格式：

解析json格式标签：

统计points个数：

coco合并标签分离为独立文件：

coco标签转换成独立标签线性插值：

导出onnx

onnx结构：

yolov7 mask python训练，tensorrt推理框架

python对比原图和结果：

训练过程：

https://github.com/chelsea456/yolov7_mask/tree/main/yolov7_mask

不要过早停止训练：

parser.add_argument('--patience', type=int, default=0, help='EarlyStopping patience (epochs without improvement)')

yaml文件解析

3、backbone（骨干网络*）
# yolov7 backbone
backbone:
  # [from, number, module, args]
  [[-1, 1, Conv, [32, 3, 1]],  # 0
  
   [-1, 1, Conv, [64, 3, 2]],  # 1-P1/2      
   [-1, 1, Conv, [64, 3, 1]],
   
   [-1, 1, Conv, [128, 3, 2]],  # 3-P2/4  
   [-1, 1, Conv, [64, 1, 1]],
   [-2, 1, Conv, [64, 1, 1]],
   [-1, 1, Conv, [64, 3, 1]],
   [-1, 1, Conv, [64, 3, 1]],
   [-1, 1, Conv, [64, 3, 1]],
   [-1, 1, Conv, [64, 3, 1]],
   [[-1, -3, -5, -6], 1, Concat, [1]],
   [-1, 1, Conv, [256, 1, 1]],  # 11
         
   [-1, 1, MP, []],
   [-1, 1, Conv, [128, 1, 1]],
   [-3, 1, Conv, [128, 1, 1]],
   [-1, 1, Conv, [128, 3, 2]],
   [[-1, -3], 1, Concat, [1]],  # 16-P3/8  
   [-1, 1, Conv, [128, 1, 1]],
   [-2, 1, Conv, [128, 1, 1]],
   [-1, 1, Conv, [128, 3, 1]],
   [-1, 1, Conv, [128, 3, 1]],
   [-1, 1, Conv, [128, 3, 1]],
   [-1, 1, Conv, [128, 3, 1]],
   [[-1, -3, -5, -6], 1, Concat, [1]],
   [-1, 1, Conv, [512, 1, 1]],  # 24
         
   [-1, 1, MP, []],
   [-1, 1, Conv, [256, 1, 1]],
   [-3, 1, Conv, [256, 1, 1]],
   [-1, 1, Conv, [256, 3, 2]],
   [[-1, -3], 1, Concat, [1]],  # 29-P4/16  
   [-1, 1, Conv, [256, 1, 1]],
   [-2, 1, Conv, [256, 1, 1]],
   [-1, 1, Conv, [256, 3, 1]],
   [-1, 1, Conv, [256, 3, 1]],
   [-1, 1, Conv, [256, 3, 1]],
   [-1, 1, Conv, [256, 3, 1]],
   [[-1, -3, -5, -6], 1, Concat, [1]],
   [-1, 1, Conv, [1024, 1, 1]],  # 37
         
   [-1, 1, MP, []],
   [-1, 1, Conv, [512, 1, 1]],
   [-3, 1, Conv, [512, 1, 1]],
   [-1, 1, Conv, [512, 3, 2]],
   [[-1, -3], 1, Concat, [1]],  # 42-P5/32  
   [-1, 1, Conv, [256, 1, 1]],
   [-2, 1, Conv, [256, 1, 1]],
   [-1, 1, Conv, [256, 3, 1]],
   [-1, 1, Conv, [256, 3, 1]],
   [-1, 1, Conv, [256, 3, 1]],
   [-1, 1, Conv, [256, 3, 1]],
   [[-1, -3, -5, -6], 1, Concat, [1]],
   [-1, 1, Conv, [1024, 1, 1]],  # 50
  ]
from 表示该层的输入从哪来。-1表示输入取自上一层，-2表示上两层，3表示第3层（从0开始数），[-1, 4]表示取自上一层和第4层，依次类推。。。。。
number 表示该层模块堆叠的次数，对于C3、BottleneckCSP等模块，表示其子模块的堆叠，具体细节可以查看源代码。当然最终的次数还要乘上depth_multiple系数。
module 表示该层的模块是什么类型。Conv就是卷积+BN+激活模块。所有的模块在 model/common.py 中都有定义。
args 表示输入到模块的参数。例如Conv：[128, 3, 2] 表示输出通道128，卷积核尺寸3，strid=2，当然最终的输出通道数还要乘上 width_multiple，对于其他模块，第一个参数值一般都是指输出通道数，具体细节可以看 model/common.py 中的定义。
 
原文链接：https://blog.csdn.net/weixin_43397302/article/details/126708227

关于预训练：

提高了3个预训练：

yolov5s-seg.pt

yolov7-seg.pt

yolov7x-seg.pt

python segment/train.py --data coco.yaml --batch 16 --weights '' --cfg yolov7-seg.yaml --epochs 300 --name yolov7-seg --img 640 --hyp hyp.scratch-high.yaml

dataloader

labels, shapes, self.segments = zip(*cache.values())

image_detect/coco128-seg.yaml at master · HoyoenKim/image_detect · GitHub

coco128-seg.yaml

download: https://ultralytics.com/assets/coco128-seg.zip

  1. # YOLOv5 by Ultralytics, GPL-3.0 license
  2. # COCO128-seg dataset https://www.kaggle.com/ultralytics/coco128 (first 128 images from COCO train2017) by Ultralytics
  3. # Example usage: python train.py --data coco128.yaml
  4. # parent
  5. # ├── yolov5
  6. # └── datasets
  7. #     └── coco128-seg  ← downloads here (7 MB)
  8.  
  9.  
  10. # Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
  11. path: ../datasets/coco128-seg  # dataset root dir
  12. train: images/train2017  # train images (relative to 'path') 128 images
  13. val: images/train2017  # val images (relative to 'path') 128 images
  14. test:  # test images (optional)
  15.  
  16. # Classes
  17. names:
  18.   0: person
  19.   1: bicycle
  20.   2: car
  21.   3: motorcycle
  22.   4: airplane
  23.   5: bus
  24.   6: train
  25.   7: truck
  26.   8: boat
  27.   9: traffic light
  28.   10: fire hydrant
  29.   11: stop sign
  30.   12: parking meter
  31.   13: bench
  32.   14: bird
  33.   15: cat
  34.   16: dog
  35.   17: horse
  36.   18: sheep
  37.   19: cow
  38.   20: elephant
  39.   21: bear
  40.   22: zebra
  41.   23: giraffe
  42.   24: backpack
  43.   25: umbrella
  44.   26: handbag
  45.   27: tie
  46.   28: suitcase
  47.   29: frisbee
  48.   30: skis
  49.   31: snowboard
  50.   32: sports ball
  51.   33: kite
  52.   34: baseball bat
  53.   35: baseball glove
  54.   36: skateboard
  55.   37: surfboard
  56.   38: tennis racket
  57.   39: bottle
  58.   40: wine glass
  59.   41: cup
  60.   42: fork
  61.   43: knife
  62.   44: spoon
  63.   45: bowl
  64.   46: banana
  65.   47: apple
  66.   48: sandwich
  67.   49: orange
  68.   50: broccoli
  69.   51: carrot
  70.   52: hot dog
  71.   53: pizza
  72.   54: donut
  73.   55: cake
  74.   56: chair
  75.   57: couch
  76.   58: potted plant
  77.   59: bed
  78.   60: dining table
  79.   61: toilet
  80.   62: tv
  81.   63: laptop
  82.   64: mouse
  83.   65: remote
  84.   66: keyboard
  85.   67: cell phone
  86.   68: microwave
  87.   69: oven
  88.   70: toaster
  89.   71: sink
  90.   72: refrigerator
  91.   73: book
  92.   74: clock
  93.   75: vase
  94.   76: scissors
  95.   77: teddy bear
  96.   78: hair drier
  97.   79: toothbrush
  98.  
  99.  
  100. # Download script/URL (optional)
  101. download: https://ultralytics.com/assets/coco128-seg.zip

segments 格式：

segments = [xyn2xy(x, w, h, padw, padh) for x in segments]

segments2boxes

  1.  
  2. def segments2boxes(segments):
  3.     # Convert segment labels to box labels, i.e. (cls, xy1, xy2, ...) to (cls, xywh)
  4.     boxes = []
  5.     for s in segments:
  6.         x, y = s.T  # segment xy
  7.         boxes.append([x.min(), y.min(), x.max(), y.max()])  # cls, xyxy
  8.     return xyxy2xywh(np.array(boxes))  # cls, xywh

解析json格式标签：

  1.     def cache_labels(self, path=Path('./labels.cache'), prefix=''):
  2.         # Cache dataset labels, check images and read shapes
  3.         x = {}  # dict
  4.         nm, nf, ne, nc, msgs = 0, 0, 0, 0, []  # number missing, found, empty, corrupt, messages
  5.         desc = f"{prefix}Scanning '{path.parent / path.stem}' images and labels..."
  6.  
  7.         if self.data_type=="json":
  8.             pbar = tqdm(zip(self.im_files, self.label_files), desc='Scanning images', total=len(self.im_files))
  9.             for i, (im_file, lb_file) in enumerate(pbar):
  10.                 try:
  11.                     annotations = np.zeros((0, 5), dtype=np.float32)
  12.                     shape = np.zeros((2,), dtype=np.int32)
  13.                     label = []
  14.                     segments=[]
  15.                     if lb_file.endswith(".json"):
  16.  
  17.                         json_file = json.load(open(lb_file, "r", encoding="utf-8"))
  18.                         imageHeight = json_file['imageHeight']
  19.                         imageWidth = json_file['imageWidth']
  20.                         shape[0] = imageWidth
  21.                         shape[1] = imageHeight
  22.                         for multi in json_file["shapes"]:
  23.                             points = np.array(multi["points"])
  24.                             xmin = (min(points[:, 0]) if min(points[:, 0]) > 0 else 0) / imageWidth
  25.                             xmax = (max(points[:, 0]) if max(points[:, 0]) > 0 else 0) / imageWidth
  26.                             ymin = (min(points[:, 1]) if min(points[:, 1]) > 0 else 0) / imageHeight
  27.                             ymax = (max(points[:, 1]) if max(points[:, 1]) > 0 else 0) / imageHeight
  28.                             label = multi["label"]
  29.                             if xmax > xmin and ymax > ymin:
  30.                                 annotation = np.zeros((1, 5), dtype=np.float32)
  31.                                 annotation[0, 1] = (xmin + xmax) / 2
  32.                                 annotation[0, 2] = (ymin + ymax) / 2
  33.                                 annotation[0, 3] = xmax - xmin
  34.                                 annotation[0, 4] = ymax - ymin
  35.                                 # cls
  36.                                 annotation[0, 0] = self.hyp['names'].index(label.lower().strip())
  37.                                 # annotation[0, 0] = _class_to_ind[label.lower().strip()]
  38.                                 # annotations = np.append(annotations, annotation, axis=0)
  39.                                 annotations = np.row_stack((annotations, annotation))
  40.                                 points = points.astype(np.float32)
  41.  
  42.                                 points[:, 0] = points[:, 0] / imageWidth
  43.                                 points[:, 1] = points[:, 1] / imageHeight
  44.                                 segments.append(points)
  45.                     if len(annotations) < 1:
  46.                         nm += 1
  47.                         print("json no obj------------")
  48.                         annotations = np.zeros((0, 5), dtype=np.float32)
  49.                     else:
  50.                         nf += 1
  51.                     x[im_file] = [annotations, shape, segments]
  52.                 except Exception as e:
  53.                     nc += 1
  54.                     print(f'{prefix}WARNING: Ignoring corrupted image and/or label {im_file}: {e}')
  55.  
  56.                 pbar.desc = f"{prefix}Scanning '{path.parent / path.stem}' images and labels... " \
  57.                             f"{nf} found, {nm} missing, {ne} empty, {nc} corrupted"
  58.             pbar.close()
  59.  
  60.             if nf == 0:
  61.                 print(f'{prefix}WARNING: No labels found in {path}.')

统计points个数：

  1. import glob
  2. import json
  3.  
  4. import numpy as np
  5.  
  6. if __name__ == '__main__':
  7.  
  8.     label_dir=r'F:\project\detect\yolov7\yolov7_mask\Y\train'
  9.  
  10.     lb_files= glob.glob(label_dir+ '/*.json', recursive=True)  # f = list(p.rglob('*.*'))  # pathlib
  11.  
  12.     annotations = np.zeros((0, 5), dtype=np.float32)
  13.     shape = np.zeros((2,), dtype=np.int32)
  14.     label = []
  15.     for lb_file in lb_files:
  16.         if lb_file.endswith(".json"):
  17.  
  18.             json_file = json.load(open(lb_file, "r", encoding="utf-8"))
  19.             imageHeight = json_file['imageHeight']
  20.             imageWidth = json_file['imageWidth']
  21.             shape[0] = imageWidth
  22.             shape[1] = imageHeight
  23.             for multi in json_file["shapes"]:
  24.                 points = np.array(multi["points"])
  25.                 xmin = (min(points[:, 0]) if min(points[:, 0]) > 0 else 0) / imageWidth
  26.                 xmax = (max(points[:, 0]) if max(points[:, 0]) > 0 else 0) / imageWidth
  27.                 ymin = (min(points[:, 1]) if min(points[:, 1]) > 0 else 0) / imageHeight
  28.                 ymax = (max(points[:, 1]) if max(points[:, 1]) > 0 else 0) / imageHeight
  29.                 label = multi["label"]
  30.                 if xmax > xmin and ymax > ymin:
  31.                     print(lb_file,len(points))
  32.                 # points = points.astype(np.float32)
  33.

coco合并标签分离为独立文件：

  1. import json
  2. import os, cv2
  3.  
  4. import numpy as np
  5.  
  6. def visualization_bbox1(json_path, img_path):  # 需要画的第num副图片， 对应的json路径和图片路径
  7.     with open(json_path ,encoding='utf-8') as annos:
  8.         annotation_json = json.load(annos)
  9.  
  10.     print('num_key is:', len(annotation_json) ,'json key is:', annotation_json.keys())  # 读出json文件的关键字
  11.     print('json num_images is:', len(annotation_json['images']))  # json文件中包含的图片数量
  12.     for img_i in range(len(annotation_json['images'])):
  13.         image_name = annotation_json['images'][img_i]['file_name']  # 读取图片名
  14.         id = annotation_json['images'][img_i]['id']  # 读取图片id
  15.  
  16.         image_path = os.path.join(img_path, str(image_name).zfill(5))  # 拼接图像路径
  17.         image = cv2.imread(image_path, 1)  # 保持原始格式的方式读取图像
  18.         num_bbox = 0  # 统计一幅图片中bbox的数量
  19.         coco_train = dict()
  20.         coco_train['flags'] = {}
  21.         coco_train['imagePath'] = os.path.basename(image_path)
  22.         coco_train['shapes'] = []
  23.         coco_train['imageData'] = None
  24.         coco_train['imageHeight'] = image.shape[0]
  25.         coco_train['imageWidth'] = image.shape[1]
  26.  
  27.         for i in range(len(annotation_json['annotations'][::])):
  28.             if annotation_json['annotations'][i - 1]['image_id'] == id:
  29.                 num_bbox = num_bbox + 1
  30.  
  31.                 box_dict={}
  32.                 box_dict["label"]="quexian"
  33.                 box_dict["shape_type"]="polygon"
  34.                 box_dict["points"]=[]
  35.                 x, y, w, h = annotation_json['annotations'][i - 1]['bbox']  # 读取边框
  36.                 image = cv2.rectangle(image, (int(x), int(y)), (int(x + w), int(y + h)), (0, 255, 255), 2)
  37.                 points = annotation_json['annotations'][i - 1]['segmentation']  # keypoints
  38.  
  39.                 data_len =len(points[0] )//2
  40.                 for index in range(data_len):
  41.                     cv2.circle(image, (int(points[0][index *2]) ,int(points[0][index * 2 +1])), 3, (0, 0, 213), -1)  # x,y,r,color
  42.                     box_dict["points"].append((int(points[0][index *2]) ,int(points[0][index * 2 +1])))
  43.  
  44.                 coco_train['shapes'].append(box_dict)
  45.  
  46.  
  47.         # cv2.imwrite(f"gt/{img_i}.jpg" ,image)
  48.  
  49.         # cv2.resizeWindow("image_name", 2500, 1250)  # 创建500*500的窗口
  50.  
  51.         if is_show:
  52.             print(os.path.basename(image_path))
  53.             cv2.namedWindow("image_name", 0)  # 创建窗口
  54.             cv2.imshow("image_name", image)
  55.             cv2.waitKey(0)
  56.         else:
  57.             train_file=image_path[:-4]+".json"
  58.             with open(train_file, 'w') as write_f:
  59.                 write_f.write(json.dumps(coco_train, indent=2, ensure_ascii=False))
  60.  
  61.  
  62. if __name__ == "__main__":
  63.     is_show=True
  64.     os.makedirs("gt" ,exist_ok=True)
  65.     train_json = r'D:\work\lbg\fenge\data/mark.json'
  66.     train_path = r'D:\work\lbg\fenge\data\trainImage'
  67.     visualization_bbox1(train_json, train_path)
  68.  
  69.

coco标签转换成独立标签线性插值：

  1.  
  2.  
  3. import json
  4. import os, cv2
  5.  
  6. import numpy as np
  7. from scipy.interpolate import interp1d
  8.  
  9. def visualization_bbox1(json_path, img_path):  # 需要画的第num副图片， 对应的json路径和图片路径
  10.     with open(json_path ,encoding='utf-8') as annos:
  11.         annotation_json = json.load(annos)
  12.  
  13.     print('num_key is:', len(annotation_json) ,'json key is:', annotation_json.keys())  # 读出json文件的关键字
  14.     print('json num_images is:', len(annotation_json['images']))  # json文件中包含的图片数量
  15.     for img_i in range(len(annotation_json['images'])):
  16.         image_name = annotation_json['images'][img_i]['file_name']  # 读取图片名
  17.         id = annotation_json['images'][img_i]['id']  # 读取图片id
  18.  
  19.         image_path = os.path.join(img_path, str(image_name).zfill(5))  # 拼接图像路径
  20.         image = cv2.imread(image_path, 1)  # 保持原始格式的方式读取图像
  21.         num_bbox = 0  # 统计一幅图片中bbox的数量
  22.         coco_train = dict()
  23.         coco_train['flags'] = {}
  24.         coco_train['imagePath'] = os.path.basename(image_path)
  25.         coco_train['shapes'] = []
  26.         coco_train['imageData'] = None
  27.         coco_train['imageHeight'] = image.shape[0]
  28.         coco_train['imageWidth'] = image.shape[1]
  29.  
  30.         for i in range(len(annotation_json['annotations'][::])):
  31.             if annotation_json['annotations'][i - 1]['image_id'] == id:
  32.                 num_bbox = num_bbox + 1
  33.  
  34.                 box_dict={}
  35.                 box_dict["label"]="quexian"
  36.                 box_dict["shape_type"]="polygon"
  37.                 box_dict["points"]=[]
  38.                 x, y, w, h = annotation_json['annotations'][i - 1]['bbox']  # 读取边框
  39.                 image = cv2.rectangle(image, (int(x), int(y)), (int(x + w), int(y + h)), (0, 255, 255), 2)
  40.                 points = annotation_json['annotations'][i - 1]['segmentation']  # keypoints
  41.  
  42.                 data_len =len(points[0] )//2
  43.                 print('points len is:', data_len)
  44.                 data_o=[]
  45.                 y_row=[]
  46.                 for index in range(data_len):
  47.                     data_o.append((points[0][index *2],points[0][index * 2 +1]))
  48.                     cv2.circle(image, (int(points[0][index *2]) ,int(points[0][index * 2 +1])), 3, (0, 0, 213), -1)  # x,y,r,color
  49.                     # box_dict["points"].append((int(points[0][index *2]) ,int(points[0][index * 2 +1])))
  50.  
  51.                 coco_train['shapes'].append(box_dict)
  52.                 data_o=np.asarray(data_o)
  53.                 data_o=data_o[np.lexsort(data_o[:, ::-1].T)]
  54.  
  55.                 for index, data_x in enumerate(data_o):
  56.                     if index==0:
  57.                         continue
  58.                     if data_x[0]<=data_o[index-1][0]:
  59.                         data_x[0] = data_o[index-1][0]+0.01
  60.                 # f1 = interp1d(data_o[:,0], data_o[:,1], kind='cubic')
  61.                 f1 = interp1d(data_o[:,0], data_o[:,1], kind='linear')
  62.  
  63.                 x_pred = np.linspace(data_o[0][0], data_o[-1][0], num=155)
  64.                 y1 = f1(x_pred)
  65.  
  66.                 for index,x_data in enumerate(x_pred):
  67.                     cv2.circle(image, (int(x_data), int(y1[index])), 2, (255, 0, 0), -1)  # x,y,r,color
  68.                     box_dict["points"].append((int(x_data), int(y1[index])))
  69.  
  70.         # cv2.imwrite(f"gt/{img_i}.jpg" ,image)
  71.         cv2.namedWindow("image_name", 0)  # 创建窗口
  72.         # cv2.resizeWindow("image_name", 2500, 1250)  # 创建500*500的窗口
  73.         cv2.imshow("image_name", image)
  74.         cv2.waitKey(1)
  75.  
  76.  
  77.         train_file=image_path[:-4]+".json"
  78.         with open(train_file, 'w') as write_f:
  79.             write_f.write(json.dumps(coco_train, indent=2, ensure_ascii=False))
  80.  
  81.  
  82.  
  83. if __name__ == "__main__":
  84.  
  85.     os.makedirs("gt" ,exist_ok=True)
  86.     train_json = r'D:\work\lbg\fenge\data/mark.json'
  87.     train_path = r'D:\work\lbg\fenge\data\trainImage'
  88.     visualization_bbox1(train_json, train_path)
  89.  
  90.

导出onnx

  1. import argparse
  2. import json
  3. import os
  4. import platform
  5. import subprocess
  6. import sys
  7. import time
  8. import warnings
  9. from pathlib import Path
  10.  
  11. import pandas as pd
  12. import torch
  13. import yaml
  14. from torch.utils.mobile_optimizer import optimize_for_mobile
  15.  
  16. FILE = Path(__file__).resolve()
  17. ROOT = FILE.parents[0]  # YOLOv5 root directory
  18. if str(ROOT) not in sys.path:
  19.     sys.path.append(str(ROOT))  # add ROOT to PATH
  20. if platform.system() != 'Windows':
  21.     ROOT = Path(os.path.relpath(ROOT, Path.cwd()))  # relative
  22.  
  23. from models.experimental import attempt_load
  24. from models.yolo import Detect
  25. from utils.dataloaders import LoadImages
  26. from utils.general import (LOGGER, Profile, check_dataset, check_img_size, check_requirements, check_version,
  27.                            check_yaml, colorstr, file_size, get_default_args, print_args, url2file)
  28. from utils.torch_utils import select_device, smart_inference_mode
  29.  
  30.  
  31. def export_formats():
  32.     # YOLOv5 export formats
  33.     x = [['PyTorch', '-', '.pt', True, True], ['TorchScript', 'torchscript', '.torchscript', True, True],
  34.         ['ONNX', 'onnx', '.onnx', True, True], ['OpenVINO', 'openvino', '_openvino_model', True, False],
  35.         ['TensorRT', 'engine', '.engine', False, True], ['CoreML', 'coreml', '.mlmodel', True, False],
  36.         ['TensorFlow SavedModel', 'saved_model', '_saved_model', True, True],
  37.         ['TensorFlow GraphDef', 'pb', '.pb', True, True], ['TensorFlow Lite', 'tflite', '.tflite', True, False],
  38.         ['TensorFlow Edge TPU', 'edgetpu', '_edgetpu.tflite', False, False],
  39.         ['TensorFlow.js', 'tfjs', '_web_model', False, False], ]
  40.     return pd.DataFrame(x, columns=['Format', 'Argument', 'Suffix', 'CPU', 'GPU'])
  41.  
  42.  
  43. def try_export(inner_func):
  44.     # YOLOv5 export decorator, i..e @try_export
  45.     inner_args = get_default_args(inner_func)
  46.  
  47.     def outer_func(*args, **kwargs):
  48.         prefix = inner_args['prefix']
  49.         try:
  50.             with Profile() as dt:
  51.                 f, model = inner_func(*args, **kwargs)
  52.             LOGGER.info(f'{prefix} export success ✅ {dt.t:.1f}s, saved as {f} ({file_size(f):.1f} MB)')
  53.             return f, model
  54.         except Exception as e:
  55.             LOGGER.info(f'{prefix} export failure ❌ {dt.t:.1f}s: {e}')
  56.             return None, None
  57.  
  58.     return outer_func
  59.  
  60.  
  61. @try_export
  62. def export_onnx(model, im, file, opset, train, dynamic, simplify, prefix=colorstr('ONNX:')):
  63.     # YOLOv5 ONNX export
  64.     check_requirements(('onnx',))
  65.     import onnx
  66.  
  67.     LOGGER.info(f'\n{prefix} starting export with onnx {onnx.__version__}...')
  68.     f = file.with_suffix('.onnx')
  69.  
  70.     torch.onnx.export(model.cpu() if dynamic else model,  # --dynamic only compatible with cpu
  71.         im.cpu() if dynamic else im, f, verbose=False, opset_version=opset,
  72.         training=torch.onnx.TrainingMode.TRAINING if train else torch.onnx.TrainingMode.EVAL,
  73.         do_constant_folding=not train, input_names=['images'], output_names=['output'],
  74.         dynamic_axes={'images': {0: 'batch', 2: 'height', 3: 'width'},  # shape(1,3,640,640)
  75.             'output': {0: 'batch', 1: 'anchors'}  # shape(1,25200,85)
  76.         } if dynamic else None)
  77.  
  78.     # Checks
  79.     model_onnx = onnx.load(f)  # load onnx model
  80.     onnx.checker.check_model(model_onnx)  # check onnx model
  81.  
  82.     # Metadata
  83.     d = {'stride': int(max(model.stride)), 'names': model.names}
  84.     for k, v in d.items():
  85.         meta = model_onnx.metadata_props.add()
  86.         meta.key, meta.value = k, str(v)
  87.     onnx.save(model_onnx, f)
  88.  
  89.     # Simplify
  90.     if simplify:
  91.         try:
  92.             cuda = torch.cuda.is_available()
  93.             check_requirements(('onnxruntime-gpu' if cuda else 'onnxruntime', 'onnx-simplifier>=0.4.1'))
  94.             import onnxsim
  95.  
  96.             LOGGER.info(f'{prefix} simplifying with onnx-simplifier {onnxsim.__version__}...')
  97.             model_onnx, check = onnxsim.simplify(model_onnx)
  98.             assert check, 'assert check failed'
  99.             onnx.save(model_onnx, f)
  100.         except Exception as e:
  101.             LOGGER.info(f'{prefix} simplifier failure: {e}')
  102.     return f, model_onnx
  103.  
  104.  
  105. @try_export
  106. def export_openvino(model, file, half, prefix=colorstr('OpenVINO:')):
  107.     # YOLOv5 OpenVINO export
  108.     check_requirements(('openvino-dev',))  # requires openvino-dev: https://pypi.org/project/openvino-dev/
  109.     import openvino.inference_engine as ie
  110.  
  111.     LOGGER.info(f'\n{prefix} starting export with openvino {ie.__version__}...')
  112.     f = str(file).replace('.pt', f'_openvino_model{os.sep}')
  113.  
  114.     cmd = f"mo --input_model {file.with_suffix('.onnx')} --output_dir {f} --data_type {'FP16' if half else 'FP32'}"
  115.     subprocess.check_output(cmd.split())  # export
  116.     with open(Path(f) / file.with_suffix('.yaml').name, 'w') as g:
  117.         yaml.dump({'stride': int(max(model.stride)), 'names': model.names}, g)  # add metadata.yaml
  118.     return f, None
  119.  
  120.  
  121. @try_export
  122. def export_coreml(model, im, file, int8, half, prefix=colorstr('CoreML:')):
  123.     # YOLOv5 CoreML export
  124.     check_requirements(('coremltools',))
  125.     import coremltools as ct
  126.  
  127.     LOGGER.info(f'\n{prefix} starting export with coremltools {ct.__version__}...')
  128.     f = file.with_suffix('.mlmodel')
  129.  
  130.     ts = torch.jit.trace(model, im, strict=False)  # TorchScript model
  131.     ct_model = ct.convert(ts, inputs=[ct.ImageType('image', shape=im.shape, scale=1 / 255, bias=[0, 0, 0])])
  132.     bits, mode = (8, 'kmeans_lut') if int8 else (16, 'linear') if half else (32, None)
  133.     if bits < 32:
  134.         if platform.system() == 'Darwin':  # quantization only supported on macOS
  135.             with warnings.catch_warnings():
  136.                 warnings.filterwarnings("ignore", category=DeprecationWarning)  # suppress numpy==1.20 float warning
  137.                 ct_model = ct.models.neural_network.quantization_utils.quantize_weights(ct_model, bits, mode)
  138.         else:
  139.             print(f'{prefix} quantization only supported on macOS, skipping...')
  140.     ct_model.save(f)
  141.     return f, ct_model
  142.  
  143.  
  144. @try_export
  145. def export_engine(model, im, file, half, dynamic, simplify, workspace=4, verbose=False, prefix=colorstr('TensorRT:')):
  146.     # YOLOv5 TensorRT export https://developer.nvidia.com/tensorrt
  147.     assert im.device.type != 'cpu', 'export running on CPU but must be on GPU, i.e. `python export.py --device 0`'
  148.     try:
  149.         import tensorrt as trt
  150.     except Exception:
  151.         if platform.system() == 'Linux':
  152.             check_requirements(('nvidia-tensorrt',), cmds=('-U --index-url https://pypi.ngc.nvidia.com',))
  153.         import tensorrt as trt
  154.  
  155.     if trt.__version__[0] == '7':  # TensorRT 7 handling https://github.com/ultralytics/yolov5/issues/6012
  156.         grid = model.model[-1].anchor_grid
  157.         model.model[-1].anchor_grid = [a[..., :1, :1, :] for a in grid]
  158.         export_onnx(model, im, file, 12, False, dynamic, simplify)  # opset 12
  159.         model.model[-1].anchor_grid = grid
  160.     else:  # TensorRT >= 8
  161.         check_version(trt.__version__, '8.0.0', hard=True)  # require tensorrt>=8.0.0
  162.         export_onnx(model, im, file, 13, False, dynamic, simplify)  # opset 13
  163.     onnx = file.with_suffix('.onnx')
  164.  
  165.     LOGGER.info(f'\n{prefix} starting export with TensorRT {trt.__version__}...')
  166.     assert onnx.exists(), f'failed to export ONNX file: {onnx}'
  167.     f = file.with_suffix('.engine')  # TensorRT engine file
  168.     logger = trt.Logger(trt.Logger.INFO)
  169.     if verbose:
  170.         logger.min_severity = trt.Logger.Severity.VERBOSE
  171.  
  172.     builder = trt.Builder(logger)
  173.     config = builder.create_builder_config()
  174.     config.max_workspace_size = workspace * 1 << 30
  175.     # config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, workspace << 30)  # fix TRT 8.4 deprecation notice
  176.  
  177.     flag = (1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
  178.     network = builder.create_network(flag)
  179.     parser = trt.OnnxParser(network, logger)
  180.     if not parser.parse_from_file(str(onnx)):
  181.         raise RuntimeError(f'failed to load ONNX file: {onnx}')
  182.  
  183.     inputs = [network.get_input(i) for i in range(network.num_inputs)]
  184.     outputs = [network.get_output(i) for i in range(network.num_outputs)]
  185.     LOGGER.info(f'{prefix} Network Description:')
  186.     for inp in inputs:
  187.         LOGGER.info(f'{prefix}\tinput "{inp.name}" with shape {inp.shape} and dtype {inp.dtype}')
  188.     for out in outputs:
  189.         LOGGER.info(f'{prefix}\toutput "{out.name}" with shape {out.shape} and dtype {out.dtype}')
  190.  
  191.     if dynamic:
  192.         if im.shape[0] <= 1:
  193.             LOGGER.warning(f"{prefix}WARNING: --dynamic model requires maximum --batch-size argument")
  194.         profile = builder.create_optimization_profile()
  195.         for inp in inputs:
  196.             profile.set_shape(inp.name, (1, *im.shape[1:]), (max(1, im.shape[0] // 2), *im.shape[1:]), im.shape)
  197.         config.add_optimization_profile(profile)
  198.  
  199.     LOGGER.info(f'{prefix} building FP{16 if builder.platform_has_fast_fp16 and half else 32} engine in {f}')
  200.     if builder.platform_has_fast_fp16 and half:
  201.         config.set_flag(trt.BuilderFlag.FP16)
  202.     with builder.build_engine(network, config) as engine, open(f, 'wb') as t:
  203.         t.write(engine.serialize())
  204.     return f, None
  205.  
  206.  
  207. @try_export
  208. def export_saved_model(model, im, file, dynamic, tf_nms=False, agnostic_nms=False, topk_per_class=100, topk_all=100,
  209.                        iou_thres=0.45, conf_thres=0.25, keras=False, prefix=colorstr('TensorFlow SavedModel:')):
  210.     # YOLOv5 TensorFlow SavedModel export
  211.     import tensorflow as tf
  212.     from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2
  213.  
  214.     from models.tf import TFModel
  215.  
  216.     LOGGER.info(f'\n{prefix} starting export with tensorflow {tf.__version__}...')
  217.     f = str(file).replace('.pt', '_saved_model')
  218.     batch_size, ch, *imgsz = list(im.shape)  # BCHW
  219.  
  220.     tf_model = TFModel(cfg=model.yaml, model=model, nc=model.nc, imgsz=imgsz)
  221.     im = tf.zeros((batch_size, *imgsz, ch))  # BHWC order for TensorFlow
  222.     _ = tf_model.predict(im, tf_nms, agnostic_nms, topk_per_class, topk_all, iou_thres, conf_thres)
  223.     inputs = tf.keras.Input(shape=(*imgsz, ch), batch_size=None if dynamic else batch_size)
  224.     outputs = tf_model.predict(inputs, tf_nms, agnostic_nms, topk_per_class, topk_all, iou_thres, conf_thres)
  225.     keras_model = tf.keras.Model(inputs=inputs, outputs=outputs)
  226.     keras_model.trainable = False
  227.     keras_model.summary()
  228.     if keras:
  229.         keras_model.save(f, save_format='tf')
  230.     else:
  231.         spec = tf.TensorSpec(keras_model.inputs[0].shape, keras_model.inputs[0].dtype)
  232.         m = tf.function(lambda x: keras_model(x))  # full model
  233.         m = m.get_concrete_function(spec)
  234.         frozen_func = convert_variables_to_constants_v2(m)
  235.         tfm = tf.Module()
  236.         tfm.__call__ = tf.function(lambda x: frozen_func(x)[:4] if tf_nms else frozen_func(x)[0], [spec])
  237.         tfm.__call__(im)
  238.         tf.saved_model.save(tfm, f,
  239.                             options=tf.saved_model.SaveOptions(experimental_custom_gradients=False) if check_version(
  240.                                 tf.__version__, '2.6') else tf.saved_model.SaveOptions())
  241.     return f, keras_model
  242.  
  243.  
  244. @try_export
  245. def export_pb(keras_model, file, prefix=colorstr('TensorFlow GraphDef:')):
  246.     # YOLOv5 TensorFlow GraphDef *.pb export https://github.com/leimao/Frozen_Graph_TensorFlow
  247.     import tensorflow as tf
  248.     from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2
  249.  
  250.     LOGGER.info(f'\n{prefix} starting export with tensorflow {tf.__version__}...')
  251.     f = file.with_suffix('.pb')
  252.  
  253.     m = tf.function(lambda x: keras_model(x))  # full model
  254.     m = m.get_concrete_function(tf.TensorSpec(keras_model.inputs[0].shape, keras_model.inputs[0].dtype))
  255.     frozen_func = convert_variables_to_constants_v2(m)
  256.     frozen_func.graph.as_graph_def()
  257.     tf.io.write_graph(graph_or_graph_def=frozen_func.graph, logdir=str(f.parent), name=f.name, as_text=False)
  258.     return f, None
  259.  
  260.  
  261. @try_export
  262. def export_tflite(keras_model, im, file, int8, data, nms, agnostic_nms, prefix=colorstr('TensorFlow Lite:')):
  263.     # YOLOv5 TensorFlow Lite export
  264.     import tensorflow as tf
  265.  
  266.     LOGGER.info(f'\n{prefix} starting export with tensorflow {tf.__version__}...')
  267.     batch_size, ch, *imgsz = list(im.shape)  # BCHW
  268.     f = str(file).replace('.pt', '-fp16.tflite')
  269.  
  270.     converter = tf.lite.TFLiteConverter.from_keras_model(keras_model)
  271.     converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS]
  272.     converter.target_spec.supported_types = [tf.float16]
  273.     converter.optimizations = [tf.lite.Optimize.DEFAULT]
  274.     if int8:
  275.         from models.tf import representative_dataset_gen
  276.         dataset = LoadImages(check_dataset(check_yaml(data))['train'], img_size=imgsz, auto=False)
  277.         converter.representative_dataset = lambda: representative_dataset_gen(dataset, ncalib=100)
  278.         converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
  279.         converter.target_spec.supported_types = []
  280.         converter.inference_input_type = tf.uint8  # or tf.int8
  281.         converter.inference_output_type = tf.uint8  # or tf.int8
  282.         converter.experimental_new_quantizer = True
  283.         f = str(file).replace('.pt', '-int8.tflite')
  284.     if nms or agnostic_nms:
  285.         converter.target_spec.supported_ops.append(tf.lite.OpsSet.SELECT_TF_OPS)
  286.  
  287.     tflite_model = converter.convert()
  288.     open(f, "wb").write(tflite_model)
  289.     return f, None
  290.  
  291.  
  292. @try_export
  293. def export_edgetpu(file, prefix=colorstr('Edge TPU:')):
  294.     # YOLOv5 Edge TPU export https://coral.ai/docs/edgetpu/models-intro/
  295.     cmd = 'edgetpu_compiler --version'
  296.     help_url = 'https://coral.ai/docs/edgetpu/compiler/'
  297.     assert platform.system() == 'Linux', f'export only supported on Linux. See {help_url}'
  298.     if subprocess.run(f'{cmd} >/dev/null', shell=True).returncode != 0:
  299.         LOGGER.info(f'\n{prefix} export requires Edge TPU compiler. Attempting install from {help_url}')
  300.         sudo = subprocess.run('sudo --version >/dev/null', shell=True).returncode == 0  # sudo installed on system
  301.         for c in ('curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | sudo apt-key add -',
  302.                   'echo "deb https://packages.cloud.google.com/apt coral-edgetpu-stable main" | sudo tee /etc/apt/sources.list.d/coral-edgetpu.list',
  303.                   'sudo apt-get update', 'sudo apt-get install edgetpu-compiler'):
  304.             subprocess.run(c if sudo else c.replace('sudo ', ''), shell=True, check=True)
  305.     ver = subprocess.run(cmd, shell=True, capture_output=True, check=True).stdout.decode().split()[-1]
  306.  
  307.     LOGGER.info(f'\n{prefix} starting export with Edge TPU compiler {ver}...')
  308.     f = str(file).replace('.pt', '-int8_edgetpu.tflite')  # Edge TPU model
  309.     f_tfl = str(file).replace('.pt', '-int8.tflite')  # TFLite model
  310.  
  311.     cmd = f"edgetpu_compiler -s -d -k 10 --out_dir {file.parent} {f_tfl}"
  312.     subprocess.run(cmd.split(), check=True)
  313.     return f, None
  314.  
  315.  
  316. @try_export
  317. def export_tfjs(file, prefix=colorstr('TensorFlow.js:')):
  318.     # YOLOv5 TensorFlow.js export
  319.     check_requirements(('tensorflowjs',))
  320.     import re
  321.  
  322.     import tensorflowjs as tfjs
  323.  
  324.     LOGGER.info(f'\n{prefix} starting export with tensorflowjs {tfjs.__version__}...')
  325.     f = str(file).replace('.pt', '_web_model')  # js dir
  326.     f_pb = file.with_suffix('.pb')  # *.pb path
  327.     f_json = f'{f}/model.json'  # *.json path
  328.  
  329.     cmd = f'tensorflowjs_converter --input_format=tf_frozen_model ' \
  330.           f'--output_node_names=Identity,Identity_1,Identity_2,Identity_3 {f_pb} {f}'
  331.     subprocess.run(cmd.split())
  332.  
  333.     json = Path(f_json).read_text()
  334.     with open(f_json, 'w') as j:  # sort JSON Identity_* in ascending order
  335.         subst = re.sub(r'{"outputs": {"Identity.?.?": {"name": "Identity.?.?"}, '
  336.                        r'"Identity.?.?": {"name": "Identity.?.?"}, '
  337.                        r'"Identity.?.?": {"name": "Identity.?.?"}, '
  338.                        r'"Identity.?.?": {"name": "Identity.?.?"}}}', r'{"outputs": {"Identity": {"name": "Identity"}, '
  339.                                                                       r'"Identity_1": {"name": "Identity_1"}, '
  340.                                                                       r'"Identity_2": {"name": "Identity_2"}, '
  341.                                                                       r'"Identity_3": {"name": "Identity_3"}}}', json)
  342.         j.write(subst)
  343.     return f, None
  344.  
  345.  
  346. @smart_inference_mode()
  347. def run(data=ROOT / 'data/coco128.yaml',  # 'dataset.yaml path'
  348.         weights=ROOT / 'yolov5s.pt',  # weights path
  349.         imgsz=(640, 640),  # image (height, width)
  350.         batch_size=1,  # batch size
  351.         device='cpu',  # cuda device, i.e. 0 or 0,1,2,3 or cpu
  352.         include=('torchscript', 'onnx'),  # include formats
  353.         half=False,  # FP16 half-precision export
  354.         inplace=False,  # set YOLOv5 Detect() inplace=True
  355.         train=False,  # model.train() mode
  356.         keras=False,  # use Keras
  357.         optimize=False,  # TorchScript: optimize for mobile
  358.         int8=False,  # CoreML/TF INT8 quantization
  359.         dynamic=False,  # ONNX/TF/TensorRT: dynamic axes
  360.         simplify=False,  # ONNX: simplify model
  361.         opset=12,  # ONNX: opset version
  362.         verbose=False,  # TensorRT: verbose log
  363.         workspace=4,  # TensorRT: workspace size (GB)
  364.         nms=False,  # TF: add NMS to model
  365.         agnostic_nms=False,  # TF: add agnostic NMS to model
  366.         topk_per_class=100,  # TF.js NMS: topk per class to keep
  367.         topk_all=100,  # TF.js NMS: topk for all classes to keep
  368.         iou_thres=0.45,  # TF.js NMS: IoU threshold
  369.         conf_thres=0.25,  # TF.js NMS: confidence threshold
  370. ):
  371.     t = time.time()
  372.     include = [x.lower() for x in include]  # to lowercase
  373.     fmts = tuple(export_formats()['Argument'][1:])  # --include arguments
  374.     flags = [x in include for x in fmts]
  375.     assert sum(flags) == len(include), f'ERROR: Invalid --include {include}, valid --include arguments are {fmts}'
  376.     jit, onnx, xml, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs = flags  # export booleans
  377.     file = Path(url2file(weights) if str(weights).startswith(('http:/', 'https:/')) else weights)  # PyTorch weights
  378.  
  379.     # Load PyTorch model
  380.     device = select_device(device)
  381.     if half:
  382.         assert device.type != 'cpu' or coreml, '--half only compatible with GPU export, i.e. use --device 0'
  383.         assert not dynamic, '--half not compatible with --dynamic, i.e. use either --half or --dynamic but not both'
  384.     model = attempt_load(weights, device=device, inplace=True, fuse=True)  # load FP32 model
  385.  
  386.     # Checks
  387.     imgsz *= 2 if len(imgsz) == 1 else 1  # expand
  388.     if optimize:
  389.         assert device.type == 'cpu', '--optimize not compatible with cuda devices, i.e. use --device cpu'
  390.  
  391.     # Input
  392.     gs = int(max(model.stride))  # grid size (max stride)
  393.     imgsz = [check_img_size(x, gs) for x in imgsz]  # verify img_size are gs-multiples
  394.     im = torch.zeros(batch_size, 3, *imgsz).to(device)  # image size(1,3,320,192) BCHW iDetection
  395.  
  396.     # Update model
  397.     model.train() if train else model.eval()  # training mode = no Detect() layer grid construction
  398.     for k, m in model.named_modules():
  399.         if isinstance(m, Detect):
  400.             m.inplace = inplace
  401.             m.dynamic = dynamic
  402.             m.export = True
  403.  
  404.     for _ in range(2):
  405.         y = model(im)  # dry runs
  406.     if half and not coreml:
  407.         im, model = im.half(), model.half()  # to FP16
  408.     shape = tuple((y[0] if isinstance(y, tuple) else y).shape)  # model output shape
  409.     LOGGER.info(f"\n{colorstr('PyTorch:')} starting from {file} with output shape {shape} ({file_size(file):.1f} MB)")
  410.  
  411.     # Exports
  412.     f = [''] * 10  # exported filenames
  413.     warnings.filterwarnings(action='ignore', category=torch.jit.TracerWarning)  # suppress TracerWarning
  414.  
  415.     if engine:  # TensorRT required before ONNX
  416.         f[1], _ = export_engine(model, im, file, half, dynamic, simplify, workspace, verbose)
  417.     if onnx or xml:  # OpenVINO requires ONNX
  418.         f[2], _ = export_onnx(model, im, file, opset, train, dynamic, simplify)
  419.     if xml:  # OpenVINO
  420.         f[3], _ = export_openvino(model, file, half)
  421.     if coreml:
  422.         f[4], _ = export_coreml(model, im, file, int8, half)
  423.  
  424.     # TensorFlow Exports
  425.     if any((saved_model, pb, tflite, edgetpu, tfjs)):
  426.         if int8 or edgetpu:  # TFLite --int8 bug https://github.com/ultralytics/yolov5/issues/5707
  427.             check_requirements(('flatbuffers==1.12',))  # required before `import tensorflow`
  428.         assert not tflite or not tfjs, 'TFLite and TF.js models must be exported separately, please pass only one type.'
  429.         f[5], model = export_saved_model(model.cpu(), im, file, dynamic, tf_nms=nms or agnostic_nms or tfjs,
  430.                                          agnostic_nms=agnostic_nms or tfjs, topk_per_class=topk_per_class,
  431.                                          topk_all=topk_all, iou_thres=iou_thres, conf_thres=conf_thres, keras=keras)
  432.         if pb or tfjs:  # pb prerequisite to tfjs
  433.             f[6], _ = export_pb(model, file)
  434.         if tflite or edgetpu:
  435.             f[7], _ = export_tflite(model, im, file, int8 or edgetpu, data=data, nms=nms, agnostic_nms=agnostic_nms)
  436.         if edgetpu:
  437.             f[8], _ = export_edgetpu(file)
  438.         if tfjs:
  439.             f[9], _ = export_tfjs(file)
  440.  
  441.     # Finish
  442.     f = [str(x) for x in f if x]  # filter out '' and None
  443.     if any(f):
  444.         h = '--half' if half else ''  # --half FP16 inference arg
  445.         LOGGER.info(f'\nExport complete ({time.time() - t:.1f}s)'
  446.                     f"\nResults saved to {colorstr('bold', file.parent.resolve())}"
  447.                     f"\nDetect:          python detect.py --weights {f[-1]} {h}"
  448.                     f"\nValidate:        python val.py --weights {f[-1]} {h}"
  449.                     f"\nPyTorch Hub:     model = torch.hub.load('ultralytics/yolov5', 'custom', '{f[-1]}')"
  450.                     f"\nVisualize:       https://netron.app")
  451.     return f  # return list of exported files/dirs
  452.  
  453.  
  454. def parse_opt():
  455.     parser = argparse.ArgumentParser()
  456.     parser.add_argument('--data', type=str, default='../data/data_y.yaml', help='dataset.yaml path')
  457.     parser.add_argument('--weights', nargs='+', type=str, default='./runs/train-seg/exp/weights/best.pt',
  458.                         help='model.pt path(s)')
  459.     parser.add_argument('--imgsz', '--img', '--img-size', nargs='+', type=int, default=[640, 640], help='image (h, w)')
  460.     parser.add_argument('--batch-size', type=int, default=1, help='batch size')
  461.     parser.add_argument('--device', default='cpu', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
  462.     parser.add_argument('--half', action='store_true', help='FP16 half-precision export')
  463.     parser.add_argument('--inplace', action='store_true', help='set YOLOv5 Detect() inplace=True')
  464.     parser.add_argument('--train', action='store_true', help='model.train() mode')
  465.     parser.add_argument('--keras', action='store_true', help='TF: use Keras')
  466.     parser.add_argument('--optimize', action='store_true', help='TorchScript: optimize for mobile')
  467.     parser.add_argument('--int8', action='store_true', help='CoreML/TF INT8 quantization')
  468.     parser.add_argument('--dynamic', action='store_true', help='ONNX/TF/TensorRT: dynamic axes')
  469.     parser.add_argument('--simplify', action='store_true',default=True, help='ONNX: simplify model')
  470.     parser.add_argument('--opset', type=int, default=12, help='ONNX: opset version')
  471.     parser.add_argument('--verbose', action='store_true', help='TensorRT: verbose log')
  472.     parser.add_argument('--workspace', type=int, default=4, help='TensorRT: workspace size (GB)')
  473.     parser.add_argument('--nms', action='store_true', help='TF: add NMS to model')
  474.     parser.add_argument('--agnostic-nms', action='store_true', help='TF: add agnostic NMS to model')
  475.     parser.add_argument('--topk-per-class', type=int, default=100, help='TF.js NMS: topk per class to keep')
  476.     parser.add_argument('--topk-all', type=int, default=100, help='TF.js NMS: topk for all classes to keep')
  477.     parser.add_argument('--iou-thres', type=float, default=0.45, help='TF.js NMS: IoU threshold')
  478.     parser.add_argument('--conf-thres', type=float, default=0.25, help='TF.js NMS: confidence threshold')
  479.     parser.add_argument('--include', nargs='+', default=['onnx'],
  480.                         help='torchscript, onnx, openvino, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs')
  481.     opt = parser.parse_args()
  482.     print_args(vars(opt))
  483.     return opt
  484.  
  485.  
  486. if __name__ == "__main__":
  487.     opt = parse_opt()
  488.     for opt.weights in (opt.weights if isinstance(opt.weights, list) else [opt.weights]):
  489.         run(**vars(opt))

onnx结构：

转ncnn：

onnx2ncnn yolov7_mask.onnx yolov7_mask.param yolov7_mask.bin

报错：

Unsupported slice axes !
ScatterND not supported yet!
Unsupported slice axes !
ScatterND not supported yet!
Unsupported slice axes !
ScatterND not supported yet!

yolov7 mask python训练，tensorrt推理框架

整套框架已跑通，有需要的私信联系 AI视觉网奇的博客_CSDN博客-python宝典,深度学习宝典,pytorch知识宝典领域博主。

python对比原图和结果：

  1.  
  2. import json
  3. import os, cv2
  4.  
  5. import numpy as np
  6.  
  7. def visualization_bbox1(json_path, img_path):  # 需要画的第num副图片， 对应的json路径和图片路径
  8.     with open(json_path ,encoding='utf-8') as annos:
  9.         annotation_json = json.load(annos)
  10.  
  11.     print('num_key is:', len(annotation_json) ,'json key is:', annotation_json.keys())  # 读出json文件的关键字
  12.     print('json num_images is:', len(annotation_json['images']))  # json文件中包含的图片数量
  13.     for img_i in range(len(annotation_json['images'])):
  14.         image_name = annotation_json['images'][img_i]['file_name']  # 读取图片名
  15.         id = annotation_json['images'][img_i]['id']  # 读取图片id
  16.  
  17.         image_path = os.path.join(img_path, str(image_name).zfill(5))  # 拼接图像路径
  18.         image = cv2.imread(image_path, 1)  # 保持原始格式的方式读取图像
  19.         num_bbox = 0  # 统计一幅图片中bbox的数量
  20.         coco_train = dict()
  21.         coco_train['flags'] = {}
  22.         coco_train['imagePath'] = os.path.basename(image_path)
  23.         coco_train['shapes'] = []
  24.         coco_train['imageData'] = None
  25.         coco_train['imageHeight'] = image.shape[0]
  26.         coco_train['imageWidth'] = image.shape[1]
  27.  
  28.         for i in range(len(annotation_json['annotations'][::])):
  29.             if annotation_json['annotations'][i - 1]['image_id'] == id:
  30.                 num_bbox = num_bbox + 1
  31.  
  32.                 box_dict={}
  33.                 box_dict["label"]="quexian"
  34.                 box_dict["shape_type"]="polygon"
  35.                 box_dict["points"]=[]
  36.                 x, y, w, h = annotation_json['annotations'][i - 1]['bbox']  # 读取边框
  37.                 image = cv2.rectangle(image, (int(x), int(y)), (int(x + w), int(y + h)), (0, 255, 255), 2)
  38.                 points = annotation_json['annotations'][i - 1]['segmentation']  # keypoints
  39.  
  40.                 data_len =len(points[0] )//2
  41.                 for index in range(data_len):
  42.                     cv2.circle(image, (int(points[0][index *2]) ,int(points[0][index * 2 +1])), 3, (0, 0, 213), -1)  # x,y,r,color
  43.                     box_dict["points"].append((int(points[0][index *2]) ,int(points[0][index * 2 +1])))
  44.  
  45.                 coco_train['shapes'].append(box_dict)
  46.  
  47.  
  48.         # cv2.imwrite(f"gt/{img_i}.jpg" ,image)
  49.  
  50.         # cv2.resizeWindow("image_name", 2500, 1250)  # 创建500*500的窗口
  51.  
  52.         if is_show:
  53.             print(os.path.basename(image_path))
  54.             # cv2.namedWindow("image_name", 0)  # 创建窗口
  55.  
  56.             if image.shape[1] > 1000:
  57.                 x_scale = 1000 / image.shape[1]
  58.                 image = cv2.resize(image, None, fx=x_scale, fy=x_scale, interpolation=cv2.INTER_AREA)
  59.  
  60.             cv2.imshow("image_name", image)
  61.  
  62.             img_result=cv2.imread(result_dir+os.path.basename(image_path))
  63.  
  64.             if img_result.shape[1] > 1000:
  65.                 x_scale = 1000 / img_result.shape[1]
  66.                 img_result = cv2.resize(img_result, None, fx=x_scale, fy=x_scale, interpolation=cv2.INTER_AREA)
  67.  
  68.             cv2.imshow("img_result", img_result)
  69.             cv2.waitKey(0)
  70.         else:
  71.             train_file=image_path[:-4]+".json"
  72.             with open(train_file, 'w') as write_f:
  73.                 write_f.write(json.dumps(coco_train, indent=2, ensure_ascii=False))
  74.  
  75.  
  76. if __name__ == "__main__":
  77.     is_show=True
  78.  
  79.     result_dir=r'D:\work\lbg\fenge\yolov7_mask-main\segment\runs\predict-seg\exp10/'
  80.     train_json = r'D:\work\lbg\fenge\data/mark.json'
  81.     train_path = r'D:\work\lbg\fenge\data\trainImage'
  82.     os.makedirs(train_path, exist_ok=True)
  83.     visualization_bbox1(train_json, train_path)
  84.  
  85.

 

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

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

  

闽ICP备14008679号