From b00f4afefa2ae87e27f21baedae943b1f79a65a8 Mon Sep 17 00:00:00 2001 From: Laughing <61612323+Laughing-q@users.noreply.github.com> Date: Fri, 6 Sep 2024 22:03:15 +0800 Subject: [PATCH] Fix `_predict_augment` and add warning (#16056) Signed-off-by: UltralyticsAssistant Co-authored-by: Glenn Jocher Co-authored-by: UltralyticsAssistant --- ultralytics/nn/tasks.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/ultralytics/nn/tasks.py b/ultralytics/nn/tasks.py index fdcec302f..07a040d90 100644 --- a/ultralytics/nn/tasks.py +++ b/ultralytics/nn/tasks.py @@ -341,11 +341,8 @@ class DetectionModel(BaseModel): def _predict_augment(self, x): """Perform augmentations on input image x and return augmented inference and train outputs.""" - if getattr(self, "end2end", False): - LOGGER.warning( - "WARNING ⚠️ End2End model does not support 'augment=True' prediction. " - "Reverting to single-scale prediction." - ) + if getattr(self, "end2end", False) or self.__class__.__name__ != "DetectionModel": + LOGGER.warning("WARNING ⚠️ Model does not support 'augment=True', reverting to single-scale prediction.") return self._predict_once(x) img_size = x.shape[-2:] # height, width s = [1, 0.83, 0.67] # scales