image_output_1.png
from PIL import Image

from pathlib import Path

from torch import nn

from torch.utils.data import Dataset, DataLoader

from torchmetrics import Accuracy

from torchvision import transforms

import matplotlib.pyplot as plt

import os

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

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)

show_image(transformed_image)