From 1a7dce1e470e8514efd806620b10b310025e564f Mon Sep 17 00:00:00 2001 From: Marcos Luciano Date: Sun, 26 Mar 2023 18:59:14 -0300 Subject: [PATCH] Update PPYOLOE conversion --- utils/gen_wts_ppyoloe.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/utils/gen_wts_ppyoloe.py b/utils/gen_wts_ppyoloe.py index d54b088..058e705 100644 --- a/utils/gen_wts_ppyoloe.py +++ b/utils/gen_wts_ppyoloe.py @@ -406,6 +406,8 @@ with open(wts_file, 'w') as fw, open(cfg_file, 'w') as fc: if yolo_head == 'PPYOLOEHead': layers.fc.write('\n# PPYOLOEHead\n') + reg_max = model.yolo_head.reg_max + 1 if hasattr(model.yolo_head, 'reg_max') else model.yolo_head.reg_range[1] + for i, feat in enumerate(layers.neck_pan_feats): if i > 0: layers.AvgPool2d(route=feat) @@ -416,7 +418,7 @@ with open(wts_file, 'w') as fw, open(cfg_file, 'w') as fc: layers.Shuffle(reshape=[model.yolo_head.num_classes, 'hw'], output='cls') layers.ESEAttn(model.yolo_head.stem_reg[i], route=-7) layers.Conv2D(model.yolo_head.pred_reg[i]) - layers.Shuffle(reshape=[4, model.yolo_head.reg_max + 1, 'hw'], transpose2=[1, 0, 2]) + layers.Shuffle(reshape=[4, reg_max + 1, 'hw'], transpose2=[1, 0, 2]) layers.SoftMax(0) layers.Conv2D(model.yolo_head.proj_conv) layers.Shuffle(reshape=['h', 'w'], output='reg')