From 40e5ac34c951964f09e988a233a9a4c4e575ef05 Mon Sep 17 00:00:00 2001 From: PacificDou Date: Mon, 6 May 2024 20:54:56 +0200 Subject: [PATCH] Fix single-image dataset failure mode (#11694) Co-authored-by: UltralyticsAssistant Co-authored-by: Glenn Jocher --- tests/test_cli.py | 4 ++-- ultralytics/data/base.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/test_cli.py b/tests/test_cli.py index 9a31db24..71bc2e48 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -54,8 +54,8 @@ def test_export(model): def test_rtdetr(task="detect", model="yolov8n-rtdetr.yaml", data="coco8.yaml"): """Test the RTDETR functionality with the Ultralytics framework.""" - # Warning: MUST use imgsz=640 - run(f"yolo train {task} model={model} data={data} --imgsz= 160 epochs =1, cache = disk") # add coma, spaces to args + # Warning: must use imgsz=640 (note also add coma, spaces, fraction=0.25 args to test single-image training) + run(f"yolo train {task} model={model} data={data} --imgsz= 160 epochs =1, cache = disk fraction=0.25") run(f"yolo predict {task} model={model} source={ASSETS / 'bus.jpg'} imgsz=160 save save_crop save_txt") diff --git a/ultralytics/data/base.py b/ultralytics/data/base.py index 2218d31f..256a2c20 100644 --- a/ultralytics/data/base.py +++ b/ultralytics/data/base.py @@ -170,7 +170,7 @@ class BaseDataset(Dataset): if self.augment: self.ims[i], self.im_hw0[i], self.im_hw[i] = im, (h0, w0), im.shape[:2] # im, hw_original, hw_resized self.buffer.append(i) - if len(self.buffer) >= self.max_buffer_length: + if 1 < len(self.buffer) >= self.max_buffer_length: # prevent empty buffer j = self.buffer.pop(0) if self.cache != "ram": self.ims[j], self.im_hw0[j], self.im_hw[j] = None, None, None