Custom transformation on Subset of CIFAR10 - PyTorch

vision
1

I am trying to create a custom transformation to part of the CIFAR10 data set which superimposing of an image over the dataset. I was able to download the data and divide it into subsets. Using the following code:

transform_train = transforms.Compose([
    transforms.RandomCrop(32, padding=4),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
])

traindata = datasets.CIFAR10('./data', train=True, download=True,
                       transform= transform_train)

partitions = 5
traindata_split = torch.utils.data.random_split(traindata, [int(traindata.data.shape[0] / partitions) for _ in range(partitions)])

then I wanted to modify part of the splits so I created the following class and functions to use as as follows:

class MyDataset(Dataset): # https://discuss.pytorch.org/t/torch-utils-data-dataset-random-split/32209/3 
    def __init__(self, subset, transform=None):
        self.subset = subset
        self.transform = transform
        
    def __getitem__(self, index):
        x, y = self.subset[index]
        if self.transform:
            x = self.transform(x)
        return x, y
        
    def __len__(self):
        return len(self.subset)

and

class ImageSuperImpose(object):
    """  Image input as PIL and output as PIL
        To be used as part of  torchvision.transforms
       Args: p, a threshold value to control image thinning        
    """
    def __init__(self, p=0):
        self.p = p                  
        
    def __call__(self, image):
        img = cv2.imread('img.jpg') 
        img = img('float32')/255
        imgSm = cv2.resize(img,(32,32))
        np_arr = image.cpu().detach().numpy().T
        sample = cv2.addWeighted(np_arr, 1, imgSm, 1, 0)
        sample = sample.T
        t = torch.from_numpy(sample)
        return sample

transform_train2 = transforms.Compose([
    transforms.RandomCrop(32, padding=4),
    transforms.RandomHorizontalFlip(),
    ImagePoisoning(), 
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
])

datasetA = MyDataset(
    traindata_split[0], transform= transform_train2
)

test_loader = torch.utils.data.DataLoader(datasetA, batch_size=128, shuffle=True)

But when I tried to train the model on the subset I got the following error:

RuntimeError: The size of tensor a (32) must match the size of tensor b (3) at non-singleton dimension 0
2

I guess the shape mismatch is raised by the different memory formats used in OpenCV and PIL.
While OpenCV reads images as channels-last ([height, width, channels]), PIL uses channels-first ([channels, height, width]).
Adding these two images would try to add the spatial size of 32 to the channels of 3, which would raise this error. You would thus have to make sure both images are in the same format before adding them.