From cb2d70343916d4aa10cb8790e119191f5fd38f19 Mon Sep 17 00:00:00 2001 From: Diego Montes Date: Thu, 9 Dec 2021 14:10:29 -0500 Subject: [PATCH 1/4] add nms and agnostic nms to export.py --- export.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/export.py b/export.py index 88d03a2c9475..b98e1385f687 100644 --- a/export.py +++ b/export.py @@ -325,6 +325,8 @@ def run(data=ROOT / 'data/coco128.yaml', # 'dataset.yaml path' int8=False, # CoreML/TF INT8 quantization dynamic=False, # ONNX/TF: dynamic axes simplify=False, # ONNX: simplify model + nms=False, # TF: add NMS to model + agnostic_nms=False, # TF: add agnostic NMS to model opset=12, # ONNX: opset version verbose=False, # TensorRT: verbose log workspace=4, # TensorRT: workspace size (GB) @@ -381,7 +383,7 @@ 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, + model = export_saved_model(model, im, file, dynamic, tf_nms=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 @@ -411,6 +413,8 @@ def parse_opt(): parser.add_argument('--int8', action='store_true', help='CoreML/TF INT8 quantization') parser.add_argument('--dynamic', action='store_true', help='ONNX/TF: dynamic axes') parser.add_argument('--simplify', action='store_true', help='ONNX: simplify model') + 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('--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)') From 13728d39f003dcfdae99ac124b22d048e7fdf50b Mon Sep 17 00:00:00 2001 From: Diego Montes Date: Thu, 9 Dec 2021 17:14:44 -0500 Subject: [PATCH 2/4] fix agnostic implies nms --- export.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/export.py b/export.py index b98e1385f687..39ace55b9026 100644 --- a/export.py +++ b/export.py @@ -383,7 +383,7 @@ 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=nms or tfjs, agnostic_nms=agnostic_nms or tfjs, + 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 From b17ddf279d0587adf8603de4ee835d3ddb6e4935 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Fri, 10 Dec 2021 18:16:36 +0100 Subject: [PATCH 3/4] reorder args to group TF args --- export.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/export.py b/export.py index 8f9cf9833d74..28eda670c1ac 100644 --- a/export.py +++ b/export.py @@ -325,11 +325,11 @@ def run(data=ROOT / 'data/coco128.yaml', # 'dataset.yaml path' int8=False, # CoreML/TF INT8 quantization dynamic=False, # ONNX/TF: dynamic axes simplify=False, # ONNX: simplify model - nms=False, # TF: add NMS to model - agnostic_nms=False, # TF: add agnostic NMS to model 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 @@ -413,11 +413,11 @@ def parse_opt(): parser.add_argument('--int8', action='store_true', help='CoreML/TF INT8 quantization') parser.add_argument('--dynamic', action='store_true', help='ONNX/TF: dynamic axes') parser.add_argument('--simplify', action='store_true', help='ONNX: simplify model') - 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('--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') From cf1e87194019f84355dc92eff6a4166f7a867d7a Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Fri, 10 Dec 2021 18:18:53 +0100 Subject: [PATCH 4/4] PEP8 120 char --- export.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/export.py b/export.py index 28eda670c1ac..7feb525711e8 100644 --- a/export.py +++ b/export.py @@ -383,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=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 + 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: