diff --git a/detect.py b/detect.py index 4549095..62a4e8a 100644 --- a/detect.py +++ b/detect.py @@ -33,6 +33,7 @@ parser.add_argument("--bench_coco", action='store_true', help="Process a stream") parser.add_argument("--coco_path", type=str, help="Path to COCO 2017 Val folder") parser.add_argument("--quiet","-q", action='store_true', help="Disable logging (except errors)") + parser.add_argument("--int8", action='store_true', help="Datatype of model int8?") args = parser.parse_args() @@ -44,10 +45,13 @@ logger.error("Please select either an input image or a stream") exit(1) - model = EdgeTPUModel(args.model, args.names, conf_thresh=args.conf_thresh, iou_thresh=args.iou_thresh) + model = EdgeTPUModel(args.model, args.names, conf_thresh=args.conf_thresh, iou_thresh=args.iou_thresh, int8=args.int8) input_size = model.get_image_size() - x = (255*np.random.random((3,*input_size))).astype(np.uint8) + if args.int8: + x = (255*np.random.random((3,*input_size))).astype(np.int8) + else: + x = (255*np.random.random((3,*input_size))).astype(np.uint8) model.forward(x) conf_thresh = 0.25