From 35b6fed8148d659de85661a34fad7135d5092404 Mon Sep 17 00:00:00 2001 From: Mohammed Yasin <32206511+Y-T-G@users.noreply.github.com> Date: Tue, 19 Nov 2024 19:11:20 +0800 Subject: [PATCH] Add `fraction`, `single_cls` and `classes` to `RTDETRDataset` (#17633) Co-authored-by: Glenn Jocher --- ultralytics/models/rtdetr/train.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/ultralytics/models/rtdetr/train.py b/ultralytics/models/rtdetr/train.py index cb11a72794..962aeaf9bc 100644 --- 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):