From bb61350a833595ffbfea3911ed527c7501a925ae Mon Sep 17 00:00:00 2001 From: Josh Date: Wed, 1 Sep 2021 09:16:23 +1200 Subject: [PATCH] add conf/iou to args --- edgetpumodel.py | 4 ++++ test_edgetpu.py | 4 +++- 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/edgetpumodel.py b/edgetpumodel.py index 7d94f55..07f0425 100644 --- a/edgetpumodel.py +++ b/edgetpumodel.py @@ -44,6 +44,9 @@ def __init__(self, model_file, names_file, conf_thresh=0.25, iou_thresh=0.45, fi self.agnostic_nms = agnostic_nms self.max_det = 1000 + logger.info("Confidence threshold: {}".format(conf_thresh)) + logger.info("IOU threshold: {}".format(iou_thresh)) + self.inference_time = None self.nms_time = None self.interpreter = None @@ -102,6 +105,7 @@ def make_interpreter(self): logger.debug("Output scale: {}".format(self.output_scale)) logger.debug("Output zero: {}".format(self.output_zero)) + logger.info("Successfully loaded {}".format(self.model_file)) def get_image_size(self): """ diff --git a/test_edgetpu.py b/test_edgetpu.py index 5859a0d..4549095 100644 --- a/test_edgetpu.py +++ b/test_edgetpu.py @@ -24,6 +24,8 @@ parser.add_argument("--model", "-m", help="weights file", required=True) parser.add_argument("--bench_speed", action='store_true', help="run speed test on dummy data") parser.add_argument("--bench_image", action='store_true', help="run detection test") + parser.add_argument("--conf_thresh", type=float, default=0.25, help="model confidence threshold") + parser.add_argument("--iou_thresh", type=float, default=0.45, help="NMS IOU threshold") parser.add_argument("--names", type=str, default='data/coco.yaml', help="Names file") parser.add_argument("--image", "-i", type=str, help="Image file to run detection on") parser.add_argument("--device", type=int, default=0, help="Image capture device to run live detection") @@ -42,7 +44,7 @@ logger.error("Please select either an input image or a stream") exit(1) - model = EdgeTPUModel(args.model, args.names) + model = EdgeTPUModel(args.model, args.names, conf_thresh=args.conf_thresh, iou_thresh=args.iou_thresh) input_size = model.get_image_size() x = (255*np.random.random((3,*input_size))).astype(np.uint8)