text_plain_output_1.png
from PIL import Image

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)

#counting number of images under each category
plt.figure(figsize=(10,6))
g=sns.countplot(train_csv['Target'])
g.set_xticklabels(g.get_xticklabels(),rotation=40);

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()])

class heroDataset(Dataset):
    def __init__(self, csv_file, root_dir, transform=None):
        self.df = csv_file
        self.transform = transform
        self.root_dir = root_dir
    def __len__(self):
        return len(self.df)
    def __getitem__(self, idx):
        row = self.df.loc[idx]
        img_id, img_label = (row['File_name'], row['Target'])
        img = Image.open(row['image_path'])
        if self.transform:
            img = self.transform(img)
        return (img, encode_label(img_label))

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

def show_sample(img, target):
    pass

show_sample(*train_dataset[2000])