Image augmentation is a bit tricky, tied to models, and hard to do for training data only

Checklist / 检查清单

• I have searched existing issues, and this is a new question or discussion topic. / 我已经搜索过现有的 issues，确认这是一个新的问题与讨论。

Question Description / 问题描述

I wrote an image augmentation plugin (code below). It does what I want, but there are a few things I find a bit tricky here, and I wonder if they could be done better:

1. I only want to augment the training dataset, not the validation one. I think the template doesn't get the information of which dataset is being used. I solve this by adding an extra field _augment_image to my dataset, and then I also need to pass --remove_unused_columns false on the command line. Of course this also means I couldn't use it with non-JSONL datasets.
2. I think the current approach of overriding a template couples more things than it ideally should. In my plugin, I define augmentation for Qwen 3.5. But, it seems to me that image augmentation should have nothing to do with the model used. As I understand it, currently there's very little ways to write an image augmentation plugin that works with any multimodal model.

Here's what I use; I have a separate Python package swiftwrap installed in the same venv as swift:

swiftwrap/_image_augmentation.py:

import numpy as np
import swift.utils
from PIL import Image as PILImage
from swift.template import StdTemplateInputs, Template
from swift.template.templates.qwen import Qwen3_5Template
from torchvision.transforms import v2 as tv_v2

# The logic:
# - If the image has more pixels than _MIN_PIXELS, scale it randomly between _MIN_PIXELS and _MAX_PIXELS
# - If the image has fewer pixels than _MIN_PIXELS, scale it anywhere between its original size and _MAX_PIXELS
#
# That is, _MIN_PIXELS is not our absolute minimum pixels; it's the threshold below which we do not scale down
# images larger than _MIN_PIXELS.

DEBUG = False

_MAX_PIXELS = 2880**2
_MIN_PIXELS = 800**2

logger = swift.utils.get_logger()


def _augment_image(image: PILImage.Image) -> PILImage.Image:
    im = tv_v2.functional.to_image(image)
    channels, width, height = im.size()

    if DEBUG:
        logger.info(
            "Got image of size: %dx%d, %d channels; min=%f, max=%f",
            width,
            height,
            channels,
            im.min().item(),
            im.max().item(),
        )
    transforms = tv_v2.Compose(
        [
            tv_v2.RandomRotation(
                degrees=15,
                expand=True,
                fill=im.max().item(),
                interpolation=tv_v2.InterpolationMode.BILINEAR,
            ),
            tv_v2.ColorJitter(brightness=0.5, contrast=0.5, saturation=0.5, hue=0.5),
        ]
    )
    im = transforms(im)

    # torchvision Resize supports maximum edge sizes but not maximum pixel counts, so no Compose
    num_pixels = width * height
    # compute max and min scale factors based on pixel constraints
    max_scale = (_MAX_PIXELS / num_pixels) ** 0.5 if num_pixels > _MAX_PIXELS else 1.0
    # We allow original size even if it's below _MIN_PIXELS
    min_scale = 1.0 if num_pixels <= _MIN_PIXELS else (_MIN_PIXELS / num_pixels) ** 0.5
    scale = np.random.uniform(min_scale, max_scale)
    im = tv_v2.functional.resize(
        im,
        [int(height * scale), int(width * scale)],
        interpolation=tv_v2.InterpolationMode.BILINEAR,
    )
    image = tv_v2.functional.to_pil_image(im)
    if DEBUG:
        logger.info(
            "Returning augmented image of size: %dx%d", image.width, image.height
        )
    return image


class _ImageAugmentationMixin(Template):
    def _preprocess_inputs(self, inputs: StdTemplateInputs) -> None:
        super()._preprocess_inputs(inputs)

        augment = inputs.extra_kwargs.get("_augment_image")
        if not isinstance(augment, bool):
            raise TypeError(
                f"Every sample must contain boolean '_augment_image'; got {augment!r}. "
                " Did you forget to run with --remove_unused_columns false?"
            )

        if not augment:
            return

        for i, image in enumerate(inputs.images):
            if not isinstance(image, PILImage.Image):
                raise TypeError(
                    f"Expected preprocessing to produce PIL.Image.Image, "
                    f"got {type(image)!r}"
                )
            inputs.images[i] = _augment_image(image)


class AugmentedQwen35Template(
    _ImageAugmentationMixin,
    Qwen3_5Template,
):
    pass

image_augmentation_plugin.py:

from copy import deepcopy

from swift.template import TEMPLATE_MAPPING, register_template

from swiftwrap._image_augmentation import AugmentedQwen35Template

meta = deepcopy(TEMPLATE_MAPPING["qwen3_5"])
meta.template_type = "qwen3_5_augmented"
meta.template_cls = AugmentedQwen35Template

register_template(meta)

Then I add "augment_image": true to all rows in my training dataset and "augment_image": false to all rows in my validation dataset and run swift with --remove_unused_columns false --external_plugins image_augmentation_plugin.py.

(I also wish --external_plugins would take Python module names, but that's not a major hurdle here!)

Start by tracing the template preprocessing and dataset-column handling used by the external plugin, especially _ImageAugmentationMixin._preprocess_inputs and the _augment_image field. Compare that flow for training and validation data and for different multimodal templates. Done should provide model-independent image augmentation, apply it only to training samples, and avoid requiring JSONL-specific fields or --remove_unused_columns false.