From 2c6317547a46a2dfc414fe1a5886fb9f63c14bf4 Mon Sep 17 00:00:00 2001 From: Diego Montes <54745152+d57montes@users.noreply.github.com> Date: Fri, 10 Dec 2021 12:24:32 -0500 Subject: [PATCH] Add nms and agnostic nms to export.py (#5938) * add nms and agnostic nms to export.py * fix agnostic implies nms * reorder args to group TF args * PEP8 120 char Co-authored-by: Glenn Jocher --- export.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/export.py b/export.py index 4f83c75c89a0..7feb525711e8 100644 --- a/export.py +++ b/export.py @@ -328,6 +328,8 @@ def run(data=ROOT / 'data/coco128.yaml', # 'dataset.yaml path' opset=14, # ONNX: opset version verbose=False, # TensorRT: verbose log workspace=4, # TensorRT: workspace size (GB) + nms=False, # TF: add NMS to model + agnostic_nms=False, # TF: add agnostic NMS to model topk_per_class=100, # TF.js NMS: topk per class to keep topk_all=100, # TF.js NMS: topk for all classes to keep iou_thres=0.45, # TF.js NMS: IoU threshold @@ -381,9 +383,9 @@ def run(data=ROOT / 'data/coco128.yaml', # 'dataset.yaml path' if any(tf_exports): pb, tflite, tfjs = tf_exports[1:] assert not (tflite and tfjs), 'TFLite and TF.js models must be exported separately, please pass only one type.' - model = export_saved_model(model, im, file, dynamic, tf_nms=tfjs, agnostic_nms=tfjs, - topk_per_class=topk_per_class, topk_all=topk_all, conf_thres=conf_thres, - iou_thres=iou_thres) # keras model + model = export_saved_model(model, im, file, dynamic, tf_nms=nms or agnostic_nms or tfjs, + agnostic_nms=agnostic_nms or tfjs, topk_per_class=topk_per_class, topk_all=topk_all, + conf_thres=conf_thres, iou_thres=iou_thres) # keras model if pb or tfjs: # pb prerequisite to tfjs export_pb(model, im, file) if tflite: @@ -414,6 +416,8 @@ def parse_opt(): parser.add_argument('--opset', type=int, default=14, help='ONNX: opset version') parser.add_argument('--verbose', action='store_true', help='TensorRT: verbose log') parser.add_argument('--workspace', type=int, default=4, help='TensorRT: workspace size (GB)') + parser.add_argument('--nms', action='store_true', help='TF: add NMS to model') + parser.add_argument('--agnostic-nms', action='store_true', help='TF: add agnostic NMS to model') parser.add_argument('--topk-per-class', type=int, default=100, help='TF.js NMS: topk per class to keep') parser.add_argument('--topk-all', type=int, default=100, help='TF.js NMS: topk for all classes to keep') parser.add_argument('--iou-thres', type=float, default=0.45, help='TF.js NMS: IoU threshold')