Skip to content

Commit

Permalink
TFLite prep (ultralytics#4436)
Browse files Browse the repository at this point in the history
  • Loading branch information
glenn-jocher authored Aug 16, 2021
1 parent bf5738f commit 5cd1e1d
Show file tree
Hide file tree
Showing 2 changed files with 10 additions and 6 deletions.
3 changes: 2 additions & 1 deletion detect.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
13 changes: 8 additions & 5 deletions utils/general.py
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand Down

0 comments on commit 5cd1e1d

Please sign in to comment.