diff --git a/detect.py b/detect.py index 49ebbe96c068220744e8275d5174150bc42922e7..cdac4f2137901cbe5c817ec265beb3a2a7cc1703 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 850ca6ba0b1f87aea14f63d5fa85cd8053858c43..0b6e8fc7fb9a7da50174725573723e9f59ef54eb 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