"""

An image-caption dataset dataloader.

Luke Melas-Kyriazi, 2021

"""

import warnings

from typing import Optional, Callable

from pathlib import Path

import numpy as np

import torch

import pandas as pd

from torch.utils.data import Dataset

from torchvision.datasets.folder import default_loader

from PIL import ImageFile

from PIL.Image import DecompressionBombWarning

ImageFile.LOAD_TRUNCATED_IMAGES = True

warnings.filterwarnings("ignore", category=UserWarning)

warnings.filterwarnings("ignore", category=DecompressionBombWarning)

class CaptionDataset(Dataset):

"""

A PyTorch Dataset class for (image, texts) tasks. Note that this dataset

returns the raw text rather than tokens. This is done on purpose, because

it's easy to tokenize a batch of text after loading it from this dataset.

"""

def __init__(self, *, images_root: str, captions_path: str, text_transform: Optional[Callable] = None,

image_transform: Optional[Callable] = None, image_transform_type: str = 'torchvision',

include_captions: bool = True):

"""

:param images_root: folder where images are stored

:param captions_path: path to csv that maps image filenames to captions

:param image_transform: image transform pipeline

:param text_transform: image transform pipeline

:param image_transform_type: image transform type, either `torchvision` or `albumentations`

:param include_captions: Returns a dictionary with `image`, `text` if `true`; otherwise returns just the images.

"""

# Base path for images

self.images_root = Path(images_root)

# Load captions as DataFrame

self.captions = pd.read_csv(captions_path, delimiter='\t', header=0)

self.captions['image_file'] = self.captions['image_file'].astype(str)

# PyTorch transformation pipeline for the image (normalizing, etc.)

self.text_transform = text_transform

self.image_transform = image_transform

self.image_transform_type = image_transform_type.lower()

assert self.image_transform_type in ['torchvision', 'albumentations']

# Total number of datapoints

self.size = len(self.captions)

# Return image+captions or just images

self.include_captions = include_captions

def verify_that_all_images_exist(self):

for image_file in self.captions['image_file']:

p = self.images_root / image_file

if not p.is_file():

print(f'file does not exist: {p}')

def _get_raw_image(self, i):

image_file = self.captions.iloc[i]['image_file']

image_path = self.images_root / image_file

image = default_loader(image_path)

return image

def _get_raw_text(self, i):

return self.captions.iloc[i]['caption']

def __getitem__(self, i):

image = self._get_raw_image(i)

caption = self._get_raw_text(i)

if self.image_transform is not None:

if self.image_transform_type == 'torchvision':

image = self.image_transform(image)

elif self.image_transform_type == 'albumentations':

image = self.image_transform(image=np.array(image))['image']

else:

raise NotImplementedError(f"{self.image_transform_type=}")

return {'image': image, 'text': caption} if self.include_captions else image

def __len__(self):

return self.size

if __name__ == "__main__":

import albumentations as A

from albumentations.pytorch import ToTensorV2

from transformers import AutoTokenizer

# Paths

images_root = './images'

captions_path = './images-list-clean.tsv'

# Create transforms

tokenizer = AutoTokenizer.from_pretrained('distilroberta-base')

def tokenize(text):

return tokenizer(text, max_length=32, truncation=True, return_tensors='pt', padding='max_length')

image_transform = A.Compose([

A.Resize(256, 256), A.CenterCrop(256, 256),

A.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]), ToTensorV2()])

# Create dataset

dataset = CaptionDataset(

images_root=images_root,

captions_path=captions_path,

image_transform=image_transform,

text_transform=tokenize,

image_transform_type='albumentations')

# Create dataloader

dataloader = torch.utils.data.DataLoader(dataset, batch_size=2)

batch = next(iter(dataloader))

print({k: (v.shape if isinstance(v, torch.Tensor) else v) for k, v in batch.items()})

# # (Optional) Check that all the images exist

# dataset = CaptionDataset(images_root=images_root, captions_path=captions_path)

# dataset.verify_that_all_images_exist()

# print('Done')