mirror of
https://github.com/THU-MIG/yolov10.git
synced 2025-05-23 05:24:22 +08:00
Do RTDETR file suffix check using pathlib instead of string manipulations (#8525)
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
parent
1db8e912a3
commit
55ce6faec8
@ -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")
|
||||
|
||||
|
Loading…
x
Reference in New Issue
Block a user