147 lines
4.1 KiB
Python
147 lines
4.1 KiB
Python
from __future__ import annotations
|
|
|
|
import random
|
|
from io import BytesIO
|
|
|
|
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 shared.database import connect
|
|
|
|
# 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):
|
|
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
|
|
# prepare augmentation functions
|
|
self.normalize = Normalize(
|
|
mean=normalize_mean,
|
|
std=normalize_std,
|
|
)
|
|
self.color_jitter = ColorJitter(
|
|
brightness=1e-1,
|
|
contrast=8e-2,
|
|
saturation=8e-2,
|
|
)
|
|
# connect to database
|
|
collection, _, _ = connect()
|
|
self.collection = collection
|
|
|
|
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,
|
|
)
|
|
img = self.load_image(img_bytes)
|
|
if self.do_augment:
|
|
img = self.augment(img)
|
|
return img
|
|
|
|
def load_image(
|
|
self,
|
|
data: bytes,
|
|
) -> 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,
|
|
size=(512, 512),
|
|
interpolation=InterpolationMode.BICUBIC,
|
|
)
|
|
return img
|
|
|
|
@staticmethod
|
|
def square_pad(
|
|
img: Tensor,
|
|
) -> Tensor:
|
|
"""
|
|
Pads image to a square with side length
|
|
equal to the largest side of the input image.
|
|
"""
|
|
assert isinstance(img, Tensor)
|
|
# B, nc, w, h = img.shape
|
|
h = img.shape[-2]
|
|
w = img.shape[-1]
|
|
if h == w:
|
|
return img
|
|
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')
|
|
|
|
def augment(
|
|
self,
|
|
img: Tensor,
|
|
) -> Tensor:
|
|
"""
|
|
Augment image with random horizontal flips,
|
|
rotations and color jitter.
|
|
"""
|
|
assert isinstance(img, Tensor)
|
|
# left-right flip
|
|
if random.random() >= 0.5:
|
|
img = hflip(img)
|
|
# rotation
|
|
rnd = random.random()
|
|
if rnd < 0.25:
|
|
img = rotate(img, angle=90)
|
|
if rnd < 0.5:
|
|
img = rotate(img, angle=180)
|
|
if rnd < 0.75:
|
|
img = rotate(img, angle=270)
|
|
# color jitter
|
|
img = self.color_jitter(img)
|
|
return img
|
|
|
|
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
|