xclele/yolov10
1
0
Fork 0
mirror of https://github.com/THU-MIG/yolov10.git synced 2025-05-23 05:24:22 +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
"""
from pathlib import Path
from ultralytics.engine.model import Model
from ultralytics.nn.tasks import RTDETRDetectionModel
@ -34,7 +35,7 @@ class RTDETR(Model):
Raises:
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.")
super().__init__(model=model, task="detect")