application_vnd.jupyter.stderr_output_1.png
from PIL import Image

from pathlib import Path

from pytorch_lightning.callbacks import ModelCheckpoint

from pytorch_lightning.callbacks.early_stopping import EarlyStopping

from sklearn.model_selection import train_test_split

from torch import nn

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

from torchmetrics.functional import accuracy, auroc

from torchmetrics.functional import mean_absolute_error, mean_squared_error

from torchvision import transforms

from tqdm import tqdm

import pandas as pd

import pytorch_lightning as pl

import torch

import torch.nn.functional as F

RANDOM_STATE = 42
NUM_CLASSES = 98
MAX_HOURS = 90
BATCH_SIZE = 128

pl.seed_everything(RANDOM_STATE)

class CountBoxes(Dataset):
    def __init__(self, df, transform=None):
        super().__init__()
        self.df = df
        self.transform = transform
        self.num_workers = 2
    def __len__(self):
        return self.df.shape[0]
    def __getitem__(self, index):
        img_id, label = (self.df.iloc[index]['image_id'], self.df.iloc[index]['count'])
        img_path = f'../input/count-the-blue-boxes/train/train/{label}/{label}_{img_id}.png'
        image = Image.open(img_path)
        image = image.convert('RGB')
        if self.transform is not None:
            image = self.transform(image)
        return (image, label)

class CountBoxesDataModule(pl.LightningDataModule):
    def __init__(self, path, batch_size=64, num_workers=4):
        super().__init__()
        self.path = Path(path)
        self.batch_size = batch_size
        self.num_workers = num_workers
        self.transform = transforms.Compose([transforms.Resize(size=(250, 500)), transforms.ToTensor()])
    def prepare_data(self):
        df = pd.read_csv(self.path / 'train.csv')
        df = df.loc[df['count'] < MAX_HOURS]
        X_train, X_valid, y_train, y_valid = train_test_split(df['image_id'], df['count'], stratify=df['count'], test_size=0.2)
        train_df = pd.concat([X_train, y_train], axis=1)
        valid_df = pd.concat([X_valid, y_valid], axis=1)
        train_df.to_pickle('train_df.pkl')
        valid_df.to_pickle('valid_df.pkl')
    def setup(self):
        self.train_df = pd.read_pickle('train_df.pkl')
        self.valid_df = pd.read_pickle('valid_df.pkl')
        self.train_transform = self.transform
        self.valid_transform = self.transform
    def train_dataloader(self):
        train_dataset = CountBoxes(df=self.train_df, transform=self.train_transform)
        return DataLoader(train_dataset, batch_size=self.batch_size, num_workers=self.num_workers, shuffle=True, pin_memory=True)
    def val_dataloader(self):
        valid_dataset = CountBoxes(df=self.valid_df, transform=self.valid_transform)
        return DataLoader(valid_dataset, batch_size=self.batch_size, num_workers=self.num_workers, shuffle=False, pin_memory=True)

datamodule = CountBoxesDataModule(path='../input/count-the-blue-boxes/', batch_size=BATCH_SIZE)
datamodule.prepare_data()
datamodule.setup()

train_dataloader = datamodule.train_dataloader()

valid_dataloader = datamodule.val_dataloader()

class CountModel(pl.LightningModule):
    def __init__(self, input_shape, num_classes, learning_rate=0.0002):
        super().__init__()
        self.save_hyperparameters()
        self.num_classes = num_classes
        self.learning_rate = learning_rate
        self.conv1 = nn.Conv2d(3, 16, 5, 1)
        self.conv2 = nn.Conv2d(16, 32, 5, 1)
        self.conv3 = nn.Conv2d(32, 64, 5, 1)
        self.conv4 = nn.Conv2d(64, 128, 5, 1)
        self.pool1 = torch.nn.MaxPool2d(2)
        self.pool2 = torch.nn.MaxPool2d(2)
        self.pool3 = torch.nn.MaxPool2d(2)
        n_sizes = self._get_conv_output(input_shape)
        self.fc1 = nn.Linear(n_sizes, 512)
        self.fc2 = nn.Linear(512, 256)
        self.fc3 = nn.Linear(256, self.num_classes)
    def _get_conv_output(self, shape):
        batch_size = 1
        input = torch.autograd.Variable(torch.rand(batch_size, *shape))
        output_feat = self._forward_features(input)
        n_size = output_feat.data.view(batch_size, -1).size(1)
        return n_size
    def _forward_features(self, x):
        x = F.relu(self.conv1(x))
        x = F.relu(self.pool1(x))
        x = F.relu(self.conv2(x))
        x = F.relu(self.pool2(x))
        x = F.relu(self.conv3(x))
        x = F.relu(self.pool3(x))
        x = F.relu(self.conv4(x))
        return x
    def forward(self, x):
        x = self._forward_features(x)
        x = x.view(x.size(0), -1)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = F.log_softmax(self.fc3(x), dim=1)
        return x
    def training_step(self, batch, batch_idx):
        x, y = batch
        logits = self(x)
        loss = F.nll_loss(logits, y)
        preds = torch.argmax(logits, dim=1)
        acc = accuracy(preds, y)
        mae = mean_absolute_error(preds, y)
        mse = mean_squared_error(preds, y)
        self.log('train_loss', loss, on_step=True, on_epoch=True, logger=True)
        self.log('train_acc', acc, on_step=True, on_epoch=True, logger=True)
        self.log('train_mae', mae, on_step=True, on_epoch=True, logger=True)
        self.log('train_mse', mse, on_step=True, on_epoch=True, logger=True)
        return loss
    def validation_step(self, batch, batch_idx):
        x, y = batch
        logits = self(x)
        loss = F.nll_loss(logits, y)
        preds = torch.argmax(logits, dim=1)
        acc = accuracy(preds, y)
        mae = mean_absolute_error(preds, y)
        mse = mean_squared_error(preds, y)
        self.log('val_loss', loss, on_step=True, on_epoch=True, logger=True)
        self.log('val_acc', acc, on_step=True, on_epoch=True, logger=True)
        self.log('val_mae', mae, on_step=True, on_epoch=True, logger=True)
        self.log('val_mse', mse, on_step=True, on_epoch=True, logger=True)
        return loss
    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), lr=self.learning_rate)
        return optimizer

checkpoint_callback = ModelCheckpoint(monitor='val_loss', filename='model-{epoch:02d}-{val_loss:.2f}', save_top_k=1, mode='min')
early_stop_callback = EarlyStopping(monitor='val_loss', patience=3, verbose=False, mode='min')

inference_model = CountModel.load_from_checkpoint(checkpoint_callback.best_model_path)

inference_model = inference_model.to('cuda')
inference_model.eval()
inference_model.freeze()

def get_predictions(dataloader):
    y_preds, y_actuals = ([], [])
    for imgs, labels in tqdm(dataloader):
        logits = inference_model(imgs.to('cuda'))
        preds = torch.argmax(logits, dim=1)
        y_actuals.extend(labels.numpy().tolist())
        y_preds.extend(preds.cpu().detach().numpy().tolist())
    return (y_preds, y_actuals)

valid_dataloader = datamodule.val_dataloader()
train_dataloader = datamodule.train_dataloader()

train_pred, train_act = get_predictions(train_dataloader)