def augment(images, labels, aug, iters):
    """Augments both image and label. Assumes input is a tensor with
       a batch dimension and values normalized to N(0,1)."""
    IMG_MEAN = np.array(
        (104.00698793, 116.66876762, 122.67891434), dtype=np.float32)

    # Transform label shape: B, C, W, H ==> B, W, H, C
    labels_are_3d = (len(labels.shape) == 4)
    if labels_are_3d:
        labels = labels.permute(0, 2, 3, 1)

    # Transform each image independently.
    aug_images, aug_labels = [], []
    for image, label in zip(images, labels):

        # Step 1: Undo normalization transformation, convert to numpy
        image = cv2.cvtColor(image.numpy().transpose(1, 2, 0) + IMG_MEAN,
                             cv2.COLOR_BGR2RGB).astype(np.uint8)
        label = label.numpy()
        label = label.astype('int64')

        # Step 2: Perform transformations on numpy images
        data = aug(image=image, mask=label)
        image, label = data['image'], data['mask']

        # Step 3: Convert back to PyTorch tensors
        image = paddle.to_tensor((cv2.cvtColor(
            image.astype(np.float32), cv2.COLOR_RGB2BGR) - IMG_MEAN).transpose(
                2, 0, 1))
        # label = np.where(label==-1, 255, label)
        label = paddle.to_tensor(label)
        if not labels_are_3d:
            label = label.astype('int64')

        # Add to list
        aug_images.append(image)
        aug_labels.append(label)

    # Stack
    images = paddle.stack(aug_images, axis=0)
    labels = paddle.stack(aug_labels, axis=0)

    # Transform label shape back: B, W, H, C ==> B, C, W, H
    if labels_are_3d:
        labels = labels.permute(0, 3, 1, 2)

    return images, labels
Augments both image and label. Assumes input is a tensor with
       a batch dimension and values normalized to N(0,1).
contrib/DomainAdaptation/utils/augmentation.py