Fix SAM CLI stem usage (#6257)

This commit is contained in:
Glenn Jocher 2023-11-14 00:39:18 +01:00 committed by GitHub
parent ec977f6ca7
commit 87233ea17c
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -405,13 +405,14 @@ def entrypoint(debug=''):
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}'.")
overrides['model'] = model overrides['model'] = model
if 'rtdetr' in model.lower(): # guess architecture stem = Path(model).stem.lower()
if 'rtdetr' in stem: # guess architecture
from ultralytics import RTDETR from ultralytics import RTDETR
model = RTDETR(model) # no task argument model = RTDETR(model) # no task argument
elif 'fastsam' in model.lower(): elif 'fastsam' in stem:
from ultralytics import FastSAM from ultralytics import FastSAM
model = FastSAM(model) model = FastSAM(model)
elif 'sam' in model.lower(): elif 'sam' in stem:
from ultralytics import SAM from ultralytics import SAM
model = SAM(model) model = SAM(model)
else: else: