from torchvision.datasets import OxfordIIITPet
from torchvision.transforms.v2 import RandAugment
from torchvision.transforms.functional import InterpolationMode

origin_data = OxfordIIITPet(
    root="data",
    transform=None
)

no0_data = OxfordIIITPet( # `no` is num_ops.
    root="data",
    transform=RandAugment(num_ops=0)
)

no1_data = OxfordIIITPet(
    root="data",
    transform=RandAugment(num_ops=1)
)

no2_data = OxfordIIITPet(
    root="data",
    transform=RandAugment(num_ops=2)
)

no5_data = OxfordIIITPet(
    root="data",
    transform=RandAugment(num_ops=5)
)

no10_data = OxfordIIITPet(
    root="data",
    transform=RandAugment(num_ops=10)
)

no25_data = OxfordIIITPet(
    root="data",
    transform=RandAugment(num_ops=25)
)

no50_data = OxfordIIITPet(
    root="data",
    transform=RandAugment(num_ops=50)
)

no100_data = OxfordIIITPet(
    root="data",
    transform=RandAugment(num_ops=100)
)

no500_data = OxfordIIITPet(
    root="data",
    transform=RandAugment(num_ops=500)
)

no1000_data = OxfordIIITPet(
    root="data",
    transform=RandAugment(num_ops=1000)
)

no0m30_data = OxfordIIITPet( # `m` is magnitude.
    root="data",
    transform=RandAugment(num_ops=0, magnitude=30)
)

no1m30_data = OxfordIIITPet(
    root="data",
    transform=RandAugment(num_ops=1, magnitude=30)
)

no2m30_data = OxfordIIITPet(
    root="data",
    transform=RandAugment(num_ops=2, magnitude=30)
)

no5m30_data = OxfordIIITPet(
    root="data",
    transform=RandAugment(num_ops=5, magnitude=30)
)

no10m30_data = OxfordIIITPet(
    root="data",
    transform=RandAugment(num_ops=10, magnitude=30)
)

no25m30_data = OxfordIIITPet(
    root="data",
    transform=RandAugment(num_ops=25, magnitude=30)
)

no0nmb1000_data = OxfordIIITPet( # `nmb` is num_magnitude_bins.
    root="data",
    transform=RandAugment(num_ops=0, num_magnitude_bins=1000)
)

no1nmb1000_data = OxfordIIITPet(
    root="data",
    transform=RandAugment(num_ops=1, num_magnitude_bins=1000)
)

no2nmb1000_data = OxfordIIITPet(
    root="data",
    transform=RandAugment(num_ops=2, num_magnitude_bins=1000)
)

no5nmb1000_data = OxfordIIITPet(
    root="data",
    transform=RandAugment(num_ops=5, num_magnitude_bins=1000)
)

no10nmb1000_data = OxfordIIITPet(
    root="data",
    transform=RandAugment(num_ops=10, num_magnitude_bins=1000)
)

no25nmb1000_data = OxfordIIITPet(
    root="data",
    transform=RandAugment(num_ops=25, num_magnitude_bins=1000)
)

no50nmb1000_data = OxfordIIITPet(
    root="data",
    transform=RandAugment(num_ops=50, num_magnitude_bins=1000)
)

no100nmb1000_data = OxfordIIITPet(
    root="data",
    transform=RandAugment(num_ops=100, num_magnitude_bins=1000)
)

no500nmb1000_data = OxfordIIITPet(
    root="data",
    transform=RandAugment(num_ops=500, num_magnitude_bins=1000)
)

no1000nmb1000_data = OxfordIIITPet(
    root="data",
    transform=RandAugment(num_ops=1000, num_magnitude_bins=1000)
)

no0m999nmb1000_data = OxfordIIITPet(
    root="data",
    transform=RandAugment(num_ops=0, magnitude=999, num_magnitude_bins=1000)
)

no1m999nmb1000_data = OxfordIIITPet(
    root="data",
    transform=RandAugment(num_ops=1, magnitude=999, num_magnitude_bins=1000)
)

no2m999nmb1000_data = OxfordIIITPet(
    root="data",
    transform=RandAugment(num_ops=2, magnitude=999, num_magnitude_bins=1000)
)

no5m999nmb1000_data = OxfordIIITPet(
    root="data",
    transform=RandAugment(num_ops=5, magnitude=999, num_magnitude_bins=1000)
)

no10m999nmb1000_data = OxfordIIITPet(
    root="data",
    transform=RandAugment(num_ops=10, magnitude=999, num_magnitude_bins=1000)
)

