Add warning for inference end2end models with `augment` arg (#13958)

Co-authored-by: UltralyticsAssistant <web@ultralytics.com>
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
pull/10163/merge
Laughing 5 months ago committed by GitHub
parent b10e0f3fa8
commit be1722b4b8
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
  1. 1
      tests/test_python.py
  2. 10
      ultralytics/nn/tasks.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)

@ -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)

Loading…
Cancel
Save