Unverified Commit f3e3f760 authored by Glenn Jocher's avatar Glenn Jocher Committed by GitHub
Browse files

TFLite prep (#4436)

parent dbc06ce2
......@@ -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
......@@ -203,11 +203,14 @@ def check_requirements(requirements='requirements.txt', exclude=()):
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
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment