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) ...@@ -67,7 +67,8 @@ def run(weights='yolov5s.pt', # model.pt path(s)
# Load model # Load model
w = weights[0] if isinstance(weights, list) else weights 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 stride, names = 64, [f'class{i}' for i in range(1000)] # assign defaults
if pt: if pt:
model = attempt_load(weights, map_location=device) # load FP32 model model = attempt_load(weights, map_location=device) # load FP32 model
......
...@@ -203,11 +203,14 @@ def check_requirements(requirements='requirements.txt', exclude=()): ...@@ -203,11 +203,14 @@ def check_requirements(requirements='requirements.txt', exclude=()):
print(emojis(s)) print(emojis(s))
def check_img_size(img_size, s=32, floor=0): def check_img_size(imgsz, s=32, floor=0):
# Verify img_size is a multiple of stride s # Verify image size is a multiple of stride s in each dimension
new_size = max(make_divisible(img_size, int(s)), floor) # ceil gs-multiple if isinstance(imgsz, int): # integer i.e. img_size=640
if new_size != img_size: new_size = max(make_divisible(imgsz, int(s)), floor)
print(f'WARNING: --img-size {img_size} must be multiple of max stride {s}, updating to {new_size}') 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 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