Albumentations in classification on Yolov11

[taitsmp]

Could someone give me very short example of how to configure specific Albumentations transforms when doing image classification using YoloV11 or

My understanding is that classification requires having the data argument to model.train be a directory not a yaml file. I'm confused on how to change the configuration of Albumentations in this case. I want to increase the percentage of black and white images used when training.

Thank you!

[UltralyticsAssistant]

👋 Hello @taitsmp, thank you for your interest in Ultralytics 🚀! We recommend a visit to the Docs for new users, where you can find many Python and CLI usage examples and where many of the most common questions may already be answered.

If this is a 🐛 Bug Report, please provide a minimum reproducible example to help us debug it.

If this is a custom training ❓ Question, please provide as much information as possible, including dataset samples, your setup details, and specific configurations. You can also verify you are following our Tips for Best Training Results for optimal usage.

Regarding Albumentations specifically, customization typically involves configuring transformations either in your training script or through the dataset pipeline. To assist you further, could you share any portion of your current workflow or code where you're incorporating Albumentations? This will help ensure a more targeted response from the community or team 😃.

Join the Ultralytics community where it suits you best. For real-time chat, head to Discord 🎧. Prefer in-depth discussions? Check out Discourse. Or dive into threads on our Subreddit to share knowledge with others.

Upgrade

Upgrade to the latest ultralytics package including all requirements in a Python>=3.8 environment with PyTorch>=1.8 to verify your issue is not already resolved in the latest version:

pip install -U ultralytics

Environments

YOLO may be run in any of the following up-to-date verified environments (with all dependencies including CUDA/CUDNN, Python and PyTorch preinstalled):

• Notebooks with free GPU:
• Google Cloud Deep Learning VM. See GCP Quickstart Guide
• Amazon Deep Learning AMI. See AWS Quickstart Guide
• Docker Image. See Docker Quickstart Guide

Status

If this badge is green, all Ultralytics CI tests are currently passing. CI tests verify correct operation of all YOLO Modes and Tasks on macOS, Windows, and Ubuntu every 24 hours and on every commit.

This is an automated response to help guide you as quickly as possible. An Ultralytics engineer will also assist you soon 😊.

[glenn-jocher]

@taitsmp you're correct that for classification, the data argument points to a directory. To customize Albumentations, you can configure the augmenter object directly in your training script. For example, to increase the probability of black-and-white transformations, you can customize Albumentations like this:

from ultralytics import YOLO
import albumentations as A

# Define custom Albumentations transforms
custom_augs = A.Compose([
    A.ToGray(p=0.5),  # 50% chance to convert images to grayscale
])

# Load model
model = YOLO("yolo11n-cls.pt") 

# Train model with custom augmentations applied
results = model.train(data="path/to/your/dataset", epochs=100, imgsz=640, augment=custom_augs)

This allows you to adjust augmentation probabilities as needed. Let me know if further clarification is required!

[taitsmp]

This didn't work for me.

 model = YOLO("yolo11m-cls.pt")
 data = os.path.abspath(base_folder)


 # Train the model
 results = model.train(
        data=data,
        epochs=training_epochs,
        imgsz=model_input_size,
        augment=A.Compose([
    A.ToGray(p=0.5),  
]),        
    )

'augment' must be a bool (i.e. 'augment=True' or 'augment=False')

[gmallios]

After going through the code, it seems that passing custom Albumentation augs is not supported currently 😕. I will try to open a PR but im not very familiar with the codebase so that may take some time.

[glenn-jocher]

Thank you for pointing that out! Indeed, custom Albumentations augmentations are not currently supported directly. We welcome contributions, and if you encounter any issues or need guidance while working on the PR, feel free to ask. The team and the community are here to help!