no25m999nmb1000_data = OxfordIIITPet(
    root="data",
    transform=RandAugment(num_ops=25, magnitude=999, num_magnitude_bins=1000)
)

no50m999nmb1000_data = OxfordIIITPet(
    root="data",
    transform=RandAugment(num_ops=50, magnitude=999, num_magnitude_bins=1000)
)

no100m999nmb1000_data = OxfordIIITPet(
    root="data",
    transform=RandAugment(num_ops=100, magnitude=999, num_magnitude_bins=1000)
)

no500m999nmb1000_data = OxfordIIITPet(
    root="data",
    transform=RandAugment(num_ops=500, magnitude=999, num_magnitude_bins=1000)
)

no1000m999nmb1000_data = OxfordIIITPet(
    root="data",
    transform=RandAugment(num_ops=1000, magnitude=999, num_magnitude_bins=1000)
)

no25fgray_data = OxfordIIITPet( # `f` is fill.
    root="data",
    transform=RandAugment(num_ops=25, fill=150)
    # transform=RandAugment(num_ops=25, fill=[150])
)

no25fpurple_data = OxfordIIITPet(
    root="data",
    transform=RandAugment(num_ops=25, fill=[160, 32, 240])
)

import matplotlib.pyplot as plt

def show_images1(data, main_title=None):
    plt.figure(figsize=[10, 5])
    plt.suptitle(t=main_title, y=0.8, fontsize=14)
    for i, (im, _) in zip(range(1, 6), data):
        plt.subplot(1, 5, i)
        plt.imshow(X=im)
        plt.xticks(ticks=[])
        plt.yticks(ticks=[])
    plt.tight_layout()
    plt.show()

show_images1(data=origin_data, main_title="origin_data")
print()
show_images1(data=no0_data, main_title="no0_data")
show_images1(data=no1_data, main_title="no1_data")
show_images1(data=no2_data, main_title="no2_data")
show_images1(data=no5_data, main_title="no5_data")
show_images1(data=no10_data, main_title="no10_data")
show_images1(data=no25_data, main_title="no25_data")
show_images1(data=no50_data, main_title="no50_data")
show_images1(data=no100_data, main_title="no100_data")
show_images1(data=no500_data, main_title="no500_data")
show_images1(data=no1000_data, main_title="no1000_data")
print()
show_images1(data=no0m30_data, main_title="no0m30_data")
show_images1(data=no1m30_data, main_title="no1m30_data")
show_images1(data=no2m30_data, main_title="no2m30_data")
show_images1(data=no5m30_data, main_title="no5m30_data")
show_images1(data=no10m30_data, main_title="no10m30_data")
show_images1(data=no25m30_data, main_title="no25m30_data")
print()
show_images1(data=no0nmb1000_data, main_title="no0nmb1000_data")
show_images1(data=no1nmb1000_data, main_title="no1nmb1000_data")
show_images1(data=no2nmb1000_data, main_title="no2nmb1000_data")
show_images1(data=no5nmb1000_data, main_title="no5nmb1000_data")
show_images1(data=no10nmb1000_data, main_title="no10nmb1000_data")
show_images1(data=no25nmb1000_data, main_title="no25nmb1000_data")
show_images1(data=no50nmb1000_data, main_title="no50nmb1000_data")
show_images1(data=no100nmb1000_data, main_title="no100nmb1000_data")
show_images1(data=no500nmb1000_data, main_title="no500nmb1000_data")
show_images1(data=no1000nmb1000_data, main_title="no1000nmb1000_data")
print()
show_images1(data=no0m999nmb1000_data, main_title="no0m999nmb1000_data")
show_images1(data=no1m999nmb1000_data, main_title="no1m999nmb1000_data")
show_images1(data=no2m999nmb1000_data, main_title="no2m999nmb1000_data")
show_images1(data=no5m999nmb1000_data, main_title="no5m999nmb1000_data")
show_images1(data=no10m999nmb1000_data, main_title="no10m999nmb1000_data")
show_images1(data=no25m999nmb1000_data, main_title="no25m999nmb1000_data")
show_images1(data=no50m999nmb1000_data, main_title="no50m999nmb1000_data")
show_images1(data=no100m999nmb1000_data, main_title="no100m999nmb1000_data")
show_images1(data=no500m999nmb1000_data, main_title="no500m999nmb1000_data")
show_images1(data=no1000m999nmb1000_data, main_title="no1000m999nmb1000_data")
print()
show_images1(data=no25fgray_data, main_title="no25fgray_data")
show_images1(data=no25fpurple_data, main_title="no25fpurple_data")

