text_plain_output_2.png application_vnd.jupyter.stderr_output_1.png
from PIL import Image

from pathlib import Path

from sklearn.model_selection import train_test_split

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

from torchvision import transforms

import matplotlib.pyplot as plt

import pandas as pd

import pytorch_lightning as pl

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

for data in train_dataloader:
    images, labels = data
    f, ax = plt.subplots(5, 5, figsize=(30, 30))
    for i in range(5 * 5):
        ax[i // 5, i % 5].imshow(images[i].permute(1, 2, 0))
        ax[i // 5, i % 5].axis('off')
        ax[i // 5, i % 5].set_title(labels[i], fontdict={'fontsize': 20})
    plt.tight_layout()
    plt.show()
    break