moved for clarity

model/src/utils/vcda_dataset.py

@@ -0,0 +1,138 @@
+"""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.data_store import connect, get_image
+
+# 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,
+        )
+        # connect to minio
+        minio_client = connect()
+        self.minio_client = minio_client
+
+    def __len__(self):
+        return len(self.data_name_list)
+
+    def __getitem__(self, idx):
+        # get image from database
+        object_name = self.data_name_list[idx]
+        image = get_image(
+            client=self.minio_client,
+            object_name=object_name,
+        )
+        tensor = self.image_to_tensor(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