From 9cd4f7019428550f5bde769f9aab1e52ededf2df Mon Sep 17 00:00:00 2001 From: Brian Bjarke Jensen Date: Sun, 14 Jul 2024 00:10:22 +0200 Subject: [PATCH] refactor for readability and implement new get_image function --- model/src/dataloader.py | 142 +++++++++++++++++++--------------------- 1 file changed, 67 insertions(+), 75 deletions(-) diff --git a/model/src/dataloader.py b/model/src/dataloader.py index 681e885..d5fc86a 100644 --- a/model/src/dataloader.py +++ b/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