Open In Colab

In [1]:
%%bash
mkdir A B C
mkdir A/sub_1 A/sub_2 A/sub_3
mkdir B/sub_1 B/sub_2 B/sub_3
mkdir C/sub_1 C/sub_2 C/sub_3
In [2]:
import os

import numpy as np
import matplotlib.pyplot as plt

import torch
import torchvision
In [3]:
for folder in ["A", "B", "C"]:
    for sub_folder in os.listdir(folder):
        for i in range(2):
            img = np.random.random((20,20))
            plt.imsave(arr=img, fname=f"{folder}/{sub_folder}/img_{i}.png")
In [4]:
A_dataset = torchvision.datasets.ImageFolder(root = "A" , transform = torchvision.transforms.ToTensor())
B_dataset = torchvision.datasets.ImageFolder(root = "B" , transform = torchvision.transforms.ToTensor())
C_dataset = torchvision.datasets.ImageFolder(root = "C" , transform = torchvision.transforms.ToTensor())

all_datasets = []
all_datasets.append(A_dataset)
all_datasets.append(B_dataset)
all_datasets.append(C_dataset)

final_training_dataset = torch.utils.data.ConcatDataset(all_datasets)
In [5]:
for ind, c in enumerate(A_dataset.classes):
    A_dataset.classes[ind] = f"A_{c}"

for ind, c in enumerate(B_dataset.classes):
    B_dataset.classes[ind] = f"B_{c}"

for ind, c in enumerate(C_dataset.classes):
    C_dataset.classes[ind] = f"C_{c}"
In [6]:
A_dataset.classes, B_dataset.classes, C_dataset.classes, 
Out[6]:
(['A_sub_1', 'A_sub_2', 'A_sub_3'],
 ['B_sub_1', 'B_sub_2', 'B_sub_3'],
 ['C_sub_1', 'C_sub_2', 'C_sub_3'])
In [7]:
full_dl = torch.utils.data.DataLoader(final_training_dataset, batch_size = 1, shuffle = False)  
In [8]:
for idx, element in enumerate(full_dl):
    img, l = element
    if len(A_dataset) - idx >=0:
        print(A_dataset.classes[l])
    elif len(A_dataset)+len(B_dataset) - idx >=0:
        print(B_dataset.classes[l])
    else:
        print(C_dataset.classes[l])
A_sub_1
A_sub_1
A_sub_2
A_sub_2
A_sub_3
A_sub_3
A_sub_1
B_sub_1
B_sub_2
B_sub_2
B_sub_3
B_sub_3
B_sub_1
C_sub_1
C_sub_2
C_sub_2
C_sub_3
C_sub_3