From ddc94a6981f985425653669671810364d6f76928 Mon Sep 17 00:00:00 2001 From: Laughing <61612323+Laughing-q@users.noreply.github.com> Date: Wed, 6 Mar 2024 01:37:33 +0800 Subject: [PATCH] OBB: Fix when training on zip data (#8680) Co-authored-by: UltralyticsAssistant Co-authored-by: Glenn Jocher --- ultralytics/engine/trainer.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/ultralytics/engine/trainer.py b/ultralytics/engine/trainer.py index f005f341..6db3eb18 100644 --- a/ultralytics/engine/trainer.py +++ b/ultralytics/engine/trainer.py @@ -129,7 +129,12 @@ class BaseTrainer: try: if self.args.task == "classify": self.data = check_cls_dataset(self.args.data) - elif self.args.data.split(".")[-1] in ("yaml", "yml") or self.args.task in ("detect", "segment", "pose"): + elif self.args.data.split(".")[-1] in ("yaml", "yml") or self.args.task in ( + "detect", + "segment", + "pose", + "obb", + ): self.data = check_det_dataset(self.args.data) if "yaml_file" in self.data: self.args.data = self.data["yaml_file"] # for validating 'yolo train data=url.zip' usage