text_plain_output_1.png
from PIL import Image

from imutils import paths

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

from torchvision.utils import make_grid

from tqdm import tqdm

import torch.nn as nn

import torch.nn.functional as F

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

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

batch_size = 50
input_size = 129 * 129
output_size = 12

train_dl = DataLoader(train_ds, batch_size, shuffle=True, num_workers=2, pin_memory=True)
val_dl = DataLoader(val_ds, batch_size * 2, num_workers=2, pin_memory=True)

for a, b in val_dl:
    break

def show_batch(dl):
    for images, labels in dl:
        fig, ax = plt.subplots(figsize=(16, 8))
        ax.set_xticks([]); ax.set_yticks([])
        data = images
        ax.imshow(make_grid(data, nrow=15).permute(1, 2, 0))
        break

def accuracy(outputs, labels):
    _, preds = torch.max(outputs, dim=1)
    return torch.tensor(torch.sum(preds == labels).item() / len(preds))

def F_score(output, label, threshold=0.5, beta=1):
    prob = output > threshold
    label = label > threshold
    TP = (prob & label).sum(1).float()
    TN = (~prob & ~label).sum(1).float()
    FP = (prob & ~label).sum(1).float()
    FN = (~prob & label).sum(1).float()
    precision = torch.mean(TP / (TP + FP + 1e-12))
    recall = torch.mean(TP / (TP + FN + 1e-12))
    F2 = (1 + beta ** 2) * precision * recall / (beta ** 2 * precision + recall + 1e-12)
    return F2.mean(0)

class ImageClassificationBase(nn.Module):
    def training_step(self, batch):
        images, labels = batch
        out = self(images)
        loss = F.binary_cross_entropy(out, labels)
        return loss
    def validation_step(self, batch):
        images, labels = batch
        out = self(images)
        loss = F.binary_cross_entropy(out, labels)
        acc = F_score(out, labels)
        return {'val_loss': loss.detach(), 'val_acc': acc}
    def validation_epoch_end(self, outputs):
        batch_losses = [x['val_loss'] for x in outputs]
        epoch_loss = torch.stack(batch_losses).mean()
        batch_accs = [x['val_acc'] for x in outputs]
        epoch_acc = torch.stack(batch_accs).mean()
        return {'val_loss': epoch_loss.item(), 'val_acc': epoch_acc.item()}
    def epoch_end(self, epoch, result):
        pass

@torch.no_grad()
def evaluate(model, val_loader):
    model.eval()
    outputs = [model.validation_step(batch) for batch in val_loader]
    return model.validation_epoch_end(outputs)
def fit(epochs, lr, model, train_loader, val_loader, opt_func=torch.optim.SGD):
    history = []
    optimizer = opt_func(model.parameters(), lr)
    for epoch in range(epochs):
        model.train()
        train_losses = []
        for batch in train_loader:
            loss = model.training_step(batch)
            train_losses.append(loss)
            loss.backward()
            optimizer.step()
            optimizer.zero_grad()
        result = evaluate(model, val_loader)
        result['train_loss'] = torch.stack(train_losses).mean().item()
        model.epoch_end(epoch, result)
        history.append(result)
    return history

def get_default_device():
    """Pick GPU if available, else CPU"""
    if torch.cuda.is_available():
        return torch.device('cuda')
    else:
        return torch.device('cpu')
def to_device(data, device):
    """Move tensor(s) to chosen device"""
    if isinstance(data, (list, tuple)):
        return [to_device(x, device) for x in data]
    return data.to(device, non_blocking=True)
class DeviceDataLoader:
    """Wrap a dataloader to move data to a device"""
    def __init__(self, dl, device):
        self.dl = dl
        self.device = device
    def __iter__(self):
        """Yield a batch of data after moving it to device"""
        for b in self.dl:
            yield to_device(b, self.device)
    def __len__(self):
        """Number of batches"""
        return len(self.dl)

device = get_default_device()
device