xclele/yolov10
1
0
Fork 0
mirror of https://github.com/THU-MIG/yolov10.git synced 2025-05-24 14:44:21 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Add CLI support for SAM, RTDETR (#3253)

Browse Source
This commit is contained in:
Glenn Jocher 2023-06-18 22:47:33 +02:00 committed by GitHub
parent 4c2033d7c3
commit 58bccb1a9f
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 10 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
ultralytics/vit/sam/model.py
View File

@ -17,6 +17,7 @@ class SAM:
# Should raise AssertionError instead? # Should raise AssertionError instead?
raise NotImplementedError('Segment anything prediction requires pre-trained checkpoint') raise NotImplementedError('Segment anything prediction requires pre-trained checkpoint')
self.model = build_sam(model) self.model = build_sam(model)
self.task = 'segment' # required
self.predictor = None # reuse predictor self.predictor = None # reuse predictor
def predict(self, source, stream=False, **kwargs): def predict(self, source, stream=False, **kwargs):

11
ultralytics/yolo/cfg/__init__.py
View File

@ -366,9 +366,16 @@ def entrypoint(debug=''):
if model is None: if model is None:
model = 'yolov8n.pt' model = 'yolov8n.pt'
LOGGER.warning(f"WARNING ⚠️ 'model' is missing. Using default 'model={model}'.") LOGGER.warning(f"WARNING ⚠️ 'model' is missing. Using default 'model={model}'.")
from ultralytics.yolo.engine.model import YOLO
overrides['model'] = model 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): if isinstance(overrides.get('pretrained'), str):
model.load(overrides['pretrained']) model.load(overrides['pretrained'])