text_plain_output_1.png image_output_1.png
from imutils import paths

from torch.utils.data import Dataset, random_split, DataLoader

from tqdm import tqdm

import torchvision.transforms as transforms

data_dir = '../input/super-hero/Q4-superheroes_image_data/'
train_dir = data_dir + 'CAX_Superhero_Train'
test_dir = data_dir + 'CAX_Superhero_Test'

label_enc = {'Ant-Man': 0, 'Aquaman': 1, 'Avengers': 2, 'Batman': 3, 'Black Panther': 4, 'Captain America': 5, 'Catwoman': 6, 'Ghost Rider': 7, 'Hulk': 8, 'Iron Man': 9, 'Spiderman': 10, 'Superman': 11}
label_deco = {0: 'Ant-Man', 1: 'Aquaman', 2: 'Avengers', 3: 'Batman', 4: 'Black Panther', 5: 'Captain America', 6: 'Catwoman', 7: 'Ghost Rider', 8: 'Hulk', 9: 'Iron Man', 10: 'Spiderman', 11: 'Superman'}

def create_img_df(dir):
    img_list = list(paths.list_images(dir))
    data = pd.DataFrame(columns=['File_name', 'Target'])
    for i, ipaths in tqdm(enumerate(img_list), total=len(img_list)):
        data.loc[i, 'image_path'] = ipaths
        data.loc[i, 'File_name'] = os.path.basename(ipaths)
        data.loc[i, 'Target'] = os.path.split(os.path.dirname(ipaths))[-1]
    return data

train_csv = create_img_df(train_dir)

def encode_label(label):
    target = torch.zeros(12, dtype=torch.float)
    target[int(label)] = 1.0
    return target
def decode_target(target, text_labels=False, threshold=0.5):
    label = None
    for i, x in enumerate(target):
        if x >= threshold:
            label = i
            break
    if text_labels:
        return f'{label_deco[label]}({label})'
    return label

transformer = transforms.Compose([transforms.Resize(130), transforms.CenterCrop(129), transforms.ToTensor()])

train_dataset = heroDataset(train_csv, train_dir, transform=transformer)

torch.manual_seed(10)
val_pct = 0.11
val_size = int(val_pct * len(train_dataset))
train_size = len(train_dataset) - val_size

train_ds, val_ds = random_split(train_dataset, [train_size, val_size])
(len(train_ds), len(val_ds))