model/src/dataloader.py

@@ -1,20 +1,21 @@
-from __future__ import annotations
+"""Definition of Visual Critial Discourse Analysis dataloader."""
 
 import random
-from io import BytesIO
 
+from PIL import Image
 from torch import Tensor
 from torch.utils.data import Dataset
-from torchvision.io import read_image
-from torchvision.transforms import ColorJitter
-from torchvision.transforms import InterpolationMode
-from torchvision.transforms import Normalize
-from torchvision.transforms.functional import hflip
-from torchvision.transforms.functional import pad
-from torchvision.transforms.functional import resize
-from torchvision.transforms.functional import rotate
+from torchvision.transforms import ColorJitter, InterpolationMode, Normalize
+from torchvision.transforms.functional import (
+    hflip,
+    pad,
+    resize,
+    rotate,
+    to_pil_image,
+    to_tensor,
+)
 
-from shared.database import connect
+from shared.data_store import connect, get_image
 
 # resnet18 original normalization values
 RESNET_NORMALIZE_MEAN = [0.485, 0.456, 0.406]
@@ -22,125 +23,116 @@ 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,
-        normalize_mean: list[float] = RESNET_NORMALIZE_MEAN,
-        normalize_std: list[float] = RESNET_NORMALIZE_STD,
     ):
         super().__init__()
         self.data_name_list = data_name_list
         self.do_augment = do_augment
         self.random_annotations = random_annotations
-        self.normalize_mean = normalize_mean
-        self.normalize_std = normalize_std
+        self.normalize_mean = RESNET_NORMALIZE_MEAN
+        self.normalize_std = RESNET_NORMALIZE_STD
         # prepare augmentation functions
         self.normalize = Normalize(
-            mean=normalize_mean,
-            std=normalize_std,
+            mean=self.normalize_mean,
+            std=self.normalize_std,
         )
         self.color_jitter = ColorJitter(
             brightness=1e-1,
             contrast=8e-2,
             saturation=8e-2,
         )
-        # connect to database
-        collection, _, _ = connect()
-        self.collection = collection
+        # 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
-        name = self.data_name_list[idx]
-        query = {
-            'name': name,
-        }
-        projection = {
-            '_id': False,
-            'image': True,
-        }
-        img_bytes = self.collection.find_one(
-            filter=query,
-            projection=projection,
+        object_name = self.data_name_list[idx]
+        image = get_image(
+            client=self.minio_client,
+            object_name=object_name,
         )
-        img = self.load_image(img_bytes)
+        tensor = self.image_to_tensor(image)
         if self.do_augment:
-            img = self.augment(img)
-        return img
+            tensor = self.augment(tensor)
+        return tensor
 
-    def load_image(
+    def image_to_tensor(
         self,
-        data: bytes,
+        image: Image.Image,
     ) -> Tensor:
         """Load images tensor from bytes."""
-        assert isinstance(data, bytes)
-        img = read_image(BytesIO(data))
-        img /= 255  # normalize 8-bit image
-        img = self.square_pad(img)
-        img = resize(
-            img=img,
+        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 img
+        return tensor
 
     @staticmethod
     def square_pad(
-        img: Tensor,
+        tensor: Tensor,
     ) -> Tensor:
-        """
-        Pads image to a square with side length
-        equal to the largest side of the input image.
-        """
-        assert isinstance(img, 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 = img.shape[-2]
-        w = img.shape[-1]
+        h = tensor.shape[-2]
+        w = tensor.shape[-1]
         if h == w:
-            return img
+            return tensor
         max_wh = max([h, w])
         hp = int((max_wh - w) / 2)
         vp = int((max_wh - h) / 2)
         padding = (hp, vp, hp, vp)
-        return pad(img, padding, 0, 'constant')
+        tensor = pad(tensor, padding, 0, 'constant')
+        return tensor
 
     def augment(
         self,
-        img: Tensor,
+        tensor: Tensor,
     ) -> Tensor:
-        """
-        Augment image with random horizontal flips,
-        rotations and color jitter.
-        """
-        assert isinstance(img, Tensor)
+        """Augment image with random horizontal flips, rotations and color
+        jitter."""
+        assert isinstance(tensor, Tensor)
         # left-right flip
         if random.random() >= 0.5:
-            img = hflip(img)
+            tensor = hflip(tensor)
         # rotation
         rnd = random.random()
         if rnd < 0.25:
-            img = rotate(img, angle=90)
+            tensor = rotate(tensor, angle=90)
         if rnd < 0.5:
-            img = rotate(img, angle=180)
+            tensor = rotate(tensor, angle=180)
         if rnd < 0.75:
-            img = rotate(img, angle=270)
+            tensor = rotate(tensor, angle=270)
         # color jitter
-        img = self.color_jitter(img)
-        return img
+        tensor = self.color_jitter(tensor)
+        return tensor
 
     def reverse_normalise(
         self,
-        img: Tensor,
-    ) -> Tensor:
-        """
-        Reverse normalization to get an image
-        that can be interpreted by humans.
-        """
-        assert isinstance(img, Tensor)
-        img *= Tensor(self.normalize_std).reshape((3, 1, 1))
-        img += Tensor(self.normalize_mean).reshape((3, 1, 1))
-        return img
+        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