model/src/dataloader.py

@@ -0,0 +1,151 @@
+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 core.database import connect
+from core.database import list_names
+
+# 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,
+        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.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
+        # load data names
+        has_annotation = False if self.random_annotations else True
+        self.data_name_list = list_names(
+            collection=self.collection,
+            has_annotation=has_annotation,
+        )
+
+    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

model/src/loss_fn.py

@@ -0,0 +1,7 @@
+from __future__ import annotations
+
+import torch.nn as nn
+
+
+def setup_criterion():
+    return nn.MSELoss()

model/src/model.py

@@ -0,0 +1,37 @@
+from __future__ import annotations
+
+import torch.nn as nn
+import torchvision
+
+
+class ResNet18Head(nn.Module):
+    def __init__(self):
+        super().__init__()
+        # copy out parts from ResNet18 with weights
+        resnet18 = torchvision.models.resnet18(
+            weights=torchvision.models.ResNet18_Weights.IMAGENET1K_V1,
+        )
+        # save relevant layers
+        self.conv1 = resnet18.conv1
+        self.bn1 = resnet18.bn1
+        self.relu = resnet18.relu
+        self.maxpool = resnet18.maxpool
+        self.layer1 = resnet18.layer1
+        self.layer2 = resnet18.layer2
+        self.layer3 = resnet18.layer3
+        self.layer4 = resnet18.layer4
+        self.avgpool = resnet18.avgpool
+        self.flat = nn.Flatten()    # size 512
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+        x = self.maxpool(x)
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+        x = self.avgpool(x)
+        x = self.flat(x)
+        return x

model/src/train.py

@@ -0,0 +1,6 @@
+from __future__ import annotations
+
+from dataloader import VCDADataset  # noqa: F401
+from loss_fn import setup_criterion  # noqa: F401
+
+from model import ResNet18Head  # noqa: F401