135 lines
3.8 KiB
Python
135 lines
3.8 KiB
Python
"""Definition of Visual Critial Discourse Analysis dataloader."""
|
|
|
|
import random
|
|
|
|
from PIL import Image
|
|
from torch import Tensor
|
|
from torch.utils.data import Dataset
|
|
from torchvision.transforms import ColorJitter, InterpolationMode, Normalize
|
|
from torchvision.transforms.functional import (
|
|
hflip,
|
|
pad,
|
|
resize,
|
|
rotate,
|
|
to_pil_image,
|
|
to_tensor,
|
|
)
|
|
|
|
from shared.repositories import ImageRepository
|
|
|
|
# resnet18 original normalization values
|
|
RESNET_NORMALIZE_MEAN = [0.485, 0.456, 0.406]
|
|
RESNET_NORMALIZE_STD = [0.229, 0.224, 0.225]
|
|
|
|
|
|
class VCDADataset(Dataset):
|
|
"""VCDA dataset class."""
|
|
|
|
def __init__(
|
|
self,
|
|
data_name_list: list[str],
|
|
do_augment: bool = False,
|
|
random_annotations: bool = False,
|
|
):
|
|
super().__init__()
|
|
self.data_name_list = data_name_list
|
|
self.do_augment = do_augment
|
|
self.random_annotations = random_annotations
|
|
self.normalize_mean = RESNET_NORMALIZE_MEAN
|
|
self.normalize_std = RESNET_NORMALIZE_STD
|
|
# prepare augmentation functions
|
|
self.normalize = Normalize(
|
|
mean=self.normalize_mean,
|
|
std=self.normalize_std,
|
|
)
|
|
self.color_jitter = ColorJitter(
|
|
brightness=1e-1,
|
|
contrast=8e-2,
|
|
saturation=8e-2,
|
|
)
|
|
|
|
def __len__(self):
|
|
return len(self.data_name_list)
|
|
|
|
def __getitem__(self, idx):
|
|
# get image from database
|
|
image_name = self.data_name_list[idx]
|
|
with ImageRepository() as repo:
|
|
image_data = repo.get_data(image_name)
|
|
assert image_data is not None
|
|
tensor = self.image_to_tensor(image_data.image)
|
|
if self.do_augment:
|
|
tensor = self.augment(tensor)
|
|
return tensor
|
|
|
|
def image_to_tensor(
|
|
self,
|
|
image: Image.Image,
|
|
) -> Tensor:
|
|
"""Load images tensor from bytes."""
|
|
tensor = to_tensor(image)
|
|
tensor /= 255 # normalize 8-bit image
|
|
tensor = self.square_pad(tensor=tensor)
|
|
tensor = resize(
|
|
img=tensor,
|
|
size=(512, 512),
|
|
interpolation=InterpolationMode.BICUBIC,
|
|
)
|
|
return tensor
|
|
|
|
@staticmethod
|
|
def square_pad(
|
|
tensor: Tensor,
|
|
) -> Tensor:
|
|
"""Pads image to a square with side length equal to the largest side of
|
|
the input image."""
|
|
assert isinstance(tensor, Tensor)
|
|
# B, nc, w, h = img.shape
|
|
h = tensor.shape[-2]
|
|
w = tensor.shape[-1]
|
|
if h == w:
|
|
return tensor
|
|
max_wh = max([h, w])
|
|
hp = int((max_wh - w) / 2)
|
|
vp = int((max_wh - h) / 2)
|
|
padding = (hp, vp, hp, vp)
|
|
tensor = pad(tensor, padding, 0, 'constant')
|
|
return tensor
|
|
|
|
def augment(
|
|
self,
|
|
tensor: Tensor,
|
|
) -> Tensor:
|
|
"""Augment image with random horizontal flips, rotations and color
|
|
jitter."""
|
|
assert isinstance(tensor, Tensor)
|
|
# left-right flip
|
|
if random.random() >= 0.5:
|
|
tensor = hflip(tensor)
|
|
# rotation
|
|
rnd = random.random()
|
|
if rnd < 0.25:
|
|
tensor = rotate(tensor, angle=90)
|
|
if rnd < 0.5:
|
|
tensor = rotate(tensor, angle=180)
|
|
if rnd < 0.75:
|
|
tensor = rotate(tensor, angle=270)
|
|
# color jitter
|
|
tensor = self.color_jitter(tensor)
|
|
return tensor
|
|
|
|
def reverse_normalise(
|
|
self,
|
|
tensor: Tensor,
|
|
) -> Image.Image:
|
|
"""Reverse normalization to get an image that can be interpreted by
|
|
humans."""
|
|
assert isinstance(tensor, Tensor)
|
|
tensor *= Tensor(self.normalize_std).reshape((3, 1, 1))
|
|
tensor += Tensor(self.normalize_mean).reshape((3, 1, 1))
|
|
image = to_pil_image(
|
|
pic=tensor,
|
|
mode='RGB',
|
|
)
|
|
return image
|