text_html_output_2.png text_html_output_1.png
from IPython.core.display import display, HTML

from PIL import Image

from io import BytesIO

from pathlib import Path

from pytorch_lightning.callbacks import ModelCheckpoint

from pytorch_lightning.loggers import WandbLogger

from pytorch_lightning.trainer import Trainer

from torch import nn

from torch.utils.data import Dataset, DataLoader

from torchmetrics import Accuracy

from torchvision import transforms

import base64

import matplotlib.pyplot as plt

import os

import pandas as pd

import pytorch_lightning as pl

import torch

import torchvision

import wandb

img_path = Path('../input/celeba-dataset/img_align_celeba/img_align_celeba')
csv_path = Path('../input/celeba-dataset/list_attr_celeba.csv')

class LitModel(pl.LightningModule):
    def __init__(self, n_classes, download_pretrained=True, wandb_logger=None, **kwargs):
        super().__init__()
        self.model = torchvision.models.resnet18(pretrained=True)
        self.model.fc = nn.Linear(512, n_classes)
        self.train_accuracy = Accuracy()
        self.val_accuracy = Accuracy()
        self.test_accuracy = Accuracy()
        self.loss = nn.CrossEntropyLoss()
        self.wandb_logger = wandb_logger
        self.opt_params = {'lr': 0.001}
    def forward(self, x):
        return self.model(x)
    def training_step(self, batch, batch_idx):
        x, y = batch
        y_hat_logits = self.forward(x)
        loss = self.loss(y_hat_logits, y)
        train_acc = self.train_accuracy(y_hat_logits, y)
        self.log('train_acc', train_acc, prog_bar=True)
        self.log('train_loss', loss)
        return {'loss': loss, 'progress_bar': {'train_acc': train_acc}}
    def training_epoch_end(self, outputs):
        avg_train_loss = torch.stack([x['loss'] for x in outputs]).mean()
        self.log('train_avg_loss', avg_train_loss, on_epoch=True)
        self.log('train_acc_epoch', self.train_accuracy, on_epoch=True)
    def validation_step(self, batch, batch_idx):
        x, y = batch
        y_hat_logits = self.forward(x)
        self.val_accuracy(y_hat_logits, y)
        val_loss = self.loss(y_hat_logits, y)
        return {'val_loss': val_loss, 'out_logits': y_hat_logits}
    def on_validation_batch_end(self, outputs, batch, batch_idx, dataloader_idx) -> None:
        if not self.wandb_logger:
            return
        if batch_idx == 0:
            outputs = outputs['out_logits']
            n = 20
            x, y = batch
            images = [img for img in x[:n]]
            assert len(outputs.shape) == 2
            assert outputs.shape[-1] == 2
            preds = outputs[:n].softmax(dim=-1)
            pred_prob, pred_class = preds.max(dim=-1)
            columns = ['image', 'pred_class', 'true_class', 'pred_prob']
            label_to_text = {0: 'Female', 1: 'Male'}
            td = []
            for image, y_pred_class, y_real, y_pred_prob in zip(x[:n], pred_class, y[:n], pred_prob):
                pred_class = label_to_text[y_pred_class.item()]
                true_class = label_to_text[y_real.item()]
                td.append([wandb.Image(image), pred_class, true_class, y_pred_prob])
            self.wandb_logger.log_table(key='samples', columns=columns, data=td)
    def validation_epoch_end(self, outputs):
        avg_val_loss = torch.stack([x['val_loss'] for x in outputs]).mean()
        self.log('val_acc_epoch', self.val_accuracy, on_epoch=True, prog_bar=True)
        self.log('val_loss_epoch', avg_val_loss, on_epoch=True, prog_bar=True)
        return {'val_loss': avg_val_loss}
    def test_step(self, batch, batch_idx):
        x, y = batch
        y_hat_logits = self.forward(x)
        self.test_accuracy(y_hat_logits, y)
        test_loss = self.loss(y_hat_logits, y)
        self.log('test_loss', test_loss, on_epoch=True)
        self.log('test_acc', self.test_accuracy, on_epoch=True)
    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), **self.opt_params)
        return optimizer

wandb_logger = WandbLogger(project='gender-detection-vit')

model = LitModel(2, wandb_logger=wandb_logger)

class CelebADataset(Dataset):
    def __init__(self, image_attr_ordered_map, base_image_dir: str, transforms):
        self.image_attr_ordered_map = image_attr_ordered_map
        self.base_image_dir = base_image_dir
        self.transforms = transforms
    def __len__(self):
        return len(self.image_attr_ordered_map)
    def __getitem__(self, index):
        img_name, labels = self.image_attr_ordered_map[index]
        image_path = os.path.join(self.base_image_dir, img_name)
        image = Image.open(image_path).convert('RGB')
        tensor_x = self.transforms(image)
        return (tensor_x, labels)

single_image_path = img_path / '000002.jpg'
orig_image = Image.open(single_image_path).convert('RGB')
orig_image = transforms.ToTensor()(orig_image)
orig_image.shape

def show_image(image_tensor):
    image_tensor = torch.permute(image_tensor.cpu(), (1, 2, 0))
    plt.axis('off')

