Repository: ultralytics/ultralytics
Pull request: `ultralytics 8.0.233` improve Classify train augmentations
<!--
copilot:all
-->
### <samp>🤖 Generated by Copilot at 24f5c97</samp>

### Summary
🧪🛠️🌈

<!--
1.  🧪 for adding a new test function
2.  🛠️ for adding new configuration options and modifying existing functions
3.  🌈 for adding support for new classification transforms and improving image processing
-->
This pull request improves the classification augmentations and transforms for the `ultralytics` package. It adds support for different torchvision versions and new configuration options for auto augmentation, random erasing and color jitter. It also updates the existing classes and functions that use the classification transforms and adds a new test function. The files affected are `ultralytics/data/augment.py`, `ultralytics/data/dataset.py`, `ultralytics/engine/predictor.py`, `ultralytics/models/yolo/classify/predict.py`, `ultralytics/utils/torch_utils.py`, `tests/test_python.py` and `ultralytics/cfg/default.yaml`.

> _`torchvision` augments_
> _classification transforms_
> _better in the fall_

### Walkthrough
*  Add and modify classification transforms for training and evaluation ([link](https://github.com/ultralytics/ultralytics/pull/4546/files?diff=unified&w=0#diff-8945a175b39a1efcabd4d2fb09b583d2f25409f05038ca9a0f1165e39ee48528L17-R25), [link](https://github.com/ultralytics/ultralytics/pull/4546/files?diff=unified&w=0#diff-8945a175b39a1efcabd4d2fb09b583d2f25409f05038ca9a0f1165e39ee48528L798-R807), [link](https://github.com/ultralytics/ultralytics/pull/4546/files?diff=unified&w=0#diff-8945a175b39a1efcabd4d2fb09b583d2f25409f05038ca9a0f1165e39ee48528R816-R978), [link](https://github.com/ultralytics/ultralytics/pull/4546/files?diff=unified&w=0#diff-8945a175b39a1efcabd4d2fb09b583d2f25409f05038ca9a0f1165e39ee48528L815-R995), [link](https://github.com/ultralytics/ultralytics/pull/4546/files?diff=unified&w=0#diff-9fa2bf7f1273283680da3daeb8bae8e20cf8d477f838575796d2bf3b415bc534L39-R39), [link](https://github.com/ultralytics/ultralytics/pull/4546/files?diff=unified&w=0#diff-9fa2bf7f1273283680da3daeb8bae8e20cf8d477f838575796d2bf3b415bc534L210-R210), [link](https://github.com/ultralytics/ultralytics/pull/4546/files?diff=unified&w=0#diff-9bd07d929cb8522458c879bd81a82d9af468ffbc215e762f39b2a50a601758d7R18), [link](https://github.com/ultralytics/ultralytics/pull/4546/files?diff=unified&w=0#diff-9bd07d929cb8522458c879bd81a82d9af468ffbc215e762f39b2a50a601758d7R30-R32))
  - Import `torchvision` and define constants for checking torchvision version in `torch_utils.py` ([link](https://github.com/ultralytics/ultralytics/pull/4546/files?diff=unified&w=0#diff-9bd07d929cb8522458c879bd81a82d9af468ffbc215e762f39b2a50a601758d7R18), [link](https://github.com/ultralytics/ultralytics/pull/4546/files?diff=unified&w=0#diff-9bd07d929cb8522458c879bd81a82d9af468ffbc215e762f39b2a50a601758d7R30-R32))
  - Rename `classify_transforms` to `classify_transforms_train` and add more parameters and functionality in `augment.py` ([link](https://github.com/ultralytics/ultralytics/pull/4546/files?diff=unified&w=0#diff-8945a175b39a1efcabd4d2fb09b583d2f25409f05038ca9a0f1165e39ee48528L798-R807), [link](https://github.com/ultralytics/ultralytics/pull/4546/files?diff=unified&w=0#diff-8945a175b39a1efcabd4d2fb09b583d2f25409f05038ca9a0f1165e39ee48528L815-R995))
  - Add `classify_transforms_eval` to create evaluation transforms in `augment.py` ([link](https://github.com/ultralytics/ultralytics/pull/4546/files?diff=unified&w=0#diff-8945a175b39a1efcabd4d2fb09b583d2f25409f05038ca9a0f1165e39ee48528R816-R978))
  - Import and use the new transform functions in `predictor.py` and `dataset.py` ([link](https://github.com/ultralytics/ultralytics/pull/4546/files?diff=unified&w=0#diff-9fa2bf7f1273283680da3daeb8bae8e20cf8d477f838575796d2bf3b415bc534L39-R39), [link](https://github.com/ultralytics/ultralytics/pull/4546/files?diff=unified&w=0#diff-9fa2bf7f1273283680da3daeb8bae8e20cf8d477f838575796d2bf3b415bc534L210-R210), [link](https://github.com/ultralytics/ultralytics/pull/4546/files?diff=unified&w=0#diff-0834dc62a5ebc93f43947ef49530b3e973013f844888646e5fb9962bfefecd90L11-R17), [link](https://github.com/ultralytics/ultralytics/pull/4546/files?diff=unified&w=0#diff-0834dc62a5ebc93f43947ef49530b3e973013f844888646e5fb9962bfefecd90L226-R246))
* Add and modify configuration options for classification augmentations in `default.yaml` ([link](https://github.com/ultralytics/ultralytics/pull/4546/files?diff=unified&w=0#diff-732d61cb4969eff986cec58058000b2949ea3b429a97ca5a1045897e754145c5R110-R112))
  - Pass the arguments from the `args` object to the transform functions in `dataset.py` ([link](https://github.com/ultralytics/ultralytics/pull/4546/files?diff=unified&w=0#diff-0834dc62a5ebc93f43947ef49530b3e973013f844888646e5fb9962bfefecd90L226-R246))
* Add and modify tests for classification transforms in `test_python.py` ([link](https://github.com/ultralytics/ultralytics/pull/4546/files?diff=unified&w=0#diff-51deb1ccb3a8b618784df2e2abfbc3b056436f1db72d1896d5102f9906978492R446-R487))
* Modify the `ClassificationDataset` and `ClassificationPredictor` classes to handle the new and legacy transforms ([link](https://github.com/ultralytics/ultralytics/pull/4546/files?diff=unified&w=0#diff-0834dc62a5ebc93f43947ef49530b3e973013f844888646e5fb9962bfefecd90L236-R253), [link](https://github.com/ultralytics/ultralytics/pull/4546/files?diff=unified&w=0#diff-0834dc62a5ebc93f43947ef49530b3e973013f844888646e5fb9962bfefecd90L251-R273), [link](https://github.com/ultralytics/ultralytics/pull/4546/files?diff=unified&w=0#diff-663b5aaaf7a04b1a7fe40326af039d2baf50df166e9738332eacbae9271ec55fL3-R5), [link](https://github.com/ultralytics/ultralytics/pull/4546/files?diff=unified&w=0#diff-663b5aaaf7a04b1a7fe40326af039d2baf50df166e9738332eacbae9271ec55fL31-R44))
  - Use `DEFAULT_MEAN` and `DEFAULT_STD` instead of hard-coded values in `dataset.py` ([link](https://github.com/ultralytics/ultralytics/pull/4546/files?diff=unified&w=0#diff-0834dc62a5ebc93f43947ef49530b3e973013f844888646e5fb9962bfefecd90L236-R253))
  - Convert the NumPy array to a PIL image before applying the torch transforms in `dataset.py` and `predict.py` ([link](https://github.com/ultralytics/ultralytics/pull/4546/files?diff=unified&w=0#diff-0834dc62a5ebc93f43947ef49530b3e973013f844888646e5fb9962bfefecd90L251-R273), [link](https://github.com/ultralytics/ultralytics/pull/4546/files?diff=unified&w=0#diff-663b5aaaf7a04b1a7fe40326af039d2baf50df166e9738332eacbae9271ec55fL3-R5), [link](https://github.com/ultralytics/ultralytics/pull/4546/files?diff=unified&w=0#diff-663b5aaaf7a04b1a7fe40326af039d2baf50df166e9738332eacbae9271ec55fL31-R44))
  - Check and store the legacy transform name in `predict.py` ([link](https://github.com/ultralytics/ultralytics/pull/4546/files?diff=unified&w=0#diff-663b5aaaf7a04b1a7fe40326af039d2baf50df166e9738332eacbae9271ec55fL31-R44))



## 🛠️ PR Summary

<sub>Made with ❤️ by [Ultralytics Actions](https://github.com/ultralytics/actions)<sub>

### 🌟 Summary
Enhanced classification model performance with improved augmentation settings.

### 📊 Key Changes
- 🔧 Updated accuracies for YOLOv8 classification models in documentation files.
- 🎨 Added new image augmentations with random erasing and auto augmentation policies like `randaugment`, `augmix`, and `autoaugment`.
- ✅ Improved and added new tests for classification augmentation transformations.
- 📈 Updated the default configuration to include new augmentation settings.

### 🎯 Purpose & Impact
- 🎯 The purpose is to improve model accuracies and provide more robust training through advanced augmentation techniques.
- 💥 This could lead to better-performing models and more variability in training, potentially improving the model's generalization capabilities. Users can now experiment with different augmentation settings tailored to their specific needs for more effective training customizations.

Problem statement:
Classification augmentations
## Summary
Execute research and testing on new/additional augmentations to use for Classification task.

## Subtasks
- [x] research new/additional augmentations for image classification
- [x] identified new augmentations to test with
- [x] execute new augmentations
- [x] collect and analyze new augmentation performance impacts
- [x] decide which augmentations to include
- [x] implement new augmentations into deployment code
- [ ] release new iteration of classification task with additional/new augmentations

Hints from issue discussion:

Commit message: add eval and no_aug train augmentations ofr cls

Patch to review:
diff --git a/ultralytics/cfg/default.yaml b/ultralytics/cfg/default.yaml
index cdeb95953d6..1d73d0280e6 100644
--- a/ultralytics/cfg/default.yaml
+++ b/ultralytics/cfg/default.yaml
@@ -107,6 +107,7 @@ fliplr: 0.5  # (float) image flip left-right (probability)
 mosaic: 1.0  # (float) image mosaic (probability)
 mixup: 0.0  # (float) image mixup (probability)
 copy_paste: 0.0  # (float) segment copy-paste (probability)
+auto_augment: randaugment # (str) auto augmentation policy for classification (randaugment, autoaugment, augmix)
 
 # Custom config.yaml ---------------------------------------------------------------------------------------------------
 cfg:  # (str, optional) for overriding defaults.yaml
diff --git a/ultralytics/data/augment.py b/ultralytics/data/augment.py
index 68160d4477a..8e88ab1a169 100644
--- a/ultralytics/data/augment.py
+++ b/ultralytics/data/augment.py
@@ -17,6 +17,10 @@
 
 from .utils import polygons2masks, polygons2masks_overlap
 
+DEFAULT_MEAN = (0.0, 0.0, 0.0)
+DEFAULT_STD = (1.0, 1.0, 1.0)
+DEFAULT_CROP_PERCENTAGE = 0.875
+
 
 # TODO: we might need a BaseTransform to make all these augments be compatible with both classification and semantic
 class BaseTransform:
@@ -792,15 +796,132 @@ def v8_transforms(dataset, imgsz, hyp, stretch=False):
         RandomFlip(direction='horizontal', p=hyp.fliplr, flip_idx=flip_idx)])  # transforms
 
 
-# Classification augmentations -----------------------------------------------------------------------------------------
-def classify_transforms(size=224, mean=(0.0, 0.0, 0.0), std=(1.0, 1.0, 1.0)):  # IMAGENET_MEAN, IMAGENET_STD
+# Classification augmentations eval -----------------------------------------------------------------------------------------
+def classify_transforms_eval(
+    size=224,
+    mean=DEFAULT_MEAN,
+    std=DEFAULT_STD,
+    interpolation: T.InterpolationMode = T.InterpolationMode.BILINEAR,
+    crop_percentage: float = DEFAULT_CROP_PERCENTAGE,
+):
+
+    if isinstance(size, (tuple, list)):
+        assert len(size) == 2
+        scale_size = tuple([math.floor(x / crop_percentage) for x in size])
+    else:
+        scale_size = math.floor(size / crop_percentage)
+        scale_size = (scale_size, scale_size)
+
+    # aspect ratio is preserved, crops center within image, no borders are added, image is lost
+    if scale_size[0] == scale_size[1]:
+        # simple case, use torchvision built-in Resize w/ shortest edge mode (scalar size arg)
+        tfl = [T.Resize(scale_size[0], interpolation=interpolation)]
+    else:
+        # resize shortest edge to matching target dim for non-square target
+        tfl = [T.Resize(scale_size)]
+    tfl += [T.CenterCrop(size)]
+
+    tfl += [T.ToTensor(), T.Normalize(
+        mean=torch.tensor(mean),
+        std=torch.tensor(std),
+    )]
+
+    return T.Compose(tfl)
+
+
+def classify_transforms_noaug_train(
+    img_size=224,
+    interpolation: T.InterpolationMode = T.InterpolationMode.BILINEAR,
+    mean=DEFAULT_MEAN,
+    std=DEFAULT_STD,
+):
+    tfl = [T.Resize(img_size, interpolation=interpolation), T.CenterCrop(img_size)]
+    tfl += [T.ToTensor(), T.Normalize(mean=torch.tensor(mean), std=torch.tensor(std))]
+    return T.Compose(tfl)
+
+
+# Classification augmentations train ---------------------------------------------------------------------------------------
+def classify_transforms_train(
+    size=224,
+    mean=DEFAULT_MEAN,
+    std=DEFAULT_STD,
+    scale=None,
+    ratio=None,
+    hflip=0.5,
+    vflip=0.0,
+    auto_augment=None,
+    color_jitter=0.4,
+    force_color_jitter=False,
+    re_prob=0.,
+    interpolation: T.InterpolationMode = T.InterpolationMode.BILINEAR,
+):
+    """
+    Classification augmentation with torchvision.
+    Inspired by timm/data/transforms_factory.py
+
+    Args:
+        size (int): image size
+        scale (tuple): scale range of the image. default is (0.08, 1.0)
+        ratio (tuple): aspect ratio range of the image. default is (3./4., 4./3.)
+        mean (tuple): mean values of RGB channels
+        std (tuple): std values of RGB channels
+        hflip (float): probability of horizontal flip
+        vflip (float): probability of vertical flip
+        auto_augment (str): auto augmentation policy. can be 'randaugment', 'augmix', 'autoaugment' or None.
+        color_jitter (float or tuple): color jitter factor or tuple of 3 values for brightness, contrast, saturation
+        force_color_jitter (bool): force to apply color jitter even if auto augment is enabled
+        re_prob (float): probability of random erasing
+        interpolation (T.InterpolationMode): interpolation mode. default is T.InterpolationMode.BILINEAR.
+
+    Returns:
+        T.Compose: torchvision transforms
+    """
     # Transforms to apply if albumentations not installed
     if not isinstance(size, int):
         raise TypeError(f'classify_transforms() size {size} must be integer, not (list, tuple)')
-    if any(mean) or any(std):
-        return T.Compose([CenterCrop(size), ToTensor(), T.Normalize(mean, std, inplace=True)])
-    else:
-        return T.Compose([CenterCrop(size), ToTensor()])
+    scale = tuple(scale or (0.08, 1.0))  # default imagenet scale range
+    ratio = tuple(ratio or (3. / 4., 4. / 3.))  # default imagenet ratio range
+    primary_tfl = [T.RandomResizedCrop(size, scale=scale, ratio=ratio, interpolation=interpolation)]
+    if hflip > 0.:
+        primary_tfl += [T.RandomHorizontalFlip(p=hflip)]
+    if vflip > 0.:
+        primary_tfl += [T.RandomVerticalFlip(p=vflip)]
+
+    secondary_tfl = []
+    disable_color_jitter = False
+    if auto_augment:
+        assert isinstance(auto_augment, str)
+        # color jitter is typically disabled if AA/RA on,
+        # this allows override without breaking old hparm cfgs
+        disable_color_jitter = not force_color_jitter
+
+        if auto_augment == 'randaugment':
+            secondary_tfl += [T.RandAugment(ainterpolation=interpolation)]
+        elif auto_augment == 'augmix':
+            secondary_tfl += [T.AugMix(interpolation=interpolation)]
+        elif auto_augment == 'autoaugment':
+            secondary_tfl += [T.AutoAugment(interpolation=interpolation)]
+        else:
+            raise ValueError(f'Invalid auto_augment policy: {auto_augment}. Should be one of "randaugment", '
+                             f'"augmix", "autoaugment" or None')
+
+    if color_jitter is not None and not disable_color_jitter:
+        # color jitter is enabled when not using AA or when forced
+        if isinstance(color_jitter, (list, tuple)):
+            # color jitter should be a 3-tuple/list if spec brightness/contrast/saturation
+            # or 4 if also augmenting hue
+            assert len(color_jitter) in (3, 4)
+        else:
+            # if it's a scalar, duplicate for brightness, contrast, and saturation, no hue
+            color_jitter = (float(color_jitter), ) * 3
+        secondary_tfl += [T.ColorJitter(*color_jitter)]
+
+    final_tfl = []
+    final_tfl += [T.ToTensor(), T.Normalize(mean=torch.tensor(mean), std=torch.tensor(std))]
+    if re_prob > 0.:
+        final_tfl.append(T.RandomErasing(p=re_prob, inplace=True))
+
+    return T.Compose(primary_tfl + secondary_tfl + final_tfl)
 
 
 def hsv2colorjitter(h, s, v):
@@ -809,17 +930,17 @@ def hsv2colorjitter(h, s, v):
 
 
 def classify_albumentations(
-        augment=True,
-        size=224,
-        scale=(0.08, 1.0),
-        hflip=0.5,
-        vflip=0.0,
-        hsv_h=0.015,  # image HSV-Hue augmentation (fraction)
-        hsv_s=0.7,  # image HSV-Saturation augmentation (fraction)
-        hsv_v=0.4,  # image HSV-Value augmentation (fraction)
-        mean=(0.0, 0.0, 0.0),  # IMAGENET_MEAN
-        std=(1.0, 1.0, 1.0),  # IMAGENET_STD
-        auto_aug=False,
+    augment=True,
+    size=224,
+    scale=(0.08, 1.0),
+    hflip=0.5,
+    vflip=0.0,
+    hsv_h=0.015,  # image HSV-Hue augmentation (fraction)
+    hsv_s=0.7,  # image HSV-Saturation augmentation (fraction)
+    hsv_v=0.4,  # image HSV-Value augmentation (fraction)
+    mean=DEFAULT_MEAN,  # IMAGENET_MEAN
+    std=DEFAULT_STD,  # IMAGENET_STD
+    auto_aug=False,
 ):
     """YOLOv8 classification Albumentations (optional, only used if package is installed)."""
     prefix = colorstr('albumentations: ')
diff --git a/ultralytics/data/dataset.py b/ultralytics/data/dataset.py
index cf3a7ce9dff..5ef1d445e4e 100644
--- a/ultralytics/data/dataset.py
+++ b/ultralytics/data/dataset.py
@@ -8,11 +8,13 @@
 import numpy as np
 import torch
 import torchvision
+from PIL import Image
 from tqdm import tqdm
 
 from ultralytics.utils import LOCAL_RANK, NUM_THREADS, TQDM_BAR_FORMAT, colorstr, is_dir_writeable
 
-from .augment import Compose, Format, Instances, LetterBox, classify_albumentations, classify_transforms, v8_transforms
+from .augment import (DEFAULT_MEAN, DEFAULT_STD, Compose, Format, Instances, LetterBox, classify_albumentations,
+                      classify_transforms_eval, classify_transforms_train, v8_transforms)
 from .base import BaseDataset
 from .utils import HELP_URL, LOGGER, get_hash, img2label_paths, verify_image, verify_image_label
 
@@ -225,7 +227,22 @@ def __init__(self, root, args, augment=False, cache=False, prefix=''):
         self.cache_disk = cache == 'disk'
         self.samples = self.verify_images()  # filter out bad images
         self.samples = [list(x) + [Path(x[0]).with_suffix('.npy'), None] for x in self.samples]  # file, index, npy, im
-        self.torch_transforms = classify_transforms(args.imgsz)
+        scale = (1.0 - args.scale, 1.0),  # (0.08, 1.0)
+        self.torch_transforms = classify_transforms_train(
+            size=args.imgsz,
+            scale=scale,
+            mean=DEFAULT_MEAN,
+            std=DEFAULT_STD,
+            hflip=args.fliplr,
+            vflip=args.flipud,
+            auto_augment=args.auto_augment) if augment else classify_transforms_eval(size=args.imgsz,
+                                                                                     scale=scale,
+                                                                                     mean=DEFAULT_MEAN,
+                                                                                     std=DEFAULT_STD,
+                                                                                     hflip=args.fliplr,
+                                                                                     vflip=args.flipud,
+                                                                                     auto_augment=args.auto_augment)
+
         self.album_transforms = classify_albumentations(
             augment=augment,
             size=args.imgsz,
@@ -235,8 +252,8 @@ def __init__(self, root, args, augment=False, cache=False, prefix=''):
             hsv_h=args.hsv_h,  # HSV-Hue augmentation (fraction)
             hsv_s=args.hsv_s,  # HSV-Saturation augmentation (fraction)
             hsv_v=args.hsv_v,  # HSV-Value augmentation (fraction)
-            mean=(0.0, 0.0, 0.0),  # IMAGENET_MEAN
-            std=(1.0, 1.0, 1.0),  # IMAGENET_STD
+            mean=DEFAULT_MEAN,  # IMAGENET_MEAN
+            std=DEFAULT_STD,  # IMAGENET_STD
             auto_aug=False) if augment else None
 
     def __getitem__(self, i):
@@ -250,10 +267,12 @@ def __getitem__(self, i):
             im = np.load(fn)
         else:  # read image
             im = cv2.imread(f)  # BGR
-        if self.album_transforms:
-            sample = self.album_transforms(image=cv2.cvtColor(im, cv2.COLOR_BGR2RGB))['image']
-        else:
-            sample = self.torch_transforms(im)
+        # if self.album_transforms:
+        #     sample = self.album_transforms(image=cv2.cvtColor(im, cv2.COLOR_BGR2RGB))['image']
+        # else:
+        # Convert NumPy array to PIL image
+        im = Image.fromarray(im)
+        sample = self.torch_transforms(im)
         return {'img': sample, 'cls': j}
 
     def __len__(self) -> int:
diff --git a/ultralytics/engine/predictor.py b/ultralytics/engine/predictor.py
index 77dfcdcfb16..eafc8e02c11 100644
--- a/ultralytics/engine/predictor.py
+++ b/ultralytics/engine/predictor.py
@@ -36,7 +36,7 @@
 
 from ultralytics.cfg import get_cfg
 from ultralytics.data import load_inference_source
-from ultralytics.data.augment import LetterBox, classify_transforms
+from ultralytics.data.augment import LetterBox, classify_transforms_eval
 from ultralytics.nn.autobackend import AutoBackend
 from ultralytics.utils import DEFAULT_CFG, LOGGER, MACOS, SETTINGS, WINDOWS, callbacks, colorstr, ops
 from ultralytics.utils.checks import check_imgsz, check_imshow
@@ -207,7 +207,7 @@ def predict_cli(self, source=None, model=None):
     def setup_source(self, source):
         """Sets up source and inference mode."""
         self.imgsz = check_imgsz(self.args.imgsz, stride=self.model.stride, min_dim=2)  # check image size
-        self.transforms = getattr(self.model.model, 'transforms', classify_transforms(
+        self.transforms = getattr(self.model.model, 'transforms', classify_transforms_eval(
             self.imgsz[0])) if self.args.task == 'classify' else None
         self.dataset = load_inference_source(source=source, imgsz=self.imgsz, vid_stride=self.args.vid_stride)
         self.source_type = self.dataset.source_type


Code review comments:
diff_hunk: @@ -225,7 +227,22 @@ def __init__(self, root, args, augment=False, cache=False, prefix=''):
         self.cache_disk = cache == 'disk'
         self.samples = self.verify_images()  # filter out bad images
         self.samples = [list(x) + [Path(x[0]).with_suffix('.npy'), None] for x in self.samples]  # file, index, npy, im
-        self.torch_transforms = classify_transforms(args.imgsz)
+        scale = (1.0 - args.scale, 1.0),  # (0.08, 1.0)
path: ultralytics/data/dataset.py
text: @author:
```suggestion
        scale = (1.0 - args.scale, 1.0)  # (0.08, 1.0)
```