From 2b98bd60f2ae80d496adcbb29a961cfd80a07b4c Mon Sep 17 00:00:00 2001 From: Shihan Liu <69880398+LSH9832@users.noreply.github.com> Date: Tue, 13 Dec 2022 22:36:46 +0800 Subject: [PATCH] Add files via upload --- detect.py | 120 ++++++++++++++++++++++ onnx2trt.py | 74 ++++++++++++++ params/model/edgeyolo-visdrone.yaml | 140 ++++++++++++++++++++++++++ params/model/edgeyolo_coco.yaml | 141 ++++++++++++++++++++++++++ params/model/edgeyolo_coco_s.yaml | 148 ++++++++++++++++++++++++++++ params/train/train_settings.yaml | 71 +++++++++++++ pth2onnx.py | 142 ++++++++++++++++++++++++++ requirements.txt | 18 ++++ 8 files changed, 854 insertions(+) create mode 100644 detect.py create mode 100644 onnx2trt.py create mode 100644 params/model/edgeyolo-visdrone.yaml create mode 100644 params/model/edgeyolo_coco.yaml create mode 100644 params/model/edgeyolo_coco_s.yaml create mode 100644 params/train/train_settings.yaml create mode 100644 pth2onnx.py create mode 100644 requirements.txt diff --git a/detect.py b/detect.py new file mode 100644 index 0000000..26942f6 --- /dev/null +++ b/detect.py @@ -0,0 +1,120 @@ +from datetime import datetime as date +from glob import glob +import os +from loguru import logger +import argparse +import cv2 + +from edgeyolo.detect import Detector, TRTDetector, draw + + +if __name__ == '__main__': + + parser = argparse.ArgumentParser("EdgeYOLO Detect parser") + parser.add_argument("-w", "--weights", type=str, default="edgeyolo_coco.pth", help="weight file") + parser.add_argument("-c", "--conf-thres", type=float, default=0.25, help="confidence threshold") + parser.add_argument("-n", "--nms-thres", type=float, default=0.55, help="nms threshold") + parser.add_argument("--fp16", action="store_true", help="fp16") + parser.add_argument("--no-fuse", action="store_true", help="do not fuse model") + parser.add_argument("--input-size", type=int, nargs="+", default=[640, 640], help="input size: [height, width]") + parser.add_argument("-s", "--source", type=str, default="E:/videos/test.avi", help="video source or image dir") + parser.add_argument("--trt", action="store_true", help="is trt model") + parser.add_argument("--legacy", action="store_true", help="if img /= 255 while training, add this command.") + parser.add_argument("--use-decoder", action="store_true", help="support original yolox model v0.2.0") + parser.add_argument("--batch-size", type=int, default=1, help="batch size") + parser.add_argument("--no-label", action="store_true", help="do not draw label") + parser.add_argument("--save-dir", type=str, default="./imgs/coco", help="image result save dir") + + args = parser.parse_args() + exist_save_dir = os.path.isdir(args.save_dir) + + # detector setup + detector = TRTDetector if args.trt else Detector + detect = detector( + weight_file=args.weights, + conf_thres=args.conf_thres, + nms_thres=args.nms_thres, + input_size=args.input_size, + fuse=not args.no_fuse, + fp16=args.fp16, + use_decoder=args.use_decoder + ) + if args.trt: + args.batch_size = detect.batch_size + + # source loader setup + if os.path.isdir(args.source): + + class DirCapture: + + def __init__(self, dir_name): + self.imgs = [] + for img_type in ["jpg", "png", "jpeg", "bmp", "webp"]: + self.imgs += sorted(glob(os.path.join(dir_name, f"*.{img_type}"))) + + def isOpened(self): + return bool(len(self.imgs)) + + def read(self): + print(self.imgs[0]) + now_img = cv2.imread(self.imgs[0]) + self.imgs = self.imgs[1:] + return now_img is not None, now_img + + source = DirCapture(args.source) + delay = 0 + else: + source = cv2.VideoCapture(int(args.source) if args.source.isdigit() else args.source) + delay = 1 + + all_dt = [] + dts_len = 300 // args.batch_size + success = True + + # start inference + while source.isOpened() and success: + + frames = [] + for _ in range(args.batch_size): + success, frame = source.read() + if not success: + if not len(frames): + cv2.destroyAllWindows() + break + else: + while len(frames) < args.batch_size: + frames.append(frames[-1]) + else: + frames.append(frame) + + if not len(frames): + break + + results = detect(frames, args.legacy) + dt = detect.dt + all_dt.append(dt) + if len(all_dt) > dts_len: + all_dt = all_dt[-dts_len:] + print(f"\r{dt * 1000 / args.batch_size:.1f}ms " + f"average:{sum(all_dt) / len(all_dt) / args.batch_size * 1000:.1f}ms", end=" ") + + key = -1 + imgs = draw(frames, results, detect.class_names, 2, draw_label=not args.no_label) + # print([im.shape for im in frames]) + for img in imgs: + # print(img.shape) + cv2.imshow("EdgeYOLO result", img) + key = cv2.waitKey(delay) + if key in [ord("q"), 27]: + break + elif key == ord(" "): + delay = 1 - delay + elif key == ord("s"): + if not exist_save_dir: + os.makedirs(args.save_dir, exist_ok=True) + file_name = f"{str(date.now()).split('.')[0].replace(':', '').replace('-', '').replace(' ', '')}.jpg" + cv2.imwrite(os.path.join(args.save_dir, file_name), img) + logger.info(f"image saved to {file_name}.") + if key in [ord("q"), 27]: + cv2.destroyAllWindows() + break diff --git a/onnx2trt.py b/onnx2trt.py new file mode 100644 index 0000000..49fbcaf --- /dev/null +++ b/onnx2trt.py @@ -0,0 +1,74 @@ +import yaml +import argparse +import os.path as osp +import os +from loguru import logger +import torch + + +def get_args(): + parser = argparse.ArgumentParser("EdgeYOLO onnx2tensorrt parser") + parser.add_argument("-o", "--onnx", type=str, default="yolov7.onnx", help="ONNX file") + parser.add_argument("-y", "--yaml", type=str, default="yolov7.yaml", help="export params file") + parser.add_argument("-w", "--workspace", type=int, default=8, help="export memory workspace(GB)") + parser.add_argument("--fp16", action="store_true", help="fp16") + parser.add_argument("--int8", action="store_true", help="int8") + parser.add_argument("--best", action="store_true", help="best") + parser.add_argument("-d", "--dist-path", type=str, default="export_output/tensorrt") + parser.add_argument("--batch-size", type=int, default=0, help="batch-size") + return parser.parse_args() + + +def main(): + args = get_args() + + assert osp.isfile(args.onnx), f"No such file named {args.onnx}." + assert osp.isfile(args.yaml), f"No such file named {args.yaml}." + + os.makedirs(args.dist_path, exist_ok=True) + + name = args.onnx.replace("\\", "/").split("/")[-1][:-len(args.onnx.split(".")[-1])] + + engine_file = osp.join(args.dist_path, name + "engine").replace("\\", "/") + pt_file = osp.join(args.dist_path, name + "pt").replace("\\", "/") + cls_file = osp.join(args.dist_path, name + "txt").replace("\\", "/") + params = yaml.load(open(args.yaml).read(), yaml.Loader) + command = f"trtexec --onnx={args.onnx}" \ + f"{' --fp16' if args.fp16 else ' --int8' if args.int8 else ' --best' if args.best else ''} " \ + f"--saveEngine={engine_file} --workspace={args.workspace*1024} " \ + f"--batch={args.batch_size if not args.batch_size > 0 else params['batch_size'] if 'batch_size' in params else 1}" + + logger.info("start converting onnx to tensorRT engine file.") + os.system(command) + + if not osp.isfile(engine_file): + logger.error("tensorRT engine file convertion failed.") + return + + logger.info(f"tensorRT engine saved to {engine_file}") + + try: + data = { + "model": { + "engine": bytearray(open(engine_file, "rb").read()), + "input_names": params["input_name"], + "output_names": params["output_name"] + }, + "names": params["names"], + "img_size": params["img_size"], + "batch_size": params["batch_size"] + } + class_str = "" + for name in params["names"]: + class_str += name + "\n" + with open(cls_file, "w") as cls_f: + cls_f.write(class_str[:-1]) + logger.info(f"class names txt pt saved to {cls_file}") + torch.save(data, pt_file) + logger.info(f"tensorRT pt saved to {pt_file}") + except Exception as e: + logger.error(f"convert2pt error: {e}") + + +if __name__ == '__main__': + main() diff --git a/params/model/edgeyolo-visdrone.yaml b/params/model/edgeyolo-visdrone.yaml new file mode 100644 index 0000000..60b83f8 --- /dev/null +++ b/params/model/edgeyolo-visdrone.yaml @@ -0,0 +1,140 @@ +# parameters +nc: 10 # number of classes +depth_multiple: 1.0 # models depth multiple +width_multiple: 1.0 # layer channel multiple + +# anchor-box-free +anchors: + - [8, 8] # P3/8 + - [16, 16] # P4/16 + - [32, 32] # P5/32 + +# edgeyolo 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 + ] + +# edgeyolo head +head: + [[-1, 1, SPPCSPC, [512]], # 51 + + [-1, 1, Conv, [256, 1, 1]], + [-1, 1, nn.Upsample, [None, 2, 'nearest']], + [37, 1, Conv, [256, 1, 1]], # route backbone P4 + [[-1, -2], 1, Concat, [1]], + + [-1, 1, Conv, [256, 1, 1]], + [-2, 1, Conv, [256, 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, -2, -3, -4, -5, -6], 1, Concat, [1]], + [-1, 1, Conv, [256, 1, 1]], # 63 + + [-1, 1, Conv, [128, 1, 1]], + [-1, 1, nn.Upsample, [None, 2, 'nearest']], + [24, 1, Conv, [128, 1, 1]], # route backbone P3 + [[-1, -2], 1, Concat, [1]], + + [-1, 1, Conv, [128, 1, 1]], + [-2, 1, Conv, [128, 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, -2, -3, -4, -5, -6], 1, Concat, [1]], + [-1, 1, Conv, [128, 1, 1]], # 75 + + [-1, 1, MP, []], + [-1, 1, Conv, [128, 1, 1]], + [-3, 1, Conv, [128, 1, 1]], + [-1, 1, Conv, [128, 3, 2]], + [[-1, -3, 63], 1, Concat, [1]], + + [-1, 1, Conv, [256, 1, 1]], + [-2, 1, Conv, [256, 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, -2, -3, -4, -5, -6], 1, Concat, [1]], + [-1, 1, Conv, [256, 1, 1]], # 88 + + [-1, 1, MP, []], + [-1, 1, Conv, [256, 1, 1]], + [-3, 1, Conv, [256, 1, 1]], + [-1, 1, Conv, [256, 3, 2]], + [[-1, -3, 51], 1, Concat, [1]], + + [-1, 1, Conv, [512, 1, 1]], + [-2, 1, Conv, [512, 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, -2, -3, -4, -5, -6], 1, Concat, [1]], + [-1, 1, Conv, [512, 1, 1]], # 101 + + [75, 1, RepConv, [256, 3, 1]], # 102 + [88, 1, RepConv, [512, 3, 1]], # 103 + [101, 1, RepConv, [1024, 3, 1]], # 104 + + [[102,103,104], 1, YOLOXDetect, [nc, anchors, Conv]], + ] diff --git a/params/model/edgeyolo_coco.yaml b/params/model/edgeyolo_coco.yaml new file mode 100644 index 0000000..c78791d --- /dev/null +++ b/params/model/edgeyolo_coco.yaml @@ -0,0 +1,141 @@ +# parameters +nc: 80 # number of classes +depth_multiple: 1.0 # models depth multiple +width_multiple: 1.0 # layer channel multiple + +# anchor-box-free +anchors: + - [8, 8] # P3/8 + - [16, 16] # P4/16 + - [32, 32] # P5/32 + +# edgeyolo 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 + ] + +# edgeyolo head +head: + [[-1, 1, SPPCSPC, [512]], # 51 + + [-1, 1, Conv, [256, 1, 1]], + [-1, 1, nn.Upsample, [None, 2, 'nearest']], + [37, 1, Conv, [256, 1, 1]], # route backbone P4 + [[-1, -2], 1, Concat, [1]], + + [-1, 1, Conv, [256, 1, 1]], + [-2, 1, Conv, [256, 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, -2, -3, -4, -5, -6], 1, Concat, [1]], + [-1, 1, Conv, [256, 1, 1]], # 63 + + [-1, 1, Conv, [128, 1, 1]], + [-1, 1, nn.Upsample, [None, 2, 'nearest']], + [24, 1, Conv, [128, 1, 1]], # route backbone P3 + [[-1, -2], 1, Concat, [1]], + + [-1, 1, Conv, [128, 1, 1]], + [-2, 1, Conv, [128, 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, -2, -3, -4, -5, -6], 1, Concat, [1]], + [-1, 1, Conv, [128, 1, 1]], # 75 + + [-1, 1, MP, []], + [-1, 1, Conv, [128, 1, 1]], + [-3, 1, Conv, [128, 1, 1]], + [-1, 1, Conv, [128, 3, 2]], + [[-1, -3, 63], 1, Concat, [1]], + + [-1, 1, Conv, [256, 1, 1]], + [-2, 1, Conv, [256, 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, -2, -3, -4, -5, -6], 1, Concat, [1]], + [-1, 1, Conv, [256, 1, 1]], # 88 + + [-1, 1, MP, []], + [-1, 1, Conv, [256, 1, 1]], + [-3, 1, Conv, [256, 1, 1]], + [-1, 1, Conv, [256, 3, 2]], + [[-1, -3, 51], 1, Concat, [1]], + + [-1, 1, Conv, [512, 1, 1]], + [-2, 1, Conv, [512, 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, -2, -3, -4, -5, -6], 1, Concat, [1]], + [-1, 1, Conv, [512, 1, 1]], # 101 + + [75, 1, RepConv, [256, 3, 1]], # 102 + [88, 1, RepConv, [512, 3, 1]], # 103 + [101, 1, RepConv, [1024, 3, 1]], # 104 + + [[102,103,104], 1, YOLOXDetect, [nc, anchors, Conv]], + ] + ] diff --git a/params/model/edgeyolo_coco_s.yaml b/params/model/edgeyolo_coco_s.yaml new file mode 100644 index 0000000..5d06ff2 --- /dev/null +++ b/params/model/edgeyolo_coco_s.yaml @@ -0,0 +1,148 @@ +# parameters +nc: 80 # number of classes +depth_multiple: 1.0 # models depth multiple +width_multiple: 1.0 # layer channel multiple + +# anchor-box-free +anchors: + - [8, 8] # P3/8 + - [16, 16] # P4/16 + - [32, 32] # P5/32 + +# edgeyolo 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, [80, 3, 2]], # 3-P2/4 + [-1, 1, Conv, [40, 1, 1]], + [-2, 1, Conv, [40, 1, 1]], + [-1, 1, Conv, [40, 3, 1]], + [-1, 1, Conv, [40, 3, 1]], + [-1, 1, Conv, [40, 3, 1]], + [-1, 1, Conv, [40, 3, 1]], + [[-1, -2], 1, Shortcut, [1]], + [[-4, -5], 1, Shortcut, [1]], + [[-1, -2, -7, -8], 1, Concat, [1]], + [-1, 1, Conv, [120, 1, 1]], # 13 + + [-1, 1, MP, []], + [-1, 1, Conv, [80, 1, 1]], + [-3, 1, Conv, [80, 1, 1]], + [-1, 1, Conv, [80, 3, 2]], + [[-1, -3], 1, Concat, [1]], # 18-P3/8 + [-1, 1, Conv, [80, 1, 1]], + [-2, 1, Conv, [80, 1, 1]], + [-1, 1, DWConv, [80, 3, 1]], + [-1, 1, Conv, [80, 3, 1]], + [-1, 1, DWConv, [80, 3, 1]], + [-1, 1, Conv, [80, 3, 1]], + [[-1, -2], 1, Shortcut, [1]], + [[-4, -5], 1, Shortcut, [1]], + [[-1, -2, -7, -8], 1, Concat, [1]], + [-1, 1, Conv, [240, 1, 1]], # 28 + + [-1, 1, MP, []], + [-1, 1, Conv, [160, 1, 1]], + [-3, 1, Conv, [160, 1, 1]], + [-1, 1, Conv, [160, 3, 2]], + [[-1, -3], 1, Concat, [1]], # 33-P4/16 + [-1, 1, Conv, [160, 1, 1]], + [-2, 1, Conv, [160, 1, 1]], + [-1, 1, DWConv, [160, 3, 1]], + [-1, 1, Conv, [160, 3, 1]], + [-1, 1, DWConv, [160, 3, 1]], + [-1, 1, Conv, [160, 3, 1]], + [[-1, -2], 1, Shortcut, [1]], + [[-4, -5], 1, Shortcut, [1]], + [[-1, -2, -7, -8], 1, Concat, [1]], + [-1, 1, Conv, [480, 1, 1]], # 43 + + [-1, 1, MP, []], + [-1, 1, Conv, [320, 1, 1]], + [-3, 1, Conv, [320, 1, 1]], + [-1, 1, Conv, [320, 3, 2]], + [[-1, -3], 1, Concat, [1]], # 48-P5/32 + [-1, 1, Conv, [160, 1, 1]], + [-2, 1, Conv, [160, 1, 1]], + [-1, 1, DWConv, [160, 3, 1]], + [-1, 1, Conv, [160, 3, 1]], + [-1, 1, DWConv, [160, 3, 1]], + [-1, 1, Conv, [160, 3, 1]], + [[-1, -2], 1, Shortcut, [1]], + [[-4, -5], 1, Shortcut, [1]], + [[-1, -2, -7, -8], 1, Concat, [1]], + [-1, 1, Conv, [480, 1, 1]], # 58 + ] + +# edgeyolo head +head: + [[-1, 1, SPPCSPC, [240]], # 59 + + [-1, 1, Conv, [120, 1, 1]], + [-1, 1, nn.Upsample, [None, 2, 'nearest']], + [43, 1, Conv, [120, 1, 1]], # route backbone P4 + [[-1, -2], 1, Concat, [1]], + + [-1, 1, Conv, [120, 1, 1]], + [-2, 1, Conv, [120, 1, 1]], + [-1, 1, DWConv, [60, 3, 1]], + [-1, 1, Conv, [60, 3, 1]], + [-1, 1, DWConv, [60, 3, 1]], + [-1, 1, Conv, [60, 3, 1]], + [[-1, -2, -3, -4, -5, -6], 1, Concat, [1]], + [-1, 1, Conv, [120, 1, 1]], # 71 + + [-1, 1, Conv, [60, 1, 1]], + [-1, 1, nn.Upsample, [None, 2, 'nearest']], + [28, 1, Conv, [60, 1, 1]], # route backbone P3 + [[-1, -2], 1, Concat, [1]], + + [-1, 1, Conv, [60, 1, 1]], + [-2, 1, Conv, [60, 1, 1]], + [-1, 1, DWConv, [30, 3, 1]], + [-1, 1, Conv, [30, 3, 1]], + [-1, 1, DWConv, [30, 3, 1]], + [-1, 1, Conv, [30, 3, 1]], + [[-1, -2, -3, -4, -5, -6], 1, Concat, [1]], + [-1, 1, Conv, [60, 1, 1]], # 83 + + [-1, 1, MP, []], + [-1, 1, Conv, [60, 1, 1]], + [-3, 1, Conv, [60, 1, 1]], + [-1, 1, Conv, [60, 3, 2]], + [[-1, -3, 71], 1, Concat, [1]], + + [-1, 1, Conv, [120, 1, 1]], + [-2, 1, Conv, [120, 1, 1]], + [-1, 1, DWConv, [60, 3, 1]], + [-1, 1, Conv, [60, 3, 1]], + [-1, 1, DWConv, [60, 3, 1]], + [-1, 1, Conv, [60, 3, 1]], + [[-1, -2, -3, -4, -5, -6], 1, Concat, [1]], + [-1, 1, Conv, [120, 1, 1]], # 96 + + [-1, 1, MP, []], + [-1, 1, Conv, [120, 1, 1]], + [-3, 1, Conv, [120, 1, 1]], + [-1, 1, Conv, [120, 3, 2]], + [[-1, -3, 59], 1, Concat, [1]], + + [-1, 1, Conv, [240, 1, 1]], + [-2, 1, Conv, [240, 1, 1]], + [-1, 1, DWConv, [120, 3, 1]], + [-1, 1, Conv, [120, 3, 1]], + [-1, 1, DWConv, [120, 3, 1]], + [-1, 1, Conv, [120, 3, 1]], + [[-1, -2, -3, -4, -5, -6], 1, Concat, [1]], + [-1, 1, Conv, [240, 1, 1]], # 109 + + [83, 1, RepConv, [120, 3, 1]], # 110 + [96, 1, RepConv, [240, 3, 1]], # 111 + [109, 1, RepConv, [480, 3, 1]], # 112 + + [[110,111,112], 1, YOLOXDetect, [nc, anchors, Conv]], + ] diff --git a/params/train/train_settings.yaml b/params/train/train_settings.yaml new file mode 100644 index 0000000..67bd1cb --- /dev/null +++ b/params/train/train_settings.yaml @@ -0,0 +1,71 @@ +# models & weights +model_cfg: "params/model/edgeyolo.yaml" +weights: "./train_output/edgeyolo/last.pth" # it contains models config, set null if not use it +use_cfg: false # force using cfg file data instead of cfg in weights + +# output +output_dir: "./train_output/edgeyolo" +save_checkpoint_for_each_epoch: true +log_file: "log.txt" + +# dataset +dataset_dir: "/dataset/coco2017" +train_dir: "train2017" +val_dir: "val2017" +train_anno_file: "annotations/instances_train2017.json" +val_anno_file: "annotations/instances_val2017.json" +class_names_file: "classes.txt" # this file is in dataset dir +use_segmentation: false # use segmentation labels to generate bbox instead of using bbox labels directly. + +# dataloader +batch_size_per_gpu: 8 +loader_num_workers: 4 +num_threads: 1 + +# device & data type +device: [0, 1, 2, 3] +fp16: false +cudnn_benchmark: false # set true if multiscale_range is set zero + +## train hyper-params +optimizer: "Adam" +max_epoch: 400 # 300 +close_mosaic_epochs: 15 + +# learning rate +lr_per_img: 0.00015625 +warmup_epochs: 5 +warmup_lr_ratio: 0.0 +final_lr_ratio: 0.01 # 0.05 (if epoch is 300) + +# training/dataset enhancement +# [cls_loss, conf_loss, iou_loss] +loss_use: ["bce", "bce", "giou"] # bce: BCE loss. bcf: Balanced Focal loss. hyb: HR loss, iou, c/g/siou is available + +input_size: [640, 640] +multiscale_range: 5 +weight_decay: 0.0005 +momentum: 0.9 +use_ema: true +enable_mixup: true +mixup_scale: [0.5, 1.5] +mosaic_scale: [0.1, 2.0] +flip_prob: 0.5 +mosaic_prob: 1 +mixup_prob: 1 +degrees: 10 + +# evaluate +eval_at_start: false +val_conf_thres: 0.01 +val_nms_thres: 0.65 +eval_only: false # running benchmark only for all weights in output_dir + +# show +print_interval: 10 + +# others +load_optimizer_params: true # load optimizer params when resume train, set false if there is an error. +train_backbone: true # set false if you only want to train head +train_start_layers: 51 # when not train backbone, train layers starts from this layer, see params/models/edgeyolo.yaml +force_start_epoch: -1 # set -1 to disable this option diff --git a/pth2onnx.py b/pth2onnx.py new file mode 100644 index 0000000..73a8831 --- /dev/null +++ b/pth2onnx.py @@ -0,0 +1,142 @@ +import argparse +import os +import sys +import time +import yaml + +sys.path.append('./') # to run '$ python *.py' files in subdirectories + +import torch +import torch.nn as nn + +from edgeyolo import EdgeYOLO +from edgeyolo.utils2.activations import SiLU +from loguru import logger +# from edgeyolo.utils2.general import set_logging + + +def replace_module(module, replaced_module_type, new_module_type, replace_func=None): + """ + Replace given type in module to a new type. mostly used in deploy. + """ + + def default_replace_func(replaced_module_type, new_module_type): + return new_module_type() + + if replace_func is None: + replace_func = default_replace_func + + model = module + if isinstance(module, replaced_module_type): + model = replace_func(replaced_module_type, new_module_type) + else: # recurrsively replace + for name, child in module.named_children(): + new_child = replace_module(child, replaced_module_type, new_module_type) + if new_child is not child: # child is already replaced + model.add_module(name, new_child) + + return model + + +if __name__ == '__main__': + + parser = argparse.ArgumentParser("EdgeYOLO onnx-export parser") + parser.add_argument('--weights', type=str, default='./edgeyolo.pth', help='weights path') + parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='image size') # height, width + parser.add_argument('--batch-size', type=int, default=1, help='batch size') + parser.add_argument('--simplify', action='store_true', help='simplify onnx model') + parser.add_argument('--opset', type=int, default=11, help='opset version') + + opt = parser.parse_args() + + os.makedirs("export_output/onnx", exist_ok=True) + print(opt) + # set_logging() + t = time.time() + + opt.grid = True + opt.img_size *= 2 if len(opt.img_size) == 1 else 1 + + device = torch.device('cpu') + exp = EdgeYOLO(None, opt.weights) + model = exp.model + model.eval() + with torch.no_grad(): + model.fuse() + labels = exp.class_names + + replace_module(model, nn.SiLU, SiLU) + + img = torch.zeros(opt.batch_size, 3, *opt.img_size).to(device) + model.model[-1].export = not opt.grid + model(img) # dry run + + # ONNX export + try: + import onnx + + logger.info('\nStarting ONNX export with onnx %s...' % onnx.__version__) + file_name = "export_output/onnx/" + os.path.basename(opt.weights[:-len(opt.weights.split(".")[-1])-1]) + f = file_name + f'_{opt.img_size[0]}x{opt.img_size[1]}.onnx' # filename + model.eval() + model.model[-1].concat = True + input_names = ["input_0"] + output_names = ["output_0"] # , "output_1", "output_2"] + torch.onnx.export(model, + img, + f, + verbose=False, + opset_version=opt.opset, + input_names=input_names, + output_names=output_names, + dynamic_axes=None) + logger.info("export end") + logger.info("check") + onnx_model = onnx.load(f) # load onnx model + onnx.checker.check_model(onnx_model) # check onnx model + + if opt.simplify: + try: + import onnxsim + logger.info('\nstart to simplify ONNX...') + onnx_model, check = onnxsim.simplify(onnx_model) + assert check, 'assert check failed' + except Exception as e: + logger.error(f'Simplifier failure: {e}') + else: + logger.warning("simplify disabled. It is recommend to use " + "command\"--simplify\" to simplify your export model.") + + onnx.save(onnx_model, f) + logger.info('ONNX export success, saved as %s' % f) + + with open(file_name + f"_{opt.img_size[0]}x{opt.img_size[1]}.yaml", "w") as fp: + yaml.dump({ + "input_name": input_names, + "output_name": output_names, + "names": labels, + "img_size": opt.img_size, + "batch_size": opt.batch_size + }, fp, yaml.Dumper) + logger.info(f"params saved to {file_name}_{opt.img_size[0]}x{opt.img_size[1]}.yaml") + + print("") + logger.info("############# - msg - ##############") + logger.info(f"input names : {input_names}") + logger.info(f"output names : {output_names}") + try: + logger.info(f"output shape : {model(img).shape}") + except AttributeError: + logger.info(f"output shape : {[m.shape for m in model(img)]}") + except Exception as e: + logger.error(e) + + logger.info(f"img size : {opt.img_size}") + logger.info(f"batch size : {opt.batch_size}") + logger.info(f"names : {labels}") + + except Exception as e: + logger.error('ONNX export failure: %s' % e) + + # Finish + logger.info('\nExport complete (%.2fs). Visualize with https://github.com/lutzroeder/netron.' % (time.time() - t)) diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..c4c131e --- /dev/null +++ b/requirements.txt @@ -0,0 +1,18 @@ +opencv-python +numpy +loguru +pycocotools +torch +tqdm +tabulate +pyyaml +torchvision +thop +pandas +pathlib +pillow +matplotlib +seaborn +scipy +onnx +onnx-simplifier \ No newline at end of file