Follow up to #1256, fix interpolation warning in auto_autoaugment as well

3 years ago · 7cedc8d474
parent 037e5e6c09
commit 7cedc8d474
1 changed files with 7 additions and 2 deletions
--- a/timm/data/auto_augment.py
+++ b/timm/data/auto_augment.py
@ -36,11 +36,16 @@ _HPARAMS_DEFAULT = dict(
    img_mean=_FILL,
 )

-_RANDOM_INTERPOLATION = (Image.BILINEAR, Image.BICUBIC)
+if hasattr(Image, "Resampling"):
+    _RANDOM_INTERPOLATION = (Image.Resampling.BILINEAR, Image.Resampling.BICUBIC)
+    _DEFAULT_INTERPOLATION = Image.Resampling.BICUBIC
+else:
+    _RANDOM_INTERPOLATION = (Image.BILINEAR, Image.BICUBIC)
+    _DEFAULT_INTERPOLATION = Image.BICUBIC


 def _interpolation(kwargs):
-    interpolation = kwargs.pop('resample', Image.BILINEAR)
+    interpolation = kwargs.pop('resample', _DEFAULT_INTERPOLATION)
    if isinstance(interpolation, (list, tuple)):
        return random.choice(interpolation)
    else: