From 5cd1e1dbef2fc2308428a4c335a029aadb26a68e Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Mon, 16 Aug 2021 17:25:06 +0200 Subject: [PATCH] TFLite prep (#4436) --- detect.py | 3 ++- utils/general.py | 13 ++++++++----- 2 files changed, 10 insertions(+), 6 deletions(-) diff --git a/detect.py b/detect.py index 49ebbe96c068..cdac4f213790 100644 --- a/detect.py +++ b/detect.py @@ -67,7 +67,8 @@ def run(weights='yolov5s.pt', # model.pt path(s) # Load model w = weights[0] if isinstance(weights, list) else weights - classify, pt, onnx = False, w.endswith('.pt'), w.endswith('.onnx') # inference type + classify, suffix = False, Path(w).suffix.lower() + pt, onnx, tflite, pb, graph_def = (suffix == x for x in ['.pt', '.onnx', '.tflite', '.pb', '']) # backend stride, names = 64, [f'class{i}' for i in range(1000)] # assign defaults if pt: model = attempt_load(weights, map_location=device) # load FP32 model diff --git a/utils/general.py b/utils/general.py index 850ca6ba0b1f..0b6e8fc7fb9a 100755 --- a/utils/general.py +++ b/utils/general.py @@ -203,11 +203,14 @@ def check_requirements(requirements='requirements.txt', exclude=()): print(emojis(s)) -def check_img_size(img_size, s=32, floor=0): - # Verify img_size is a multiple of stride s - new_size = max(make_divisible(img_size, int(s)), floor) # ceil gs-multiple - if new_size != img_size: - print(f'WARNING: --img-size {img_size} must be multiple of max stride {s}, updating to {new_size}') +def check_img_size(imgsz, s=32, floor=0): + # Verify image size is a multiple of stride s in each dimension + if isinstance(imgsz, int): # integer i.e. img_size=640 + new_size = max(make_divisible(imgsz, int(s)), floor) + else: # list i.e. img_size=[640, 480] + new_size = [max(make_divisible(x, int(s)), floor) for x in imgsz] + if new_size != imgsz: + print(f'WARNING: --img-size {imgsz} must be multiple of max stride {s}, updating to {new_size}') return new_size