Add CLI support for SAM, RTDETR (#3253)

This commit is contained in:
Glenn Jocher
2023-06-18 22:47:33 +02:00
committed by GitHub
parent 4c2033d7c3
commit 58bccb1a9f
2 changed files with 10 additions and 2 deletions

View File

@ -366,9 +366,16 @@ def entrypoint(debug=''):
if model is None:
model = 'yolov8n.pt'
LOGGER.warning(f"WARNING ⚠️ 'model' is missing. Using default 'model={model}'.")
from ultralytics.yolo.engine.model import YOLO
overrides['model'] = model
model = YOLO(model, task=task)
if 'rtdetr' in model.lower(): # guess architecture
from ultralytics import RTDETR
model = RTDETR(model) # no task argument
elif 'sam' in model.lower():
from ultralytics import SAM
model = SAM(model)
else:
from ultralytics import YOLO
model = YOLO(model, task=task)
if isinstance(overrides.get('pretrained'), str):
model.load(overrides['pretrained'])