mirror of
https://github.com/THU-MIG/yolov10.git
synced 2025-05-23 21:44:22 +08:00
b507e3a032
Co-authored-by: ChristopherRogers1991 <ChristopherRogers1991@gmail.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Sungjoo(Dennis) Hwang <48212469+Denny-Hwang@users.noreply.github.com>
31 lines
886 B
Python
31 lines
886 B
Python
# Ultralytics YOLO 🚀, AGPL-3.0 license
|
|
"""
|
|
RT-DETR model interface
|
|
"""
|
|
from ultralytics.engine.model import Model
|
|
from ultralytics.nn.tasks import RTDETRDetectionModel
|
|
|
|
from .predict import RTDETRPredictor
|
|
from .train import RTDETRTrainer
|
|
from .val import RTDETRValidator
|
|
|
|
|
|
class RTDETR(Model):
|
|
"""
|
|
RTDETR model interface.
|
|
"""
|
|
|
|
def __init__(self, model='rtdetr-l.pt') -> None:
|
|
if model and not model.split('.')[-1] in ('pt', 'yaml', 'yml'):
|
|
raise NotImplementedError('RT-DETR only supports creating from *.pt file or *.yaml file.')
|
|
super().__init__(model=model, task='detect')
|
|
|
|
@property
|
|
def task_map(self):
|
|
return {
|
|
'detect': {
|
|
'predictor': RTDETRPredictor,
|
|
'validator': RTDETRValidator,
|
|
'trainer': RTDETRTrainer,
|
|
'model': RTDETRDetectionModel}}
|