Add suffix checks (#4711)
* Add suffix checks * Cleanup * Cleanup2 * Cleanup3
This commit is contained in:
1 parent
8e5f9ddbdb
commit
a2b3c71636
7 files changed
+41
-21
No files matched your search
@@ -21,8 +21,9 @@ sys.path.append(FILE.parents[0].as_posix()) # add yolov5/ to path
|
||||
|
||||
from models.experimental import attempt_load
|
||||
from utils.datasets import LoadStreams, LoadImages
|
||||
from utils.general import check_img_size, check_requirements, check_imshow, colorstr, is_ascii, non_max_suppression, \
|
||||
apply_classifier, scale_coords, xyxy2xywh, strip_optimizer, set_logging, increment_path, save_one_box
|
||||
from utils.general import check_img_size, check_imshow, check_requirements, check_suffix, colorstr, is_ascii, \
|
||||
non_max_suppression, apply_classifier, scale_coords, xyxy2xywh, strip_optimizer, set_logging, increment_path, \
|
||||
save_one_box
|
||||
from utils.plots import Annotator, colors
|
||||
from utils.torch_utils import select_device, load_classifier, time_sync
|
||||
|
||||
@@ -68,8 +69,9 @@ def run(weights='yolov5s.pt', # model.pt path(s)
|
||||
|
||||
# Load model
|
||||
w = weights[0] if isinstance(weights, list) else weights
|
||||
classify, suffix = False, Path(w).suffix.lower()
|
||||
pt, onnx, tflite, pb, saved_model = (suffix == x for x in ['.pt', '.onnx', '.tflite', '.pb', '']) # backend
|
||||
classify, suffix, suffixes = False, Path(w).suffix.lower(), ['.pt', '.onnx', '.tflite', '.pb', '']
|
||||
check_suffix(w, suffixes) # check weights have acceptable suffix
|
||||
pt, onnx, tflite, pb, saved_model = (suffix == x for x in suffixes) # backend booleans
|
||||
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
|
||||
|
||||
Reference in new issue
Block a user