Add `fraction`, `single_cls` and `classes` to `RTDETRDataset` (#17633)

Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
6 days ago · 35b6fed814
parent 14495a46f7
commit 35b6fed814
1 changed files with 3 additions and 0 deletions
--- a/ultralytics/models/rtdetr/train.py
+++ b/ultralytics/models/rtdetr/train.py
@ -68,8 +68,11 @@ class RTDETRTrainer(DetectionTrainer):
            hyp=self.args,
            rect=False,
            cache=self.args.cache or None,
+            single_cls=self.args.single_cls or False,
            prefix=colorstr(f"{mode}: "),
+            classes=self.args.classes,
            data=self.data,
+            fraction=self.args.fraction if mode == "train" else 1.0,
        )

    def get_validator(self):