image_transforms = {'train': transforms.Compose([transforms.RandomHorizontalFlip(), transforms.ToTensor()]), 'val': transforms.Compose([transforms.Resize(size=(218, 178)), transforms.ToTensor()])}

transformed_image = Image.open(single_image_path).convert('RGB')
transformed_image = image_transforms['train'](transformed_image)

class CelebADataModule(pl.LightningDataModule):
    def __init__(self, csv_path: str, image_dir: str, train_transforms, val_transforms, batch_size=16, **dataloader_kwargs):
        super(CelebADataModule, self).__init__()
        dataloader_kwargs.setdefault('num_workers', 2)
        self.csv_path = csv_path
        self.image_dir = image_dir
        self.transforms = {'train': train_transforms, 'val': val_transforms}
        self.bs = batch_size
        self.dataloader_kwargs = dataloader_kwargs
        self._label_to_node_idx = {-1: 0, 1: 1}
        self.image_labels = None
        self.train_dataset = None
        self.val_dataset = None
        self.test_dataset = None
    def setup(self, stage=None):
        attributes = pd.read_csv(self.csv_path)
        self.image_labels = []
        for i in range(len(attributes)):
            ith_sample = attributes.iloc[i]
            image_name = ith_sample[0]
            label = ith_sample['Male']
            label = self._label_to_node_idx[label]
            self.image_labels.append((image_name, label))
        total_d = len(self.image_labels)
        train_data = self.image_labels[:total_d - 5000]
        test_data = self.image_labels[total_d - 5000:total_d - 2500]
        val_data = self.image_labels[total_d - 2500:]
        self.train_dataset = CelebADataset(train_data, self.image_dir, self.transforms['train'])
        self.test_dataset = CelebADataset(test_data, self.image_dir, self.transforms['val'])
        self.val_dataset = CelebADataset(val_data, self.image_dir, self.transforms['val'])
    def train_dataloader(self):
        return DataLoader(self.train_dataset, batch_size=self.bs, shuffle=True, **self.dataloader_kwargs)
    def val_dataloader(self):
        return DataLoader(self.val_dataset, batch_size=self.bs, **self.dataloader_kwargs)
    def test_dataloader(self):
        return DataLoader(self.test_dataset, batch_size=self.bs, **self.dataloader_kwargs)

data_mod = CelebADataModule(csv_path=csv_path, image_dir=img_path, train_transforms=image_transforms['train'], val_transforms=image_transforms['val'], batch_size=32, num_workers=3)

checkpoint_callback = ModelCheckpoint(monitor='val_loss_epoch')

trainer = Trainer(gpus=1, max_epochs=6, logger=wandb_logger, callbacks=[checkpoint_callback])

trainer.fit(model, data_mod)

trainer.validate(model, data_mod)

best_model_path = trainer.checkpoint_callback.best_model_path
model = LitModel.load_from_checkpoint(best_model_path, n_classes=2)

from io import BytesIO
import base64
gender_target = {0: 'Female', 1: 'Male'}
def img_to_display(filename):
    i = Image.open(filename)
    i.thumbnail((200, 200), Image.LANCZOS)
    with BytesIO() as buffer:
        i.save(buffer, 'jpeg')
        return base64.b64encode(buffer.getvalue()).decode()
def display_result(filename, prediction, target):
    """
    Display the results in HTML
    """
    gender = 'Male'
    gender_icon = 'https://i.imgur.com/nxWan2u.png'
    if prediction[1] <= 0.5:
        gender_icon = 'https://i.imgur.com/oAAb8rd.png'
        gender = 'Female'
    display_html = '\n    <div style="overflow: auto;  border: 2px solid #D8D8D8;\n        padding: 5px; width: 420px;" >\n        <img src="data:image/jpeg;base64,{}" style="float: left;" width="200" height="200">\n        <div style="padding: 10px 0px 0px 20px; overflow: auto;">\n            <img src="{}" style="float: left;" width="40" height="40">\n            <h3 style="margin-left: 50px; margin-top: 2px;">{}</h3>\n            <p style="margin-left: 50px; margin-top: -6px; font-size: 12px">{} prob.</p>\n            <p style="margin-left: 50px; margin-top: -16px; font-size: 12px">Real Target: {}</p>\n            <p style="margin-left: 50px; margin-top: -16px; font-size: 12px">Filename: {}</p>\n        </div>\n    </div>\n    '.format(gender_icon, gender, '{0:.2f}%'.format(round(torch.max(prediction).item() * 100, 2)), gender_target[target.item()], filename)

def visualize_results(model, images, labels):
    model.eval()
    preds = model(images).softmax(dim=-1)
    for idx, (image, pred, label) in enumerate(zip(images, preds, labels)):
        test_vis_image_dir = Path('test-visualisation-images')
        if not test_vis_image_dir.exists():
            test_vis_image_dir.mkdir()
        filename = test_vis_image_dir / f'{idx}.jpg'
        torchvision.utils.save_image(image, filename)

data_mod.setup()
train_data = data_mod.test_dataloader()
img, label = next(iter(train_data))
img, label = (img[:8], label[:8])
visualize_results(model.cpu(), img, label)