From be1722b4b85f0b283b60c57c14a8d14ab7ec03e7 Mon Sep 17 00:00:00 2001 From: Laughing <61612323+Laughing-q@users.noreply.github.com> Date: Tue, 25 Jun 2024 17:56:26 +0800 Subject: [PATCH] Add warning for inference end2end models with `augment` arg (#13958) Co-authored-by: UltralyticsAssistant Co-authored-by: Glenn Jocher --- tests/test_python.py | 1 + ultralytics/nn/tasks.py | 10 ++++++++-- 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/tests/test_python.py b/tests/test_python.py index 8c1695a7..b06969aa 100644 --- a/tests/test_python.py +++ b/tests/test_python.py @@ -585,4 +585,5 @@ def test_yolov10(): # train/val/predict model.train(data="coco8.yaml", epochs=1, imgsz=32, close_mosaic=1, cache="disk") model.val(data="coco8.yaml", imgsz=32) + model.predict(imgsz=32, save_txt=True, save_crop=True, augment=True) model(SOURCE) diff --git a/ultralytics/nn/tasks.py b/ultralytics/nn/tasks.py index b6de9359..d52ecad7 100644 --- a/ultralytics/nn/tasks.py +++ b/ultralytics/nn/tasks.py @@ -151,8 +151,8 @@ class BaseModel(nn.Module): def _predict_augment(self, x): """Perform augmentations on input image x and return augmented inference.""" LOGGER.warning( - f"WARNING ⚠️ {self.__class__.__name__} does not support augmented inference yet. " - f"Reverting to single-scale inference instead." + f"WARNING ⚠️ {self.__class__.__name__} does not support 'augment=True' prediction. " + f"Reverting to single-scale prediction." ) return self._predict_once(x) @@ -337,6 +337,12 @@ class DetectionModel(BaseModel): def _predict_augment(self, x): """Perform augmentations on input image x and return augmented inference and train outputs.""" + if self.end2end: + LOGGER.warning( + "WARNING ⚠️ End2End model does not support 'augment=True' prediction. " + "Reverting to single-scale prediction." + ) + return self._predict_once(x) img_size = x.shape[-2:] # height, width s = [1, 0.83, 0.67] # scales f = [None, 3, None] # flips (2-ud, 3-lr)