mixgen.py

"""
MixGen: A New Multi-Modal Data Augmentation
https://arxiv.org/abs/2206.08358
Apache-2.0 License, Copyright 2022 Amazon
"""
import random
import numpy as np


def mixgen(image, text, num, lam=0.5):
    for i in range(num):
        # image mixup
        image[i,:] = (lam * image[i,:] + (1 - lam) * image[i+num,:]).unsqueeze(dim=0)
        # text concat
        text[i] = text[i] + " " + text[i+num]
    return image, text


def mixgen_batch(image, text, num, lam=0.5):
    batch_size = image.size()[0]
    index = np.random.permutation(batch_size)
    for i in range(batch_size):
        # image mixup
        image[i,:] = (lam * image[i,:] + (1 - lam) * image[index[i],:]).unsqueeze(dim=0)
        # text concat
        text[i] = text[i] + " " + text[index[i]]
    return image, text