From 6f0284763b0f66467dc04e5a5d87e5a68d1d49cd Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Wed, 21 Sep 2022 19:53:26 +0200
Subject: [PATCH] Update albumentations (#9503)

* Add `RandomResizedCrop(ratio)`

* Update ratio

* Update ratio

* Update ratio

* Update ratio

* Update ratio

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Create augmentations.py

Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Update augmentations.py

Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com>

Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
 utils/augmentations.py | 27 +++++++++++++++------------
 utils/dataloaders.py   |  2 +-
 2 files changed, 16 insertions(+), 13 deletions(-)

diff --git a/utils/augmentations.py b/utils/augmentations.py
index f49110f43c6a..7c8e0bcdede6 100644
--- a/utils/augmentations.py
+++ b/utils/augmentations.py
@@ -21,7 +21,7 @@
 
 class Albumentations:
     # YOLOv5 Albumentations class (optional, only used if package is installed)
-    def __init__(self):
+    def __init__(self, size=640):
         self.transform = None
         prefix = colorstr('albumentations: ')
         try:
@@ -29,6 +29,7 @@ def __init__(self):
             check_version(A.__version__, '1.0.3', hard=True)  # version requirement
 
             T = [
+                A.RandomResizedCrop(height=size, width=size, scale=(0.8, 1.0), ratio=(0.9, 1.11), p=0.0),
                 A.Blur(p=0.01),
                 A.MedianBlur(p=0.01),
                 A.ToGray(p=0.01),
@@ -303,15 +304,17 @@ def box_candidates(box1, box2, wh_thr=2, ar_thr=100, area_thr=0.1, eps=1e-16):
     return (w2 > wh_thr) & (h2 > wh_thr) & (w2 * h2 / (w1 * h1 + eps) > area_thr) & (ar < ar_thr)  # candidates
 
 
-def classify_albumentations(augment=True,
-                            size=224,
-                            scale=(0.08, 1.0),
-                            hflip=0.5,
-                            vflip=0.0,
-                            jitter=0.4,
-                            mean=IMAGENET_MEAN,
-                            std=IMAGENET_STD,
-                            auto_aug=False):
+def classify_albumentations(
+        augment=True,
+        size=224,
+        scale=(0.08, 1.0),
+        ratio=(0.75, 1.0 / 0.75),  # 0.75, 1.33
+        hflip=0.5,
+        vflip=0.0,
+        jitter=0.4,
+        mean=IMAGENET_MEAN,
+        std=IMAGENET_STD,
+        auto_aug=False):
     # YOLOv5 classification Albumentations (optional, only used if package is installed)
     prefix = colorstr('albumentations: ')
     try:
@@ -319,7 +322,7 @@ def classify_albumentations(augment=True,
         from albumentations.pytorch import ToTensorV2
         check_version(A.__version__, '1.0.3', hard=True)  # version requirement
         if augment:  # Resize and crop
-            T = [A.RandomResizedCrop(height=size, width=size, scale=scale)]
+            T = [A.RandomResizedCrop(height=size, width=size, scale=scale, ratio=ratio)]
             if auto_aug:
                 # TODO: implement AugMix, AutoAug & RandAug in albumentation
                 LOGGER.info(f'{prefix}auto augmentations are currently not supported')
@@ -338,7 +341,7 @@ def classify_albumentations(augment=True,
         return A.Compose(T)
 
     except ImportError:  # package not installed, skip
-        pass
+        LOGGER.warning(f'{prefix}⚠️ not found, install with `pip install albumentations` (recommended)')
     except Exception as e:
         LOGGER.info(f'{prefix}{e}')
 
diff --git a/utils/dataloaders.py b/utils/dataloaders.py
index 5b03b4eb9759..ee79bd0bc5a5 100644
--- a/utils/dataloaders.py
+++ b/utils/dataloaders.py
@@ -404,7 +404,7 @@ def __init__(self,
         self.mosaic_border = [-img_size // 2, -img_size // 2]
         self.stride = stride
         self.path = path
-        self.albumentations = Albumentations() if augment else None
+        self.albumentations = Albumentations(size=img_size) if augment else None
 
         try:
             f = []  # image files