From 416a8e0108398af4acac829e7f7ad260fd5478eb Mon Sep 17 00:00:00 2001 From: Marcos Luciano Date: Tue, 28 Jun 2022 01:14:44 -0300 Subject: [PATCH] Update gen_wts --- readme.md | 24 ++-- utils/gen_wts_yoloV5.py | 270 +++++++++++++++++++--------------------- utils/gen_wts_yolor.py | 45 ++++--- 3 files changed, 169 insertions(+), 170 deletions(-) diff --git a/readme.md b/readme.md index bf6a309..56ef4b2 100644 --- a/readme.md +++ b/readme.md @@ -478,7 +478,7 @@ wget https://github.com/ultralytics/yolov5/releases/download/v6.1/yolov5n.pt #### 4. Generate cfg and wts files (example for YOLOv5n) ``` -python3 gen_wts_yoloV5.py -w yolov5n.pt +python3 gen_wts_yoloV5.py -w yolov5n.pt -c models/yolov5n.yaml ``` #### 5. Copy generated cfg and wts files to DeepStream-Yolo folder @@ -567,28 +567,22 @@ deepstream-app -c deepstream_app_config.txt -c or --yaml ``` -* Model width **(default = 640 / 1280 [P6])** +* Inference size [size] or [height , weight] **(default [640] / [1280] if --p6)** ``` --mw or --width +-s or --size ``` -* Model height **(default = 640 / 1280 [P6])** +Example for 1280 ``` --mh or --height +-s 1280 ``` -* Model channels **(default = 3)** +or ``` --mc or --channels -``` - -* P6 model - -``` ---p6 +-s 1280 1280 ``` ## @@ -687,7 +681,9 @@ To change the `iou-threshold`, `score-threshold` and `topk` values, modify the ` **NOTE**: Lower `topk` values will result in more performance. -**NOTE**: Make sure to set cluster-mode=4 and pre-cluster-threshold=0 in config_infer file. +**NOTE**: Make sure to set cluster-mode=4 in config_infer file. + +**NOTE**: You are still able to change the `pre-cluster-threshold` values in the `config_infer.txt` file. ``` [property] diff --git a/utils/gen_wts_yoloV5.py b/utils/gen_wts_yoloV5.py index 6c176b6..f979a40 100644 --- a/utils/gen_wts_yoloV5.py +++ b/utils/gen_wts_yoloV5.py @@ -17,70 +17,62 @@ class YoloLayers(): break return route - def route(self, layers=""): - return "\n[route]\n" + \ - "layers=%s\n" % layers + def route(self, layers=''): + return '\n[route]\n' + \ + 'layers=%s\n' % layers def reorg(self): - return "\n[reorg]\n" + return '\n[reorg]\n' - def shortcut(self, route=-1, activation="linear"): - return "\n[shortcut]\n" + \ - "from=%d\n" % route + \ - "activation=%s\n" % activation + def shortcut(self, route=-1, activation='linear'): + return '\n[shortcut]\n' + \ + 'from=%d\n' % route + \ + 'activation=%s\n' % activation def maxpool(self, stride=1, size=1): - return "\n[maxpool]\n" + \ - "stride=%d\n" % stride + \ - "size=%d\n" % size + return '\n[maxpool]\n' + \ + 'stride=%d\n' % stride + \ + 'size=%d\n' % size def upsample(self, stride=1): - return "\n[upsample]\n" + \ - "stride=%d\n" % stride + return '\n[upsample]\n' + \ + 'stride=%d\n' % stride - def convolutional(self, bn=False, size=1, stride=1, pad=1, filters=1, groups=1, activation="linear"): - b = "batch_normalize=1\n" if bn is True else "" - g = "groups=%d\n" % groups if groups > 1 else "" - return "\n[convolutional]\n" + \ + def convolutional(self, bn=False, size=1, stride=1, pad=1, filters=1, groups=1, activation='linear'): + b = 'batch_normalize=1\n' if bn is True else '' + g = 'groups=%d\n' % groups if groups > 1 else '' + return '\n[convolutional]\n' + \ b + \ - "filters=%d\n" % filters + \ - "size=%d\n" % size + \ - "stride=%d\n" % stride + \ - "pad=%d\n" % pad + \ + 'filters=%d\n' % filters + \ + 'size=%d\n' % size + \ + 'stride=%d\n' % stride + \ + 'pad=%d\n' % pad + \ g + \ - "activation=%s\n" % activation + 'activation=%s\n' % activation - def yolo(self, mask="", anchors="", classes=80, num=3): - return "\n[yolo]\n" + \ - "mask=%s\n" % mask + \ - "anchors=%s\n" % anchors + \ - "classes=%d\n" % classes + \ - "num=%d\n" % num + \ - "scale_x_y=2.0\n" + \ - "beta_nms=0.6\n" + \ - "new_coords=1\n" + def yolo(self, mask='', anchors='', classes=80, num=3): + return '\n[yolo]\n' + \ + 'mask=%s\n' % mask + \ + 'anchors=%s\n' % anchors + \ + 'classes=%d\n' % classes + \ + 'num=%d\n' % num + \ + 'scale_x_y=2.0\n' + \ + 'beta_nms=0.6\n' + \ + 'new_coords=1\n' def parse_args(): - parser = argparse.ArgumentParser(description="PyTorch YOLOv5 conversion") - parser.add_argument("-w", "--weights", required=True, help="Input weights (.pt) file path (required)") - parser.add_argument("-c", "--yaml", help="Input cfg (.yaml) file path") - parser.add_argument("-mw", "--width", type=int, help="Model width (default = 640 / 1280 [P6])") - parser.add_argument("-mh", "--height", type=int, help="Model height (default = 640 / 1280 [P6])") - parser.add_argument("-mc", "--channels", type=int, help="Model channels (default = 3)") - parser.add_argument("--p6", action="store_true", help="P6 model") + parser = argparse.ArgumentParser(description='PyTorch YOLOv5 conversion') + parser.add_argument('-w', '--weights', required=True, help='Input weights (.pt) file path (required)') + parser.add_argument('-c', '--yaml', help='Input cfg (.yaml) file path') + parser.add_argument( + '-s', '--size', nargs='+', type=int, default=[640], help='Inference size [H,W] (default [640])') args = parser.parse_args() if not os.path.isfile(args.weights): - raise SystemExit("Invalid weights file") + raise SystemExit('Invalid weights file') if not args.yaml: - args.yaml = "" - if not args.width: - args.width = 1280 if args.p6 else 640 - if not args.height: - args.height = 1280 if args.p6 else 640 - if not args.channels: - args.channels = 3 - return args.weights, args.yaml, args.width, args.height, args.channels, args.p6 + args.yaml = '' + return args.weights, args.yaml, args.size def get_width(x, gw, divisor=8): @@ -96,40 +88,40 @@ def get_depth(x, gd): return max(r, 1) -pt_file, yaml_file, model_width, model_height, model_channels, p6 = parse_args() +pt_file, yaml_file, inference_size = parse_args() -model_name = pt_file.split(".pt")[0] -wts_file = model_name + ".wts" if "yolov5" in model_name else "yolov5_" + model_name + ".wts" -cfg_file = model_name + ".cfg" if "yolov5" in model_name else "yolov5_" + model_name + ".cfg" +model_name = os.path.basename(pt_file).split('.pt')[0] +wts_file = model_name + '.wts' if 'yolov5' in model_name else 'yolov5_' + model_name + '.wts' +cfg_file = model_name + '.cfg' if 'yolov5' in model_name else 'yolov5_' + model_name + '.cfg' -if yaml_file == "": - yaml_file = "models/" + model_name + ".yaml" +if yaml_file == '': + yaml_file = 'models/' + model_name + '.yaml' if not os.path.isfile(yaml_file): - yaml_file = "models/hub/" + model_name + ".yaml" + yaml_file = 'models/hub/' + model_name + '.yaml' if not os.path.isfile(yaml_file): - raise SystemExit("YAML file not found") + raise SystemExit('YAML file not found') elif not os.path.isfile(yaml_file): - raise SystemExit("Invalid YAML file") + raise SystemExit('Invalid YAML file') -device = select_device("cpu") -model = torch.load(pt_file, map_location=device)["model"].float() +device = select_device('cpu') +model = torch.load(pt_file, map_location=device)['model'].float() anchor_grid = model.model[-1].anchors * model.model[-1].stride[..., None, None] -delattr(model.model[-1], "anchor_grid") -model.model[-1].register_buffer("anchor_grid", anchor_grid) +delattr(model.model[-1], 'anchor_grid') +model.model[-1].register_buffer('anchor_grid', anchor_grid) model.to(device).eval() nc = 0 -anchors = "" +anchors = '' masks = [] yolo_idx = 0 spp_idx = 0 for k, v in model.state_dict().items(): - if "anchor_grid" in k: - yolo_idx = int(k.split(".")[1]) + if 'anchor_grid' in k: + yolo_idx = int(k.split('.')[1]) vr = v.cpu().numpy().tolist() a = v.reshape(-1).cpu().numpy().astype(float).tolist() anchors = str(a)[1:-1] @@ -140,116 +132,116 @@ for k, v in model.state_dict().items(): mask.append(num) num += 1 masks.append(mask) - elif ".%d.m.0.weight" % yolo_idx in k: + elif '.%d.m.0.weight' % yolo_idx in k: vr = v.cpu().numpy().tolist() nc = int((len(vr) / len(masks[0])) - 5) -with open(cfg_file, "w") as c: - with open(yaml_file, "r", encoding="utf-8") as f: - c.write("[net]\n") - c.write("width=%d\n" % model_width) - c.write("height=%d\n" % model_height) - c.write("channels=%d\n" % model_channels) - c.write("letter_box=1\n") +with open(cfg_file, 'w') as c: + with open(yaml_file, 'r', encoding='utf-8') as f: + c.write('[net]\n') + c.write('width=%d\n' % (inference_size[0] if len(inference_size) == 1 else inference_size[1])) + c.write('height=%d\n' % inference_size[0]) + c.write('channels=3\n') + c.write('letter_box=1\n') depth_multiple = 0 width_multiple = 0 layers = [] yoloLayers = YoloLayers() f = yaml.load(f, Loader=yaml.FullLoader) for topic in f: - if topic == "depth_multiple": + if topic == 'depth_multiple': depth_multiple = f[topic] - elif topic == "width_multiple": + elif topic == 'width_multiple': width_multiple = f[topic] - elif topic == "backbone" or topic == "head": + elif topic == 'backbone' or topic == 'head': for v in f[topic]: - if v[2] == "Focus": - layer = "\n# Focus\n" + if v[2] == 'Focus': + layer = '\n# Focus\n' blocks = 0 layer += yoloLayers.reorg() blocks += 1 layer += yoloLayers.convolutional(bn=True, filters=get_width(v[3][0], width_multiple), size=v[3][1], - activation="silu") + activation='silu') blocks += 1 layers.append([layer, blocks]) - if v[2] == "Conv": - layer = "\n# Conv\n" + if v[2] == 'Conv': + layer = '\n# Conv\n' blocks = 0 layer += yoloLayers.convolutional(bn=True, filters=get_width(v[3][0], width_multiple), size=v[3][1], - stride=v[3][2], activation="silu") + stride=v[3][2], activation='silu') blocks += 1 layers.append([layer, blocks]) - elif v[2] == "C3": - layer = "\n# C3\n" + elif v[2] == 'C3': + layer = '\n# C3\n' blocks = 0 # SPLIT layer += yoloLayers.convolutional(bn=True, filters=get_width(v[3][0], width_multiple) / 2, - activation="silu") + activation='silu') blocks += 1 - layer += yoloLayers.route(layers="-2") + layer += yoloLayers.route(layers='-2') blocks += 1 layer += yoloLayers.convolutional(bn=True, filters=get_width(v[3][0], width_multiple) / 2, - activation="silu") + activation='silu') blocks += 1 # Residual Block if len(v[3]) == 1 or v[3][1] is True: for _ in range(get_depth(v[1], depth_multiple)): layer += yoloLayers.convolutional(bn=True, filters=get_width(v[3][0], width_multiple) / 2, - activation="silu") + activation='silu') blocks += 1 layer += yoloLayers.convolutional(bn=True, filters=get_width(v[3][0], width_multiple) / 2, - size=3, activation="silu") + size=3, activation='silu') blocks += 1 layer += yoloLayers.shortcut(route=-3) blocks += 1 # Merge - layer += yoloLayers.route(layers="-1, -%d" % (3 * get_depth(v[1], depth_multiple) + 3)) + layer += yoloLayers.route(layers='-1, -%d' % (3 * get_depth(v[1], depth_multiple) + 3)) blocks += 1 else: for _ in range(get_depth(v[1], depth_multiple)): layer += yoloLayers.convolutional(bn=True, filters=get_width(v[3][0], width_multiple) / 2, - activation="silu") + activation='silu') blocks += 1 layer += yoloLayers.convolutional(bn=True, filters=get_width(v[3][0], width_multiple) / 2, - size=3, activation="silu") + size=3, activation='silu') blocks += 1 # Merge - layer += yoloLayers.route(layers="-1, -%d" % (2 * get_depth(v[1], depth_multiple) + 3)) + layer += yoloLayers.route(layers='-1, -%d' % (2 * get_depth(v[1], depth_multiple) + 3)) blocks += 1 # Transition layer += yoloLayers.convolutional(bn=True, filters=get_width(v[3][0], width_multiple), - activation="silu") + activation='silu') blocks += 1 layers.append([layer, blocks]) - elif v[2] == "SPP": + elif v[2] == 'SPP': spp_idx = len(layers) - layer = "\n# SPP\n" + layer = '\n# SPP\n' blocks = 0 layer += yoloLayers.convolutional(bn=True, filters=get_width(v[3][0], width_multiple) / 2, - activation="silu") + activation='silu') blocks += 1 layer += yoloLayers.maxpool(size=v[3][1][0]) blocks += 1 - layer += yoloLayers.route(layers="-2") + layer += yoloLayers.route(layers='-2') blocks += 1 layer += yoloLayers.maxpool(size=v[3][1][1]) blocks += 1 - layer += yoloLayers.route(layers="-4") + layer += yoloLayers.route(layers='-4') blocks += 1 layer += yoloLayers.maxpool(size=v[3][1][2]) blocks += 1 - layer += yoloLayers.route(layers="-6, -5, -3, -1") + layer += yoloLayers.route(layers='-6, -5, -3, -1') blocks += 1 layer += yoloLayers.convolutional(bn=True, filters=get_width(v[3][0], width_multiple), - activation="silu") + activation='silu') blocks += 1 layers.append([layer, blocks]) - elif v[2] == "SPPF": + elif v[2] == 'SPPF': spp_idx = len(layers) - layer = "\n# SPPF\n" + layer = '\n# SPPF\n' blocks = 0 layer += yoloLayers.convolutional(bn=True, filters=get_width(v[3][0], width_multiple) / 2, - activation="silu") + activation='silu') blocks += 1 layer += yoloLayers.maxpool(size=v[3][1]) blocks += 1 @@ -257,35 +249,35 @@ with open(cfg_file, "w") as c: blocks += 1 layer += yoloLayers.maxpool(size=v[3][1]) blocks += 1 - layer += yoloLayers.route(layers="-4, -3, -2, -1") + layer += yoloLayers.route(layers='-4, -3, -2, -1') blocks += 1 layer += yoloLayers.convolutional(bn=True, filters=get_width(v[3][0], width_multiple), - activation="silu") + activation='silu') blocks += 1 layers.append([layer, blocks]) - elif v[2] == "nn.Upsample": - layer = "\n# nn.Upsample\n" + elif v[2] == 'nn.Upsample': + layer = '\n# nn.Upsample\n' blocks = 0 layer += yoloLayers.upsample(stride=v[3][1]) blocks += 1 layers.append([layer, blocks]) - elif v[2] == "Concat": + elif v[2] == 'Concat': route = v[0][1] route = yoloLayers.get_route(route, layers) if route > 0 else \ yoloLayers.get_route(len(layers) + route, layers) - layer = "\n# Concat\n" + layer = '\n# Concat\n' blocks = 0 - layer += yoloLayers.route(layers="-1, %d" % (route - 1)) + layer += yoloLayers.route(layers='-1, %d' % (route - 1)) blocks += 1 layers.append([layer, blocks]) - elif v[2] == "Detect": + elif v[2] == 'Detect': for i, n in enumerate(v[0]): route = yoloLayers.get_route(n, layers) - layer = "\n# Detect\n" + layer = '\n# Detect\n' blocks = 0 - layer += yoloLayers.route(layers="%d" % (route - 1)) + layer += yoloLayers.route(layers='%d' % (route - 1)) blocks += 1 - layer += yoloLayers.convolutional(filters=((nc + 5) * len(masks[i])), activation="logistic") + layer += yoloLayers.convolutional(filters=((nc + 5) * len(masks[i])), activation='logistic') blocks += 1 layer += yoloLayers.yolo(mask=str(masks[i])[1:-1], anchors=anchors, classes=nc, num=num) blocks += 1 @@ -293,44 +285,44 @@ with open(cfg_file, "w") as c: for layer in layers: c.write(layer[0]) -with open(wts_file, "w") as f: - wts_write = "" +with open(wts_file, 'w') as f: + wts_write = '' conv_count = 0 - cv1 = "" - cv3 = "" + cv1 = '' + cv3 = '' cv3_idx = 0 for k, v in model.state_dict().items(): - if "num_batches_tracked" not in k and "anchors" not in k and "anchor_grid" not in k: + if 'num_batches_tracked' not in k and 'anchors' not in k and 'anchor_grid' not in k: vr = v.reshape(-1).cpu().numpy() - idx = int(k.split(".")[1]) - if ".cv1." in k and ".m." not in k and idx != spp_idx: - cv1 += "{} {} ".format(k, len(vr)) + idx = int(k.split('.')[1]) + if '.cv1.' in k and '.m.' not in k and idx != spp_idx: + cv1 += '{} {} '.format(k, len(vr)) for vv in vr: - cv1 += " " - cv1 += struct.pack(">f", float(vv)).hex() - cv1 += "\n" + cv1 += ' ' + cv1 += struct.pack('>f', float(vv)).hex() + cv1 += '\n' conv_count += 1 - elif cv1 != "" and ".m." in k: + elif cv1 != '' and '.m.' in k: wts_write += cv1 - cv1 = "" - if ".cv3." in k: - cv3 += "{} {} ".format(k, len(vr)) + cv1 = '' + if '.cv3.' in k: + cv3 += '{} {} '.format(k, len(vr)) for vv in vr: - cv3 += " " - cv3 += struct.pack(">f", float(vv)).hex() - cv3 += "\n" + cv3 += ' ' + cv3 += struct.pack('>f', float(vv)).hex() + cv3 += '\n' cv3_idx = idx conv_count += 1 - elif cv3 != "" and cv3_idx != idx: + elif cv3 != '' and cv3_idx != idx: wts_write += cv3 - cv3 = "" + cv3 = '' cv3_idx = 0 - if ".cv3." not in k and not (".cv1." in k and ".m." not in k and idx != spp_idx): - wts_write += "{} {} ".format(k, len(vr)) + if '.cv3.' not in k and not ('.cv1.' in k and '.m.' not in k and idx != spp_idx): + wts_write += '{} {} '.format(k, len(vr)) for vv in vr: - wts_write += " " - wts_write += struct.pack(">f", float(vv)).hex() - wts_write += "\n" + wts_write += ' ' + wts_write += struct.pack('>f', float(vv)).hex() + wts_write += '\n' conv_count += 1 - f.write("{}\n".format(conv_count)) + f.write('{}\n'.format(conv_count)) f.write(wts_write) diff --git a/utils/gen_wts_yolor.py b/utils/gen_wts_yolor.py index 95def7a..cfc6043 100644 --- a/utils/gen_wts_yolor.py +++ b/utils/gen_wts_yolor.py @@ -7,39 +7,50 @@ from models.models import Darknet def parse_args(): - parser = argparse.ArgumentParser(description="PyTorch YOLOR conversion (main branch)") - parser.add_argument("-w", "--weights", required=True, help="Input weights (.pt) file path (required)") - parser.add_argument("-c", "--cfg", required=True, help="Input cfg (.cfg) file path (required)") + parser = argparse.ArgumentParser(description='PyTorch YOLOR conversion (main branch)') + parser.add_argument('-w', '--weights', required=True, help='Input weights (.pt) file path (required)') + parser.add_argument('-c', '--cfg', help='Input cfg (.cfg) file path') args = parser.parse_args() if not os.path.isfile(args.weights): - raise SystemExit("Invalid weights file") + raise SystemExit('Invalid weights file') if not os.path.isfile(args.cfg): - raise SystemExit("Invalid cfg file") + raise SystemExit('Invalid cfg file') return args.weights, args.cfg pt_file, cfg_file = parse_args() -wts_file = "%s.wts" % cfg_file.rsplit("/")[1].split(".cfg")[0] -device = select_device("cpu") +model_name = os.path.basename(pt_file).split('.pt')[0] +wts_file = model_name + '.wts' if 'yolor' in model_name else 'yolor_' + model_name + '.wts' +new_cfg_file = model_name + '.cfg' if 'yolor' in model_name else 'yolor_' + model_name + '.cfg' + +if cfg_file == '': + cfg_file = 'cfg/' + model_name + '.cfg' + if not os.path.isfile(cfg_file): + raise SystemExit('CFG file not found') +elif not os.path.isfile(cfg_file): + raise SystemExit('Invalid CFG file') + +device = select_device('cpu') model = Darknet(cfg_file).to(device) -model.load_state_dict(torch.load(pt_file, map_location=device)["model"]) +model.load_state_dict(torch.load(pt_file, map_location=device)['model']) model.to(device).eval() -with open(wts_file, "w") as f: - wts_write = "" +with open(wts_file, 'w') as f: + wts_write = '' conv_count = 0 for k, v in model.state_dict().items(): - if "num_batches_tracked" not in k: + if 'num_batches_tracked' not in k: vr = v.reshape(-1).cpu().numpy() - wts_write += "{} {} ".format(k, len(vr)) + wts_write += '{} {} '.format(k, len(vr)) for vv in vr: - wts_write += " " - wts_write += struct.pack(">f", float(vv)).hex() - wts_write += "\n" + wts_write += ' ' + wts_write += struct.pack('>f', float(vv)).hex() + wts_write += '\n' conv_count += 1 - f.write("{}\n".format(conv_count)) + f.write('{}\n'.format(conv_count)) f.write(wts_write) -os.system("cp %s ./" % cfg_file) +if not os.path.isfile(new_cfg_file): + os.system('cp %s %s' % (cfg_file, new_cfg_file))