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

Do RTDETR file suffix check using pathlib instead of string manipulations (#8525)

Browse Source 
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
Robin Brown 2024-02-29 23:28:01 +00:00 committed by GitHub
parent 1db8e912a3
commit 55ce6faec8
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 2 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
ultralytics/models/rtdetr/model.py
View File

@ -6,6 +6,7 @@ hybrid encoder and IoU-aware query selection for enhanced detection accuracy.
For more information on RT-DETR, visit: https://arxiv.org/pdf/2304.08069.pdf For more information on RT-DETR, visit: https://arxiv.org/pdf/2304.08069.pdf
""" """
from pathlib import Path
from ultralytics.engine.model import Model from ultralytics.engine.model import Model
from ultralytics.nn.tasks import RTDETRDetectionModel from ultralytics.nn.tasks import RTDETRDetectionModel
@ -34,7 +35,7 @@ class RTDETR(Model):
Raises: Raises:
NotImplementedError: If the model file extension is not 'pt', 'yaml', or 'yml'. NotImplementedError: If the model file extension is not 'pt', 'yaml', or 'yml'.
""" """
if model and model.split(".")[-1] not in ("pt", "yaml", "yml"): if model and Path(model).suffix not in (".pt", ".yaml", ".yml"):
raise NotImplementedError("RT-DETR only supports creating from *.pt, *.yaml, or *.yml files.") raise NotImplementedError("RT-DETR only supports creating from *.pt, *.yaml, or *.yml files.")
super().__init__(model=model, task="detect") super().__init__(model=model, task="detect")