Spaces:
Sleeping
Sleeping
File size: 2,216 Bytes
e1832f4 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 |
# Mikel Broström 🔥 Yolo Tracking 🧾 AGPL-3.0 license
from boxmot.utils import logger as LOGGER
from boxmot.utils.checks import RequirementsChecker
checker = RequirementsChecker()
UL_MODELS = ['yolov8', 'yolov9', 'yolov10', 'yolo11', 'yolo12', 'rtdetr', 'sam']
def is_ultralytics_model(yolo_name):
return any(yolo in str(yolo_name) for yolo in UL_MODELS)
def is_yolox_model(yolo_name):
return 'yolox' in str(yolo_name)
def default_imgsz(yolo_name):
if is_ultralytics_model(yolo_name):
return [640, 640]
elif is_yolox_model(yolo_name):
return [800, 1440]
else:
return [640, 640]
def get_yolo_inferer(yolo_model):
if is_yolox_model(yolo_model):
try:
import yolox # for linear_assignment
assert yolox.__version__
except (ImportError, AssertionError, AttributeError):
checker.check_packages(('yolox',))
checker.check_packages(('tabulate',)) # needed dependency
checker.check_packages(('thop',)) # needed dependency
from .yolox import YoloXStrategy
return YoloXStrategy
elif 'yolov8' in str(yolo_model):
# ultralytics already installed when running track.py
from .yolov8 import Yolov8Strategy
return Yolov8Strategy
elif 'rf-detr' in str(yolo_model):
try:
import rfdetr
except (ImportError, AssertionError, AttributeError):
checker.check_packages(('onnxruntime',)) # needed dependency
checker.check_packages(('rfdetr',)) # needed dependency
from .rfdetr import RFDETRStrategy
return RFDETRStrategy
elif 'yolo_nas' in str(yolo_model):
try:
import super_gradients # for linear_assignment
assert super_gradients.__version__
except (ImportError, AssertionError, AttributeError):
checker.check_packages(('super-gradients==3.1.3',)) # install
from .yolonas import YoloNASStrategy
return YoloNASStrategy
else:
LOGGER.error('Failed to infer inference mode from yolo model name')
LOGGER.error('Your model name has to contain either yolox, yolo_nas or yolov8')
exit()
|