# ↓ ↓ ↓ ↓ ↓ ↓ The code below is identical to the code above. ↓ ↓ ↓ ↓ ↓ ↓
def show_images2(data, main_title=None, no=2, m=9, nmb=31,
                 ip=InterpolationMode.NEAREST, f=None):
    plt.figure(figsize=[10, 5])
    plt.suptitle(t=main_title, y=0.8, fontsize=14)
    if main_title != "origin_data":
        for i, (im, _) in zip(range(1, 6), data):
            plt.subplot(1, 5, i)
            ra = RandAugment(num_ops=no, magnitude=m,
                             num_magnitude_bins=nmb,
                             interpolation=ip, fill=f)
            plt.imshow(X=ra(im))
            plt.xticks(ticks=[])
            plt.yticks(ticks=[])
    else:
        for i, (im, _) in zip(range(1, 6), data):
            plt.subplot(1, 5, i)
            plt.imshow(X=im)
            plt.xticks(ticks=[])
            plt.yticks(ticks=[])
    plt.tight_layout()
    plt.show()

show_images2(data=origin_data, main_title="origin_data")
print()
show_images2(data=origin_data, main_title="no0_data", no=0)
show_images2(data=origin_data, main_title="no1_data", no=1)
show_images2(data=origin_data, main_title="no2_data", no=2)
show_images2(data=origin_data, main_title="no5_data", no=5)
show_images2(data=origin_data, main_title="no10_data", no=10)
show_images2(data=origin_data, main_title="no25_data", no=25)
show_images2(data=origin_data, main_title="no50_data", no=50)
show_images2(data=origin_data, main_title="no100_data", no=100)
show_images2(data=origin_data, main_title="no500_data", no=500)
show_images2(data=origin_data, main_title="no1000_data", no=1000)
print()
show_images2(data=origin_data, main_title="no0m30_data", no=0, m=30)
show_images2(data=origin_data, main_title="no1m30_data", no=1, m=30)
show_images2(data=origin_data, main_title="no2m30_data", no=2, m=30)
show_images2(data=origin_data, main_title="no5m30_data", no=5, m=30)
show_images2(data=origin_data, main_title="no10m30_data", no=10, m=30)
show_images2(data=origin_data, main_title="no25m30_data", no=25, m=30)
print()
show_images2(data=origin_data, main_title="no0nmb1000_data", no=0, nmb=1000)
show_images2(data=origin_data, main_title="no1nmb1000_data", no=1, nmb=1000)
show_images2(data=origin_data, main_title="no2nmb1000_data", no=2, nmb=1000)
show_images2(data=origin_data, main_title="no5nmb1000_data", no=5, nmb=1000)
show_images2(data=origin_data, main_title="no10nmb1000_data", no=10, nmb=1000)
show_images2(data=origin_data, main_title="no25nmb1000_data", no=25, nmb=1000)
show_images2(data=origin_data, main_title="no50nmb1000_data", no=50, nmb=1000)
show_images2(data=origin_data, main_title="no100nmb1000_data", no=100, 
             nmb=1000)
show_images2(data=origin_data, main_title="no500nmb1000_data", no=500, 
             nmb=1000)
show_images2(data=origin_data, main_title="no1000nmb1000_data", no=1000, 
             nmb=1000)
print()
show_images2(data=origin_data, main_title="no0m999nmb1000_data", no=0, m=999,
             nmb=1000)
show_images2(data=origin_data, main_title="no1m999nmb1000_data", no=1, m=999,
             nmb=1000)
show_images2(data=origin_data, main_title="no2m999nmb1000_data", no=2, m=999,
             nmb=1000)
show_images2(data=origin_data, main_title="no5m999nmb1000_data", no=5, m=999,
             nmb=1000)
show_images2(data=origin_data, main_title="no10m999nmb1000_data", no=10, m=999,
             nmb=1000)
show_images2(data=origin_data, main_title="no25m999nmb1000_data", no=25, m=999,
             nmb=1000)
show_images2(data=origin_data, main_title="no50m999nmb1000_data", no=50, m=999,
             nmb=1000)
show_images2(data=origin_data, main_title="no100m999nmb1000_data", no=100, 
             m=999, nmb=1000)
show_images2(data=origin_data, main_title="no500m999nmb1000_data", no=500, 
             m=999, nmb=1000)
show_images2(data=origin_data, main_title="no1000m999nmb1000_data", no=1000,
             m=999, nmb=1000)
print()
show_images2(data=origin_data, main_title="no25fgray_data", no=25, f=150)
show_images2(data=origin_data, main_title="no25fpurple_data", no=25,
             f=[160, 32, 240])