Merge pull request 'image_through_model' (#49) from image_through_model into main

Reviewed-on: http://192.168.1.2:3000/brian/visual_critical_discourse_analysis/pulls/49

.gitignore

@@ -177,3 +177,6 @@ ipython_config.py
 
 # Remove previous ipynb_checkpoints
 #   git rm -r .ipynb_checkpoints/
+
+# Logging data
+runs/*

.pre-commit-config.yaml

@@ -8,17 +8,16 @@ repos:
     -   id: debug-statements
     -   id: double-quote-string-fixer
     -   id: name-tests-test
-    -   id: requirements-txt-fixer
 -   repo: https://github.com/asottile/setup-cfg-fmt
     rev: v2.5.0
     hooks:
     -   id: setup-cfg-fmt
--   repo: https://github.com/asottile/reorder-python-imports
+-   repo: https://github.com/pre-commit/mirrors-isort
-    rev: v3.12.0
+    rev: v5.10.1
     hooks:
-    -   id: reorder-python-imports
+      - id: isort
-        exclude: ^(pre_commit/resources/|testing/resources/python3_hooks_repo/)
+        language_version: python3.10
-        args: [--py39-plus, --add-import, 'from __future__ import annotations']
+        args: [ --tc ]
 -   repo: https://github.com/asottile/add-trailing-comma
     rev: v3.1.0
     hooks:
@@ -28,10 +27,6 @@ repos:
     hooks:
     -   id: pyupgrade
         args: [--py39-plus]
--   repo: https://github.com/hhatto/autopep8
-    rev: v2.0.4
-    hooks:
-    -   id: autopep8
 -   repo: https://github.com/PyCQA/flake8
     rev: 7.0.0
     hooks:
@@ -43,5 +38,24 @@ repos:
     rev: v1.8.0
     hooks:
     -   id: mypy
-        additional_dependencies: [types-all]
         exclude: ^testing/resources/
+-   repo: https://github.com/psf/black
+    rev: 24.4.2
+    hooks:
+    - id: black
+      language_version: python3.12
+      args: [ --skip-string-normalization ]
+-   repo: https://github.com/myint/docformatter
+    rev: v1.5.0
+    hooks:
+      - id: docformatter
+        name: docformatter
+        description: 'Formats docstrings to follow PEP 257.'
+        entry: docformatter
+        language: python
+        types: [ python ]
+-     repo: https://github.com/jendrikseipp/vulture
+      rev: 'v2.6'
+      hooks:
+      - id: vulture
+        entry: vulture . --min-confidence 90 --exclude */.venv/*.py

Dockerfile.model

@@ -1,5 +1,5 @@
 # build stage
-FROM python:3.12-slim-bookworm as BUILDER
+FROM python:3.12-slim-bookworm AS builder
 
 # set environment variables
 ENV PYTHONUNBUFFERED=1 \
@@ -44,7 +44,7 @@ COPY --from=builder ${VIRTUAL_ENV} ${VIRTUAL_ENV}
 ENV PATH="${VIRTUAL_ENV}/bin:${PATH}"
 
 # create home directory and app user
-RUN mkdir -p /home/app
+RUN mkdir -p $APP_HOME
 
 # add code while changing ownership
 WORKDIR $APP_HOME

Dockerfile.web_ui

@@ -1,5 +1,5 @@
 # build stage
-FROM python:3.12-slim-bookworm as BUILDER
+FROM python:3.12-slim-bookworm AS builder
 
 # set environment variables
 ENV PYTHONUNBUFFERED=1 \

model/docker-compose.server.yml

@@ -9,3 +9,8 @@ services:
       - ../server.env
     environment:
       - ENV=TEST
+    deploy:
+      resources:
+        limits:
+          cpus: '2.000'
+          memory: 8G

model/src/configs/resnet18.json

@@ -0,0 +1,14 @@
+{
+    "datasets": {
+        "train": {
+            "filelist": "datasets/train.csv"
+        },
+        "val": {
+            "filelist": "datasets/val.csv"
+        }
+    },
+    "backbone": {
+        "class": "model.src.models.VisualCommunicationModel",
+        "model_name": "pretrained"
+    }
+}

model/src/data_store/__init__.py

@@ -1,5 +0,0 @@
-from __future__ import annotations
-
-from .connect import connect
-from .get import get
-from .put import put

model/src/data_store/connect.py

@@ -1,30 +0,0 @@
-"""Definition of connect function."""
-from __future__ import annotations
-
-import os
-
-from minio import Minio
-
-
-def connect(
-) -> Minio:
-    """Connect to MinIO server"""
-    minio_endpoint = os.getenv('MINIO_ENDPOINT', default=None)
-    assert isinstance(minio_endpoint, str)
-    minio_access_key = os.getenv('MINIO_ACCESS_KEY', default=None)
-    assert isinstance(minio_access_key, str)
-    minio_secret_key = os.getenv('MINIO_SECRET_KEY', default=None)
-    assert isinstance(minio_secret_key, str)
-    minio_bucket_name = os.getenv('MINIO_BUCKET_NAME', default=None)
-    assert isinstance(minio_bucket_name, str)
-    # connect client
-    client = Minio(
-        endpoint=minio_endpoint,
-        access_key=minio_access_key,
-        secret_key=minio_secret_key,
-        secure=False,
-    )
-    # ensure bucket exists
-    if not client.bucket_exists(bucket_name=minio_bucket_name):
-        client.make_bucket(bucket_name=minio_bucket_name)
-    return client

model/src/data_store/get.py

@@ -1,27 +0,0 @@
-"""Definition of get function."""
-from __future__ import annotations
-
-import os
-from io import BytesIO
-
-from minio import Minio
-
-
-def get(
-    client: Minio,
-    object_name: str,
-) -> BytesIO:
-    """Get buffer from bucket in MinIO."""
-    assert isinstance(client, Minio)
-    assert isinstance(object_name, str)
-    bucket_name = os.getenv('MINIO_BUCKET_NAME', default=None)
-    assert isinstance(bucket_name, str)
-    # get buffer
-    try:
-        response = client.get_object(bucket_name, object_name)
-        buffer = BytesIO(response.data)
-    finally:
-        response.close()
-        response.release_conn()
-    buffer.seek(0)
-    return buffer

model/src/data_store/put.py

@@ -1,30 +0,0 @@
-"""Definition of put function."""
-from __future__ import annotations
-
-import os
-from io import BytesIO
-
-from minio import Minio
-
-
-def put(
-    client: Minio,
-    buffer: BytesIO,
-    object_name: str,
-) -> None:
-    """Put buffer in bucket in MinIO."""
-    assert isinstance(client, Minio)
-    assert isinstance(buffer, BytesIO)
-    assert isinstance(object_name, str)
-    bucket_name = os.getenv('MINIO_BUCKET_NAME', default=None)
-    assert isinstance(bucket_name, str)
-    # prepare for saving
-    num_bytes = buffer.tell()
-    buffer.seek(0)
-    # send data to bucket
-    client.put_object(
-        bucket_name=bucket_name,
-        object_name=object_name,
-        length=num_bytes,
-        data=buffer,
-    )

model/src/dataloader.py

@@ -1,146 +0,0 @@
-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

model/src/main.py

@@ -1,45 +1,48 @@
 """Main script to be run by service."""
-from __future__ import annotations
 
-from hashlib import md5
+import json
-from io import BytesIO
+import logging
+from traceback import print_exc
 
 import torch
-from data_store import connect
-from data_store import put
 from models import VisualCommunicationModel
-from torchinfo import summary
+from tqdm import tqdm
-from utils import check_env
+from utils import DEVICE, VCDADataset, load_model
 
+from shared.data_store import connect_minio
+from shared.dto import ModelData
 from shared.utils import setup_logging
 
-
-DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-
 if __name__ == '__main__':
-    # check that environment variables are set
-    check_env()
     # setup logging
     setup_logging()
-    # instantiate model
-    model = VisualCommunicationModel().to(DEVICE)
-    # show model weights
-    summary(model)
-    # for param in model.parameters():
-    #     logging.info(param.data)
-    # save model to buffer
-    buffer = BytesIO()
-    torch.save(model.state_dict(), buffer)
-    print(f"buffer size: {len(buffer.getvalue())}")
-    # calculate buffer hash
-    hash_str = md5(buffer.getbuffer()).hexdigest()
-    print(f"hash string: {hash_str}")
     # connect to minio
-    client = connect()
+    minio_client = connect_minio()
-    # put buffer in minio bucket
+    # instantiate model
-    put(
+    model: VisualCommunicationModel = load_model(client=minio_client)
-        client=client,
+    model.eval()
-        buffer=buffer,
+    # setup dataset
-        object_name=hash_str,
+    dataset = VCDADataset(
+        minio_client=minio_client,
+        data_name_list=[
+            '02dbaf48d713e4e6d3a6b98fd2dc866e',
+        ],
+        do_augment=False,
     )
-    print('saved to Minio')
+
+    # make prediction
+    with torch.no_grad():
+        for i in tqdm(range(len(dataset))):
+            try:
+                # get image
+                image = dataset[i]
+                image = torch.unsqueeze(image, 0)  # add artificial batch dimension
+                image = image.to(DEVICE)
+                # make prediction
+                pred: ModelData = model(image)
+            except Exception:
+                print_exc()
+                continue
+            else:
+                print(json.dumps(pred.model_dump(), indent=4))
+    logging.debug('finished')

model/src/model_name.txt

@@ -0,0 +1 @@
+2bbe2be17a663f4299657378ac5e993e

model/src/models/angle.py

@@ -1,5 +1,3 @@
-from __future__ import annotations
-
 from .fully_connected import FullyConnectedModel
 
 

model/src/models/contact.py

@@ -1,5 +1,3 @@
-from __future__ import annotations
-
 from .fully_connected import FullyConnectedModel
 
 

model/src/models/distance.py

@@ -1,5 +1,3 @@
-from __future__ import annotations
-
 from .fully_connected import FullyConnectedModel
 
 

model/src/models/framing.py

@@ -1,5 +1,3 @@
-from __future__ import annotations
-
 from .fully_connected import FullyConnectedModel
 
 

model/src/models/fully_connected.py

@@ -1,26 +1,24 @@
-from __future__ import annotations
+from torch import nn
-
-import torch.nn as nn
 
 
 class FullyConnectedModel(nn.Module):
+    """Fully connected layers model template for intrepreting feature space-
+    output from ResNet18 head."""
+
     def __init__(self, num_out_features: int):
         super().__init__()
         # define layers
-        self.fc1 = nn.Linear(in_features=16*16*512, out_features=512)
+        self.fc1 = nn.Linear(in_features=512, out_features=128)
         self.af1 = nn.ReLU()
-        self.fc2 = nn.Linear(in_features=512, out_features=128)
+        self.fc2 = nn.Linear(in_features=128, out_features=32)
         self.af2 = nn.ReLU()
-        self.fc3 = nn.Linear(in_features=128, out_features=32)
+        self.fc3 = nn.Linear(in_features=32, out_features=num_out_features)
-        self.af3 = nn.ReLU()
-        self.fc4 = nn.Linear(in_features=32, out_features=num_out_features)
 
     def forward(self, x):
+        """Pass input through model."""
         x = self.fc1(x)
         x = self.af1(x)
         x = self.fc2(x)
         x = self.af2(x)
         x = self.fc3(x)
-        x = self.af3(x)
-        x = self.fc4(x)
         return x

model/src/models/information_value.py

@@ -1,5 +1,3 @@
-from __future__ import annotations
-
 from .fully_connected import FullyConnectedModel
 
 

model/src/models/modality_color.py

@@ -1,5 +1,3 @@
-from __future__ import annotations
-
 from .fully_connected import FullyConnectedModel
 
 

model/src/models/modality_depth.py

@@ -1,5 +1,3 @@
-from __future__ import annotations
-
 from .fully_connected import FullyConnectedModel
 
 

model/src/models/modality_lighting.py

@@ -1,5 +1,3 @@
-from __future__ import annotations
-
 from .fully_connected import FullyConnectedModel
 
 

model/src/models/point_of_view.py

@@ -1,5 +1,3 @@
-from __future__ import annotations
-
 from .fully_connected import FullyConnectedModel
 
 

model/src/models/resnet18_head.py

@@ -5,11 +5,15 @@ import torchvision
 
 
 class ResNet18Head(nn.Module):
-    def __init__(self):
+    def __init__(self, download_resnet_weights: bool = False):
         super().__init__()
         # copy out parts from ResNet18 with weights
+        if download_resnet_weights:
+            weights = torchvision.models.ResNet18_Weights.IMAGENET1K_V1
+        else:
+            weights = None
         resnet18 = torchvision.models.resnet18(
-            weights=torchvision.models.ResNet18_Weights.IMAGENET1K_V1,
+            weights=weights,
         )
         # save relevant layers
         self.conv1 = resnet18.conv1

model/src/models/salience.py

@@ -1,5 +1,3 @@
-from __future__ import annotations
-
 from .fully_connected import FullyConnectedModel
 
 

model/src/models/visual_communication.py

@@ -1,8 +1,11 @@
 """Definition of VisualCommunicationModel class."""
+
 from __future__ import annotations
 
 from torch import nn
 
+from shared.dto import ModelData
+
 from .angle import AngleTail
 from .contact import ContactTail
 from .distance import DistanceTail
@@ -15,26 +18,15 @@ from .point_of_view import PointOfViewTail
 from .resnet18_head import ResNet18Head
 from .salience import SalienceTail
 from .visual_syntax import VisualSyntaxTail
-from shared.dto import AngleData
-from shared.dto import ContactData
-from shared.dto import DistanceData
-from shared.dto import FramingData
-from shared.dto import InformationValueData
-from shared.dto import ModalityColorData
-from shared.dto import ModalityDepthData
-from shared.dto import ModalityLightingData
-from shared.dto import PointOfViewData
-from shared.dto import SalienceData
-from shared.dto import VisualSyntaxData
 
 
 class VisualCommunicationModel(nn.Module):
     """Visual communication model."""
 
-    def __init__(self):
+    def __init__(self, download_resnet_weights: bool = False):
         super().__init__()
         # store other models
-        self.resnet_head = ResNet18Head()
+        self.resnet_head = ResNet18Head(download_resnet_weights)
         self.visual_syntax_tail = VisualSyntaxTail()
         self.contact_tail = ContactTail()
         self.angle_tail = AngleTail()
@@ -47,76 +39,24 @@ class VisualCommunicationModel(nn.Module):
         self.framing_tail = FramingTail()
         self.salience_tail = SalienceTail()
 
-    def forward(self, x):
+    def forward(self, x) -> ModelData:
         """Calculate model output on data."""
         # generate visual representation
-        vis_rep = self.resnet_head(x)
+        features = self.resnet_head(x)
-        # prepare result map
+        # make predictions
-        results = {}
+        prediction_dict = {
-        # predict visual syntax
+            'visual_syntax': self.visual_syntax_tail(features),
-        results['visual_syntax'] = VisualSyntaxData.from_list(
+            'contact': self.contact_tail(features),
-            self.visual_syntax_tail(
+            'angle': self.angle_tail(features),
-                vis_rep,
+            'point_of_view': self.point_of_view_tail(features),
-            ).cpu(),
+            'distance': self.distance_tail(features),
-        )
+            'modality_lighting': self.modality_lighting_tail(features),
-        # predict contact
+            'modality_color': self.modality_color_tail(features),
-        results['contact'] = ContactData.from_list(
+            'modality_depth': self.modality_depth_tail(features),
-            self.contact_tail(
+            'information_value': self.information_value_tail(features),
-                vis_rep,
+            'framing': self.framing_tail(features),
-            ).cpu(),
+            'salience': self.salience_tail(features),
-        )
+        }
-        # predict angle
+        # convert to respective classes
-        results['angle'] = AngleData.from_list(
+        data = ModelData.from_prediction_dict(prediction_dict)
-            self.angle_tail(
+        return data
-                vis_rep,
-            ).cpu(),
-        )
-        # predict point of view
-        results['point_of_view'] = PointOfViewData.from_list(
-            self.point_of_view_tail(
-                vis_rep,
-            ).cpu(),
-        )
-        # predict distance
-        results['distance'] = DistanceData.from_list(
-            self.distance_tail(
-                vis_rep,
-            ).cpu(),
-        )
-        # predict modality lighting
-        results['modality_lighting'] = ModalityLightingData.from_list(
-            self.modality_lighting_tail(
-                vis_rep,
-            ).cpu(),
-        )
-        # predict modality color
-        results['modality_color'] = ModalityColorData.from_list(
-            self.modality_color_tail(
-                vis_rep,
-            ).cpu(),
-        )
-        # predict modality depth
-        results['modality_depth'] = ModalityDepthData.from_list(
-            self.modality_depth_tail(
-                vis_rep,
-            ).cpu(),
-        )
-        # predict information value
-        results['information_value'] = InformationValueData.from_list(
-            self.information_value_tail(
-                vis_rep,
-            ).cpu(),
-        )
-        # predict framing
-        results['framing'] = FramingData.from_list(
-            self.framing_tail(
-                vis_rep,
-            ).cpu(),
-        )
-        # predict salience
-        results['salience'] = SalienceData.from_list(
-            self.salience_tail(
-                vis_rep,
-            ).cpu(),
-        )
-        return results

model/src/models/visual_syntax.py

@@ -1,5 +1,3 @@
-from __future__ import annotations
-
 from .fully_connected import FullyConnectedModel
 
 

model/src/utils/__init__.py

@@ -1,3 +1,3 @@
-from __future__ import annotations
+from .get_class import get_class
-
+from .load_model import DEVICE, load_model
-from .check_env import check_env
+from .vcda_dataset import VCDADataset

model/src/utils/check_env.py

@@ -1,25 +0,0 @@
-"""Definition of check_env function."""
-from __future__ import annotations
-
-import os
-
-
-def check_env() -> None:
-    """Check necessary environment variables are set."""
-    necesasary_var_list = {
-        'MONGO_HOST',
-        'MONGO_DB',
-        'MONGO_COLLECTION',
-        'MONGO_USER',
-        'MONGO_PASSWORD',
-        'MINIO_ENDPOINT',
-        'MINIO_ACCESS_KEY',
-        'MINIO_SECRET_KEY',
-    }
-    for env_var in necesasary_var_list:
-        # ensure env var set
-        assert (
-            env_var in os.environ
-        ), (
-            f"environment variable not set: {env_var}"
-        )

model/src/utils/get_class.py

@@ -0,0 +1,11 @@
+from importlib import import_module
+
+from torch.nn import Module
+
+
+def get_class(path: str) -> Module:
+    parts = path.split('.')
+    module_path = '.'.join(parts[:-1])
+    class_name = parts[-1]
+    module = import_module(module_path)
+    return getattr(module, class_name)

model/src/utils/get_model_name.py

@@ -0,0 +1,19 @@
+"""Definition of get_model_name function."""
+
+import logging
+from pathlib import Path
+
+
+def get_model_name(
+    path: Path,
+) -> str:
+    """Get model name from model_name file."""
+    assert isinstance(path, Path)
+    assert path.exists(), f'{path} does not exist'
+    # read file contents
+    with open(file=path, encoding='utf-8') as fh:
+        model_name = fh.read()
+    # remove newline character
+    model_name = model_name.strip()
+    logging.debug('finished')
+    return model_name

model/src/utils/load_model.py

@@ -0,0 +1,40 @@
+"""Definition of load_model function."""
+
+import logging
+from pathlib import Path
+
+import torch
+from minio import Minio
+
+from model.src.models import VisualCommunicationModel
+from shared.data_store import get_model
+
+from .get_model_name import get_model_name
+
+DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+
+def load_model(
+    client: Minio,
+) -> VisualCommunicationModel:
+    """Instantiate model with weights loaded from latest model saved in
+    MinIO."""
+    assert isinstance(client, Minio)
+    # instantiate model
+    model = VisualCommunicationModel()
+    # get model object name
+    model_name_path = Path('model_name.txt')
+    model_object_name = get_model_name(path=model_name_path)
+    logging.info('using model: %s', model_object_name)
+    # load model from minio
+    model_checkpoint = get_model(
+        client=client,
+        object_name=model_object_name,
+    )
+    model.load_state_dict(model_checkpoint)
+    # clean memory
+    model_checkpoint.clear()
+    # move model to selected device
+    model = model.to(DEVICE)
+    logging.debug('finished')
+    return model

model/src/utils/train.py

@@ -6,7 +6,7 @@ import torch
 import torch.nn as nn
 from dataloader import VCDADataset  # noqa: F401
 
-from .models import VisualCommunicationModel
+from ..models import VisualCommunicationModel
 
 PRE_WARMUP_LR = 1e-10
 POST_WARMUP_LR = 1e-5

model/src/utils/vcda_dataset.py

@@ -0,0 +1,138 @@
+"""Definition of Visual Critial Discourse Analysis dataloader."""
+
+import random
+
+from minio import Minio
+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 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,
+        minio_client: Minio,
+        data_name_list: list[str],
+        do_augment: bool = False,
+        random_annotations: bool = False,
+    ):
+        super().__init__()
+        self.minio_client = minio_client
+        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,
+        )
+
+    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

model/train.py

@@ -0,0 +1,166 @@
+import json
+import os
+from argparse import ArgumentParser
+from datetime import datetime, timezone
+from pathlib import Path
+
+import torch
+from dotenv import load_dotenv
+from ignite.engine import Events, create_supervised_evaluator, create_supervised_trainer
+from ignite.handlers import global_step_from_engine
+from ignite.handlers.checkpoint import Checkpoint
+from ignite.handlers.param_scheduler import create_lr_scheduler_with_warmup
+from ignite.handlers.tensorboard_logger import TensorboardLogger
+from ignite.handlers.tqdm_logger import ProgressBar
+from ignite.metrics import Average, Loss, RunningAverage
+from torch import nn
+from torch.optim.lr_scheduler import ExponentialLR
+from torch.utils.data import DataLoader
+
+from model.src.utils import VCDADataset, get_class
+from shared.data_store import connect_minio
+
+
+def parse_arguments():
+    parser = ArgumentParser()
+    parser.add_argument('-c', '--config', type=Path, required=True)
+    parser.add_argument('-d', '--device', type=str, default='cpu')
+    parser.add_argument('-b', '--batch-size', type=int, default=32)
+    parser.add_argument('--epochs', type=int, default=50)
+    parser.add_argument('--epoch-length', type=int, default=128)
+    parser.add_argument('--checkpoint', type=Path)
+    parser.add_argument('--lr', type=float, default=1e-5)
+    parser.add_argument('--lr-gamma', type=float, default=0.95)
+    parser.add_argument('--lr-warmup-start', type=float, default=1e-8)
+    parser.add_argument('--lr-warmup-duration', type=int, default=5)
+    parser.add_argument('--train-seed', type=int, default=1312)
+    parser.add_argument('--val-seed', type=int, default=1313)
+    parser.add_argument('--loader-workers', type=int, default=8)
+    return parser.parse_args()
+
+
+args = parse_arguments()
+
+load_dotenv('server.env')
+
+# read model config
+with open(args.config) as fh:
+    config = json.load(fh)
+model_name = args.config.stem
+run_name = datetime.now(timezone.utc).strftime('%Y-%m-%d-%H%M') + '_' + model_name
+device = torch.device(args.device)
+
+# create model
+model_class = get_class(config['backbone']['class'])
+model = model_class().to(device)
+optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
+loss_fn = nn.CrossEntropyLoss()
+
+# create datasets and loaders
+minio_client = connect_minio()
+with open('model/src/dataset/train.csv', encoding='utf-8') as fh:
+    train_data_name_list = fh.read().split('\n')
+train_dataset = VCDADataset(
+    minio_client=minio_client,
+    data_name_list=train_data_name_list,
+)
+train_loader = DataLoader(dataset=train_dataset, num_workers=args.loader_workers)
+
+with open('model/src/dataset/val.csv', encoding='utf-8') as fh:
+    val_data_name_list = fh.read().split('\n')
+val_dataset = VCDADataset(minio_client=minio_client, data_name_list=val_data_name_list)
+val_loader = DataLoader(dataset=val_dataset, num_workers=args.loader_workers)
+
+# create trainer and evaluator
+trainer = create_supervised_trainer(
+    model=model,
+    optimizer=optimizer,
+    loss_fn=loss_fn,
+    device=device,
+)
+Average(output_transform=lambda x: x).attach(trainer, name='loss')
+RunningAverage(output_transform=lambda x: x, alpha=0.5).attach(
+    trainer,
+    'running_avg_loss',
+)
+ProgressBar(desc='Train', ncols=80).attach(trainer, ['running_avg_loss'])
+
+val_metrics = {
+    'loss': Loss(loss_fn=loss_fn, device=device),
+}
+evaluator = create_supervised_evaluator(
+    model=model,
+    metrics=val_metrics,  # type: ignore
+    device=device,
+)
+ProgressBar(desc='Val', ncols=80).attach(evaluator)
+
+# create lr scheduler
+lr_scheduler = ExponentialLR(
+    optimizer=optimizer,
+    gamma=args.lr_gamma,
+)
+lr_handler = create_lr_scheduler_with_warmup(
+    lr_scheduler=lr_scheduler,
+    warmup_start_value=args.lr_warmup_start,
+    warmup_duration=args.lr_warmup_duration,
+    warmup_end_value=args.lr,
+)
+trainer.add_event_handler(Events.EPOCH_STARTED, lr_handler)
+
+
+# log metrics to TensorBoard
+@trainer.on(Events.EPOCH_COMPLETED)
+def evaluate():
+    evaluator.run(val_loader, epoch_length=args.val_epoch_length)
+
+
+tb_logger = TensorboardLogger(log_dir=f'runs/logs/{run_name}')
+tb_logger.attach_opt_params_handler(
+    engine=trainer,
+    event_name=Events.EPOCH_COMPLETED,
+    optimizer=optimizer,
+)
+for tag, engine in [('train', trainer), ('val', evaluator)]:
+    tb_logger.attach_output_handler(
+        engine,
+        event_name=Events.EPOCH_COMPLETED,
+        tag=tag,
+        metric_names='all',
+        global_step_transform=global_step_from_engine(trainer),
+    )
+
+# Set up checkpoint saving
+to_save = {
+    'model': model,
+    'optimizer': optimizer,
+    'trainer': trainer,
+    'lr_scheduler': lr_scheduler,
+    'lr_handler': lr_handler,
+}
+checkpoint_handler = Checkpoint(
+    to_save,
+    f"runs/checkpoints/{run_name}",
+    n_saved=3,
+    filename_prefix='best',
+    score_function=lambda engine: -engine.state.metrics['loss'],
+    score_name='neg_val_loss',
+    global_step_transform=global_step_from_engine(trainer),
+)
+evaluator.add_event_handler(Events.COMPLETED, checkpoint_handler)
+
+if args.checkpoint:
+    Checkpoint.load_objects(to_load=to_save, checkpoint=str(args.checkpoint))
+
+# save model config
+os.makedirs('runs/configs/', exist_ok=True)
+with open(f"runs/configs/{run_name}.json", 'w', encoding='utf-8') as fh:
+    json.dump(config, fh)
+
+# start training
+trainer.run(
+    train_loader,
+    max_epochs=args.epochs,
+    epoch_length=args.epoch_length,
+)
+tb_logger.close()

other/transfer_images_minio_subfolder.py

@@ -0,0 +1,38 @@
+"""Script to move minio images to subfolder."""
+
+from pathlib import Path
+
+from dotenv import load_dotenv
+from PIL import Image
+
+from shared.data_store import connect_minio, get, put_image
+from shared.utils import setup_logging
+
+if __name__ == '__main__':
+    # load in env file
+    env_path = Path(__file__).parent.parent / 'server.env'
+    assert env_path.exists()
+    load_dotenv(env_path)
+    # setup logging
+    setup_logging()
+    # connect to minio
+    minio_client = connect_minio()
+    # list images in bucket
+    BUCKET_NAME = 'visual-critical-discourse-analysis'
+    obj_list = minio_client.list_objects(
+        bucket_name=BUCKET_NAME,
+    )
+    # begin moving images
+    for obj in obj_list:
+        # get image from minio
+        buffer = get(
+            client=minio_client,
+            object_name=obj.object_name,
+        )
+        # convert data to image
+        image = Image.open(buffer)
+        # put image into minio subfolder
+        put_image(
+            client=minio_client,
+            image=image,
+        )

other/transfer_images_mongo_minio.py

@@ -1,4 +1,5 @@
 """Script to move all images from MongoDB to MinIO."""
+
 from __future__ import annotations
 
 import logging
@@ -9,12 +10,9 @@ from bson import ObjectId
 from dotenv import load_dotenv
 from pymongo.collection import Collection
 
-from shared.data_store import connect as connect_minio
+from shared.data_store import connect_minio, put
-from shared.data_store import put
+from shared.database import VisualCommunication, connect_mongodb
-from shared.database import connect
+from shared.utils import check_env, setup_logging
-from shared.database import VisualCommunication
-from shared.utils import check_env
-from shared.utils import setup_logging
 
 
 def list_mongo_document_ids(
@@ -84,7 +82,7 @@ if __name__ == '__main__':
     # connect to minIO
     minio_client = connect_minio()
     # connect to MongoDB
-    collection, db, client = connect()
+    collection, db, client = connect_mongodb()
     # list documents in mongoDB
     id_list = list_mongo_document_ids(collection)
     for doc_id in id_list:
@@ -116,7 +114,8 @@ if __name__ == '__main__':
         except Exception as exc:
             logging.debug(exc)
             logging.error(
-                'failed updating visual communication %s', vis_com.name,
+                'failed updating visual communication %s',
+                vis_com.name,
             )
             continue
         logging.debug('updated visual communication %s', vis_com.name)

poetry.lock

@@ -1,4 +1,20 @@
-# This file is automatically @generated by Poetry 1.8.2 and should not be changed by hand.
+# This file is automatically @generated by Poetry 1.8.3 and should not be changed by hand.
+
+[[package]]
+name = "absl-py"
+version = "2.1.0"
+description = "Abseil Python Common Libraries, see https://github.com/abseil/abseil-py."
+optional = false
+python-versions = ">=3.7"
+files = [
+    {file = "absl-py-2.1.0.tar.gz", hash = "sha256:7820790efbb316739cde8b4e19357243fc3608a152024288513dd968d7d959ff"},
+    {file = "absl_py-2.1.0-py3-none-any.whl", hash = "sha256:526a04eadab8b4ee719ce68f204172ead1027549089702d99b9059f129ff1308"},
+]
+
+[package.source]
+type = "legacy"
+url = "http://192.168.1.2:5001/index"
+reference = "threadripper"
 
 [[package]]
 name = "annotated-types"
@@ -670,6 +686,69 @@ type = "legacy"
 url = "http://192.168.1.2:5001/index"
 reference = "threadripper"
 
+[[package]]
+name = "grpcio"
+version = "1.66.1"
+description = "HTTP/2-based RPC framework"
+optional = false
+python-versions = ">=3.8"
+files = [
+    {file = "grpcio-1.66.1-cp310-cp310-linux_armv7l.whl", hash = "sha256:4877ba180591acdf127afe21ec1c7ff8a5ecf0fe2600f0d3c50e8c4a1cbc6492"},
+    {file = "grpcio-1.66.1-cp310-cp310-macosx_12_0_universal2.whl", hash = "sha256:3750c5a00bd644c75f4507f77a804d0189d97a107eb1481945a0cf3af3e7a5ac"},
+    {file = "grpcio-1.66.1-cp310-cp310-manylinux_2_17_aarch64.whl", hash = "sha256:a013c5fbb12bfb5f927444b477a26f1080755a931d5d362e6a9a720ca7dbae60"},
+    {file = "grpcio-1.66.1-cp310-cp310-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:b1b24c23d51a1e8790b25514157d43f0a4dce1ac12b3f0b8e9f66a5e2c4c132f"},
+    {file = "grpcio-1.66.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:b7ffb8ea674d68de4cac6f57d2498fef477cef582f1fa849e9f844863af50083"},
+    {file = "grpcio-1.66.1-cp310-cp310-musllinux_1_1_i686.whl", hash = "sha256:307b1d538140f19ccbd3aed7a93d8f71103c5d525f3c96f8616111614b14bf2a"},
+    {file = "grpcio-1.66.1-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:1c17ebcec157cfb8dd445890a03e20caf6209a5bd4ac5b040ae9dbc59eef091d"},
+    {file = "grpcio-1.66.1-cp310-cp310-win32.whl", hash = "sha256:ef82d361ed5849d34cf09105d00b94b6728d289d6b9235513cb2fcc79f7c432c"},
+    {file = "grpcio-1.66.1-cp310-cp310-win_amd64.whl", hash = "sha256:292a846b92cdcd40ecca46e694997dd6b9be6c4c01a94a0dfb3fcb75d20da858"},
+    {file = "grpcio-1.66.1-cp311-cp311-linux_armv7l.whl", hash = "sha256:c30aeceeaff11cd5ddbc348f37c58bcb96da8d5aa93fed78ab329de5f37a0d7a"},
+    {file = "grpcio-1.66.1-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:8a1e224ce6f740dbb6b24c58f885422deebd7eb724aff0671a847f8951857c26"},
+    {file = "grpcio-1.66.1-cp311-cp311-manylinux_2_17_aarch64.whl", hash = "sha256:a66fe4dc35d2330c185cfbb42959f57ad36f257e0cc4557d11d9f0a3f14311df"},
+    {file = "grpcio-1.66.1-cp311-cp311-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:e3ba04659e4fce609de2658fe4dbf7d6ed21987a94460f5f92df7579fd5d0e22"},
+    {file = "grpcio-1.66.1-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:4573608e23f7e091acfbe3e84ac2045680b69751d8d67685ffa193a4429fedb1"},
+    {file = "grpcio-1.66.1-cp311-cp311-musllinux_1_1_i686.whl", hash = "sha256:7e06aa1f764ec8265b19d8f00140b8c4b6ca179a6dc67aa9413867c47e1fb04e"},
+    {file = "grpcio-1.66.1-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:3885f037eb11f1cacc41f207b705f38a44b69478086f40608959bf5ad85826dd"},
+    {file = "grpcio-1.66.1-cp311-cp311-win32.whl", hash = "sha256:97ae7edd3f3f91480e48ede5d3e7d431ad6005bfdbd65c1b56913799ec79e791"},
+    {file = "grpcio-1.66.1-cp311-cp311-win_amd64.whl", hash = "sha256:cfd349de4158d797db2bd82d2020554a121674e98fbe6b15328456b3bf2495bb"},
+    {file = "grpcio-1.66.1-cp312-cp312-linux_armv7l.whl", hash = "sha256:a92c4f58c01c77205df6ff999faa008540475c39b835277fb8883b11cada127a"},
+    {file = "grpcio-1.66.1-cp312-cp312-macosx_10_9_universal2.whl", hash = "sha256:fdb14bad0835914f325349ed34a51940bc2ad965142eb3090081593c6e347be9"},
+    {file = "grpcio-1.66.1-cp312-cp312-manylinux_2_17_aarch64.whl", hash = "sha256:f03a5884c56256e08fd9e262e11b5cfacf1af96e2ce78dc095d2c41ccae2c80d"},
+    {file = "grpcio-1.66.1-cp312-cp312-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:2ca2559692d8e7e245d456877a85ee41525f3ed425aa97eb7a70fc9a79df91a0"},
+    {file = "grpcio-1.66.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:84ca1be089fb4446490dd1135828bd42a7c7f8421e74fa581611f7afdf7ab761"},
+    {file = "grpcio-1.66.1-cp312-cp312-musllinux_1_1_i686.whl", hash = "sha256:d639c939ad7c440c7b2819a28d559179a4508783f7e5b991166f8d7a34b52815"},
+    {file = "grpcio-1.66.1-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:b9feb4e5ec8dc2d15709f4d5fc367794d69277f5d680baf1910fc9915c633524"},
+    {file = "grpcio-1.66.1-cp312-cp312-win32.whl", hash = "sha256:7101db1bd4cd9b880294dec41a93fcdce465bdbb602cd8dc5bd2d6362b618759"},
+    {file = "grpcio-1.66.1-cp312-cp312-win_amd64.whl", hash = "sha256:b0aa03d240b5539648d996cc60438f128c7f46050989e35b25f5c18286c86734"},
+    {file = "grpcio-1.66.1-cp38-cp38-linux_armv7l.whl", hash = "sha256:ecfe735e7a59e5a98208447293ff8580e9db1e890e232b8b292dc8bd15afc0d2"},
+    {file = "grpcio-1.66.1-cp38-cp38-macosx_10_9_universal2.whl", hash = "sha256:4825a3aa5648010842e1c9d35a082187746aa0cdbf1b7a2a930595a94fb10fce"},
+    {file = "grpcio-1.66.1-cp38-cp38-manylinux_2_17_aarch64.whl", hash = "sha256:f517fd7259fe823ef3bd21e508b653d5492e706e9f0ef82c16ce3347a8a5620c"},
+    {file = "grpcio-1.66.1-cp38-cp38-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:f1fe60d0772831d96d263b53d83fb9a3d050a94b0e94b6d004a5ad111faa5b5b"},
+    {file = "grpcio-1.66.1-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:31a049daa428f928f21090403e5d18ea02670e3d5d172581670be006100db9ef"},
+    {file = "grpcio-1.66.1-cp38-cp38-musllinux_1_1_i686.whl", hash = "sha256:6f914386e52cbdeb5d2a7ce3bf1fdfacbe9d818dd81b6099a05b741aaf3848bb"},
+    {file = "grpcio-1.66.1-cp38-cp38-musllinux_1_1_x86_64.whl", hash = "sha256:bff2096bdba686019fb32d2dde45b95981f0d1490e054400f70fc9a8af34b49d"},
+    {file = "grpcio-1.66.1-cp38-cp38-win32.whl", hash = "sha256:aa8ba945c96e73de29d25331b26f3e416e0c0f621e984a3ebdb2d0d0b596a3b3"},
+    {file = "grpcio-1.66.1-cp38-cp38-win_amd64.whl", hash = "sha256:161d5c535c2bdf61b95080e7f0f017a1dfcb812bf54093e71e5562b16225b4ce"},
+    {file = "grpcio-1.66.1-cp39-cp39-linux_armv7l.whl", hash = "sha256:d0cd7050397b3609ea51727b1811e663ffda8bda39c6a5bb69525ef12414b503"},
+    {file = "grpcio-1.66.1-cp39-cp39-macosx_10_9_universal2.whl", hash = "sha256:0e6c9b42ded5d02b6b1fea3a25f036a2236eeb75d0579bfd43c0018c88bf0a3e"},
+    {file = "grpcio-1.66.1-cp39-cp39-manylinux_2_17_aarch64.whl", hash = "sha256:c9f80f9fad93a8cf71c7f161778ba47fd730d13a343a46258065c4deb4b550c0"},
+    {file = "grpcio-1.66.1-cp39-cp39-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:5dd67ed9da78e5121efc5c510f0122a972216808d6de70953a740560c572eb44"},
+    {file = "grpcio-1.66.1-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:48b0d92d45ce3be2084b92fb5bae2f64c208fea8ceed7fccf6a7b524d3c4942e"},
+    {file = "grpcio-1.66.1-cp39-cp39-musllinux_1_1_i686.whl", hash = "sha256:4d813316d1a752be6f5c4360c49f55b06d4fe212d7df03253dfdae90c8a402bb"},
+    {file = "grpcio-1.66.1-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:9c9bebc6627873ec27a70fc800f6083a13c70b23a5564788754b9ee52c5aef6c"},
+    {file = "grpcio-1.66.1-cp39-cp39-win32.whl", hash = "sha256:30a1c2cf9390c894c90bbc70147f2372130ad189cffef161f0432d0157973f45"},
+    {file = "grpcio-1.66.1-cp39-cp39-win_amd64.whl", hash = "sha256:17663598aadbedc3cacd7bbde432f541c8e07d2496564e22b214b22c7523dac8"},
+    {file = "grpcio-1.66.1.tar.gz", hash = "sha256:35334f9c9745add3e357e3372756fd32d925bd52c41da97f4dfdafbde0bf0ee2"},
+]
+
+[package.extras]
+protobuf = ["grpcio-tools (>=1.66.1)"]
+
+[package.source]
+type = "legacy"
+url = "http://192.168.1.2:5001/index"
+reference = "threadripper"
+
 [[package]]
 name = "gunicorn"
 version = "21.2.0"
@@ -789,6 +868,26 @@ type = "legacy"
 url = "http://192.168.1.2:5001/index"
 reference = "threadripper"
 
+[[package]]
+name = "markdown"
+version = "3.7"
+description = "Python implementation of John Gruber's Markdown."
+optional = false
+python-versions = ">=3.8"
+files = [
+    {file = "Markdown-3.7-py3-none-any.whl", hash = "sha256:7eb6df5690b81a1d7942992c97fad2938e956e79df20cbc6186e9c3a77b1c803"},
+    {file = "markdown-3.7.tar.gz", hash = "sha256:2ae2471477cfd02dbbf038d5d9bc226d40def84b4fe2986e49b59b6b472bbed2"},
+]
+
+[package.extras]
+docs = ["mdx-gh-links (>=0.2)", "mkdocs (>=1.5)", "mkdocs-gen-files", "mkdocs-literate-nav", "mkdocs-nature (>=0.6)", "mkdocs-section-index", "mkdocstrings[python]"]
+testing = ["coverage", "pyyaml"]
+
+[package.source]
+type = "legacy"
+url = "http://192.168.1.2:5001/index"
+reference = "threadripper"
+
 [[package]]
 name = "markupsafe"
 version = "2.1.5"
@@ -1256,6 +1355,7 @@ description = "Nvidia JIT LTO Library"
 optional = false
 python-versions = ">=3"
 files = [
+    {file = "nvidia_nvjitlink_cu12-12.5.40-py3-none-manylinux2014_aarch64.whl", hash = "sha256:004186d5ea6a57758fd6d57052a123c73a4815adf365eb8dd6a85c9eaa7535ff"},
     {file = "nvidia_nvjitlink_cu12-12.5.40-py3-none-manylinux2014_x86_64.whl", hash = "sha256:d9714f27c1d0f0895cd8915c07a87a1d0029a0aa36acaf9156952ec2a8a12189"},
     {file = "nvidia_nvjitlink_cu12-12.5.40-py3-none-win_amd64.whl", hash = "sha256:c3401dc8543b52d3a8158007a0c1ab4e9c768fcbd24153a48c86972102197ddd"},
 ]
@@ -1521,6 +1621,31 @@ type = "legacy"
 url = "http://192.168.1.2:5001/index"
 reference = "threadripper"
 
+[[package]]
+name = "protobuf"
+version = "5.28.0"
+description = ""
+optional = false
+python-versions = ">=3.8"
+files = [
+    {file = "protobuf-5.28.0-cp310-abi3-win32.whl", hash = "sha256:66c3edeedb774a3508ae70d87b3a19786445fe9a068dd3585e0cefa8a77b83d0"},
+    {file = "protobuf-5.28.0-cp310-abi3-win_amd64.whl", hash = "sha256:6d7cc9e60f976cf3e873acb9a40fed04afb5d224608ed5c1a105db4a3f09c5b6"},
+    {file = "protobuf-5.28.0-cp38-abi3-macosx_10_9_universal2.whl", hash = "sha256:532627e8fdd825cf8767a2d2b94d77e874d5ddb0adefb04b237f7cc296748681"},
+    {file = "protobuf-5.28.0-cp38-abi3-manylinux2014_aarch64.whl", hash = "sha256:018db9056b9d75eb93d12a9d35120f97a84d9a919bcab11ed56ad2d399d6e8dd"},
+    {file = "protobuf-5.28.0-cp38-abi3-manylinux2014_x86_64.whl", hash = "sha256:6206afcb2d90181ae8722798dcb56dc76675ab67458ac24c0dd7d75d632ac9bd"},
+    {file = "protobuf-5.28.0-cp38-cp38-win32.whl", hash = "sha256:eef7a8a2f4318e2cb2dee8666d26e58eaf437c14788f3a2911d0c3da40405ae8"},
+    {file = "protobuf-5.28.0-cp38-cp38-win_amd64.whl", hash = "sha256:d001a73c8bc2bf5b5c1360d59dd7573744e163b3607fa92788b7f3d5fefbd9a5"},
+    {file = "protobuf-5.28.0-cp39-cp39-win32.whl", hash = "sha256:dde9fcaa24e7a9654f4baf2a55250b13a5ea701493d904c54069776b99a8216b"},
+    {file = "protobuf-5.28.0-cp39-cp39-win_amd64.whl", hash = "sha256:853db610214e77ee817ecf0514e0d1d052dff7f63a0c157aa6eabae98db8a8de"},
+    {file = "protobuf-5.28.0-py3-none-any.whl", hash = "sha256:510ed78cd0980f6d3218099e874714cdf0d8a95582e7b059b06cabad855ed0a0"},
+    {file = "protobuf-5.28.0.tar.gz", hash = "sha256:dde74af0fa774fa98892209992295adbfb91da3fa98c8f67a88afe8f5a349add"},
+]
+
+[package.source]
+type = "legacy"
+url = "http://192.168.1.2:5001/index"
+reference = "threadripper"
+
 [[package]]
 name = "py"
 version = "1.11.0"
@@ -1892,6 +2017,26 @@ type = "legacy"
 url = "http://192.168.1.2:5001/index"
 reference = "threadripper"
 
+[[package]]
+name = "pytorch-ignite"
+version = "0.5.1"
+description = "A lightweight library to help with training neural networks in PyTorch."
+optional = false
+python-versions = "*"
+files = [
+    {file = "pytorch_ignite-0.5.1-py3-none-any.whl", hash = "sha256:08c289f59c4185448cf0d97c0ef35f409f640fe1e7355f8fa13e504dcbe47b6d"},
+    {file = "pytorch_ignite-0.5.1.tar.gz", hash = "sha256:f41050bb0f85da6a4e5483cbc584ddf300382f537cd2f31cbb05ad62b8a6682e"},
+]
+
+[package.dependencies]
+packaging = "*"
+torch = ">=1.3,<3"
+
+[package.source]
+type = "legacy"
+url = "http://192.168.1.2:5001/index"
+reference = "threadripper"
+
 [[package]]
 name = "pytz"
 version = "2024.1"
@@ -2103,6 +2248,50 @@ type = "legacy"
 url = "http://192.168.1.2:5001/index"
 reference = "threadripper"
 
+[[package]]
+name = "tensorboard"
+version = "2.17.1"
+description = "TensorBoard lets you watch Tensors Flow"
+optional = false
+python-versions = ">=3.9"
+files = [
+    {file = "tensorboard-2.17.1-py3-none-any.whl", hash = "sha256:253701a224000eeca01eee6f7e978aea7b408f60b91eb0babdb04e78947b773e"},
+]
+
+[package.dependencies]
+absl-py = ">=0.4"
+grpcio = ">=1.48.2"
+markdown = ">=2.6.8"
+numpy = ">=1.12.0"
+packaging = "*"
+protobuf = ">=3.19.6,<4.24.0 || >4.24.0"
+setuptools = ">=41.0.0"
+six = ">1.9"
+tensorboard-data-server = ">=0.7.0,<0.8.0"
+werkzeug = ">=1.0.1"
+
+[package.source]
+type = "legacy"
+url = "http://192.168.1.2:5001/index"
+reference = "threadripper"
+
+[[package]]
+name = "tensorboard-data-server"
+version = "0.7.2"
+description = "Fast data loading for TensorBoard"
+optional = false
+python-versions = ">=3.7"
+files = [
+    {file = "tensorboard_data_server-0.7.2-py3-none-any.whl", hash = "sha256:7e0610d205889588983836ec05dc098e80f97b7e7bbff7e994ebb78f578d0ddb"},
+    {file = "tensorboard_data_server-0.7.2-py3-none-macosx_10_9_x86_64.whl", hash = "sha256:9fe5d24221b29625dbc7328b0436ca7fc1c23de4acf4d272f1180856e32f9f60"},
+    {file = "tensorboard_data_server-0.7.2-py3-none-manylinux_2_31_x86_64.whl", hash = "sha256:ef687163c24185ae9754ed5650eb5bc4d84ff257aabdc33f0cc6f74d8ba54530"},
+]
+
+[package.source]
+type = "legacy"
+url = "http://192.168.1.2:5001/index"
+reference = "threadripper"
+
 [[package]]
 name = "torch"
 version = "2.2.2"
@@ -2228,6 +2417,31 @@ type = "legacy"
 url = "http://192.168.1.2:5001/index"
 reference = "threadripper"
 
+[[package]]
+name = "tqdm"
+version = "4.66.4"
+description = "Fast, Extensible Progress Meter"
+optional = false
+python-versions = ">=3.7"
+files = [
+    {file = "tqdm-4.66.4-py3-none-any.whl", hash = "sha256:b75ca56b413b030bc3f00af51fd2c1a1a5eac6a0c1cca83cbb37a5c52abce644"},
+    {file = "tqdm-4.66.4.tar.gz", hash = "sha256:e4d936c9de8727928f3be6079590e97d9abfe8d39a590be678eb5919ffc186bb"},
+]
+
+[package.dependencies]
+colorama = {version = "*", markers = "platform_system == \"Windows\""}
+
+[package.extras]
+dev = ["pytest (>=6)", "pytest-cov", "pytest-timeout", "pytest-xdist"]
+notebook = ["ipywidgets (>=6)"]
+slack = ["slack-sdk"]
+telegram = ["requests"]
+
+[package.source]
+type = "legacy"
+url = "http://192.168.1.2:5001/index"
+reference = "threadripper"
+
 [[package]]
 name = "trio"
 version = "0.25.1"
@@ -2339,6 +2553,22 @@ type = "legacy"
 url = "http://192.168.1.2:5001/index"
 reference = "threadripper"
 
+[[package]]
+name = "types-tqdm"
+version = "4.66.0.20240417"
+description = "Typing stubs for tqdm"
+optional = false
+python-versions = ">=3.8"
+files = [
+    {file = "types-tqdm-4.66.0.20240417.tar.gz", hash = "sha256:16dce9ef522ea8d40e4f5b8d84dd8a1166eefc13ceee7a7e158bf0f1a1421a31"},
+    {file = "types_tqdm-4.66.0.20240417-py3-none-any.whl", hash = "sha256:248aef1f9986b7b8c2c12b3cb4399fc17dba0a29e7e3f3f9cd704babb879383d"},
+]
+
+[package.source]
+type = "legacy"
+url = "http://192.168.1.2:5001/index"
+reference = "threadripper"
+
 [[package]]
 name = "typing-extensions"
 version = "4.12.1"
@@ -2481,4 +2711,4 @@ reference = "threadripper"
 [metadata]
 lock-version = "2.0"
 python-versions = "^3.12"
-content-hash = "adcc6b2cdc33ee2aa56b419d6123e3d1553b08641cd81c22b187ecd0096812b1"
+content-hash = "e311eaec3058b444c5bb980a3700ca9de273d7948ace5d5b94363f4f967b622a"

pyproject.toml

@@ -20,6 +20,7 @@ types-requests = "^2.32.0.20240602"
 types-retry = "^0.9.9.4"
 flake8-pyproject = "^1.2.3"
 pandas-stubs = "^2.2.2.240603"
+types-tqdm = "^4.66.0.20240417"
 
 
 [tool.poetry.group.dev.dependencies]
@@ -34,6 +35,9 @@ torch = "^2.2.1"
 torchvision = "^0.17.1"
 torchinfo = "^1.8.0"
 minio = "^7.2.7"
+tqdm = "^4.66.4"
+pytorch-ignite = "^0.5.1"
+tensorboard = "^2.17.1"
 
 
 [tool.poetry.group.shared.dependencies]
@@ -59,8 +63,13 @@ priority = "primary"
 requires = ["poetry-core"]
 build-backend = "poetry.core.masonry.api"
 
+[tool.isort]
+profile = "black"
+
 [tool.flake8]
 per-file-ignores = "__init__.py:F401"
+max-line-length = 88
+extend-ignore = "E203"
 
 [tool.mypy]
 exclude = "image_download"

shared/data_store/__init__.py

@@ -1,6 +1,8 @@
-from __future__ import annotations
+from .connect_minio import connect_minio
-
-from .connect import connect
 from .delete import delete
 from .get import get
+from .get_image import get_image
+from .get_model import get_model
 from .put import put
+from .put_image import put_image
+from .put_model import put_model

shared/data_store/connect.py

@@ -1,45 +0,0 @@
-"""Definition of connect function."""
-from __future__ import annotations
-
-import os
-
-from minio import Minio
-
-
-def connect(
-) -> Minio:
-    """Connect to MinIO server"""
-    # ensure necessary env vars available
-    necesasary_var_list = [
-        'MINIO_ENDPOINT',
-        'MINIO_ACCESS_KEY',
-        'MINIO_SECRET_KEY',
-        'MINIO_BUCKET_NAME',
-    ]
-    for env_var in necesasary_var_list:
-        # ensure env var set
-        assert (
-            env_var in os.environ
-        ), (
-            f"environment variable not set: {env_var}"
-        )
-    # prepare arguments
-    minio_endpoint = os.getenv('MINIO_ENDPOINT', default=None)
-    assert isinstance(minio_endpoint, str)
-    minio_access_key = os.getenv('MINIO_ACCESS_KEY', default=None)
-    assert isinstance(minio_access_key, str)
-    minio_secret_key = os.getenv('MINIO_SECRET_KEY', default=None)
-    assert isinstance(minio_secret_key, str)
-    minio_bucket_name = os.getenv('MINIO_BUCKET_NAME', default=None)
-    assert isinstance(minio_bucket_name, str)
-    # connect client
-    client = Minio(
-        endpoint=minio_endpoint,
-        access_key=minio_access_key,
-        secret_key=minio_secret_key,
-        secure=False,
-    )
-    # ensure bucket exists
-    if not client.bucket_exists(bucket_name=minio_bucket_name):
-        client.make_bucket(bucket_name=minio_bucket_name)
-    return client

shared/data_store/connect_minio.py

@@ -0,0 +1,42 @@
+"""Definition of connect function."""
+
+import logging
+import os
+
+from minio import Minio
+
+
+def connect_minio() -> Minio:
+    """Connect to MinIO server."""
+    # ensure necessary env vars available
+    env_var_list = [
+        'MINIO_ENDPOINT',
+        'MINIO_ACCESS_KEY',
+        'MINIO_SECRET_KEY',
+        'MINIO_BUCKET_NAME',
+    ]
+    for env_var in env_var_list:
+        # ensure env var set
+        assert (
+            env_var in os.environ
+        ), f"environment variable not set: {
+            env_var
+        }"
+    # prepare arguments
+    minio_endpoint = os.getenv('MINIO_ENDPOINT', default='')
+    minio_access_key = os.getenv('MINIO_ACCESS_KEY', default='')
+    minio_secret_key = os.getenv('MINIO_SECRET_KEY', default='')
+    minio_bucket_name = os.getenv('MINIO_BUCKET_NAME', default='')
+    # connect client
+    client = Minio(
+        endpoint=minio_endpoint,
+        access_key=minio_access_key,
+        secret_key=minio_secret_key,
+        secure=False,
+    )
+    # ensure bucket exists
+    if not client.bucket_exists(bucket_name=minio_bucket_name):
+        logging.info('creating bucket: %s', minio_bucket_name)
+        client.make_bucket(bucket_name=minio_bucket_name)
+    logging.debug('finished')
+    return client

shared/data_store/delete.py

@@ -1,4 +1,5 @@
 """Definition of delete function."""
+
 from __future__ import annotations
 
 import logging
@@ -14,8 +15,8 @@ def delete(
     """Delete object from MinIO."""
     assert isinstance(client, Minio)
     assert isinstance(object_name, str)
-    bucket_name = os.getenv('MINIO_BUCKET_NAME', default=None)
+    bucket_name = os.getenv('MINIO_BUCKET_NAME', default='')
-    assert isinstance(bucket_name, str)
+    assert len(bucket_name) > 0
     # remove object
     try:
         client.remove_object(

shared/data_store/get.py

@@ -1,8 +1,9 @@
 """Definition of get function."""
-from __future__ import annotations
 
+import logging
 import os
 from io import BytesIO
+from traceback import print_exc
 
 from minio import Minio
 
@@ -14,14 +15,25 @@ def get(
     """Get buffer from bucket in MinIO."""
     assert isinstance(client, Minio)
     assert isinstance(object_name, str)
-    bucket_name = os.getenv('MINIO_BUCKET_NAME', default=None)
+    bucket_name = os.getenv('MINIO_BUCKET_NAME', default='')
-    assert isinstance(bucket_name, str)
     # get buffer
     try:
-        response = client.get_object(bucket_name, object_name)
+        response = client.get_object(
-        buffer = BytesIO(response.data)
+            bucket_name=bucket_name,
+            object_name=object_name,
+        )
+        assert response.status == 200
+        buffer = BytesIO()
+        chunk_size = 2**14
+        while chunk := response.read(chunk_size):
+            buffer.write(chunk)
+        buffer.seek(0)
+        logging.debug('got %s', object_name)
+        return buffer
+    except Exception as exc:
+        logging.error('failed getting data from MinIO')
+        print_exc()
+        raise exc
     finally:
         response.close()
         response.release_conn()
-    buffer.seek(0)
-    return buffer

shared/data_store/get_image.py

@@ -0,0 +1,41 @@
+"""Definition of get_image function."""
+
+import logging
+import os
+from io import BytesIO
+
+from minio import Minio
+from PIL import Image
+
+
+def get_image(
+    client: Minio,
+    object_name: str,
+) -> Image.Image:
+    """Get image from image subfolder in bucket in Minio."""
+    assert isinstance(client, Minio)
+    assert isinstance(object_name, str)
+    assert len(object_name) > 0
+    # prepare arguments
+    assert 'MINIO_BUCKET_NAME' in os.environ
+    bucket_name = os.getenv('MINIO_BUCKET_NAME', default='')
+    subfolder = 'images'
+    # get object from bucket
+    object_name = f'{subfolder}/{object_name}'
+    try:
+        response = client.get_object(
+            bucket_name=bucket_name,
+            object_name=object_name,
+        )
+        buffer = BytesIO(response.data)
+    except Exception as exc:
+        logging.error('failed getting data from MinIO')
+        raise exc
+    finally:
+        response.close()
+        response.release_conn()
+    buffer.seek(0)
+    # convert data to image
+    image = Image.open(buffer)
+    logging.debug('got data from %s', object_name)
+    return image

shared/data_store/get_model.py

@@ -0,0 +1,30 @@
+"""Definition of get_model function."""
+
+import logging
+from collections import OrderedDict
+
+import torch
+from minio import Minio
+
+from .get import get
+
+
+def get_model(
+    client: Minio,
+    object_name: str,
+) -> OrderedDict:
+    """Get model from model subfolder in bucket in Minio."""
+    assert isinstance(client, Minio)
+    assert isinstance(object_name, str)
+    subfolder = 'models'
+    object_name = f'{subfolder}/{object_name}'
+    # get buffer
+    buffer = get(
+        client=client,
+        object_name=object_name,
+    )
+    # convert data to model checkpoint
+    buffer.seek(0)
+    model_content = torch.load(buffer)
+    logging.debug('finished')
+    return model_content

shared/data_store/put.py

@@ -1,4 +1,5 @@
 """Definition of put function."""
+
 from __future__ import annotations
 
 import logging
@@ -16,8 +17,8 @@ def put(
     """Put buffer in bucket in MinIO and return MD5 checksum as object name."""
     assert isinstance(client, Minio)
     assert isinstance(buffer, BytesIO)
-    bucket_name = os.getenv('MINIO_BUCKET_NAME', default=None)
+    bucket_name = os.getenv('MINIO_BUCKET_NAME', default='')
-    assert isinstance(bucket_name, str)
+    assert len(bucket_name) > 0
     # get md5 of image
     checksum = md5(buffer.getbuffer()).hexdigest()
     # prepare for saving

shared/data_store/put_image.py

@@ -0,0 +1,44 @@
+"""Definition of put_image function."""
+
+import logging
+import os
+from hashlib import md5
+from io import BytesIO
+
+from minio import Minio
+from PIL import Image
+
+
+def put_image(
+    client: Minio,
+    image: Image.Image,
+) -> str:
+    """Put image in image subfolder in bucket in Minio and return MD5 checksum
+    used as object name."""
+    assert isinstance(client, Minio)
+    assert isinstance(image, Image.Image)
+    bucket_name = os.getenv('MINIO_BUCKET_NAME', default='')
+    assert len(bucket_name) > 0
+    subfolder = 'images'
+    # save image to buffer
+    buffer = BytesIO()
+    image.save(buffer, 'png')
+    # get md5 of image
+    checksum = md5(buffer.getbuffer()).hexdigest()
+    # prepare for saving
+    num_bytes = buffer.tell()
+    buffer.seek(0)
+    # send data to bucket
+    object_name = f'{subfolder}/{checksum}'
+    try:
+        client.put_object(
+            bucket_name=bucket_name,
+            object_name=object_name,
+            length=num_bytes,
+            data=buffer,
+        )
+    except Exception as exc:
+        logging.error('failed saving data to MinIO')
+        raise exc
+    logging.debug('saved data to %s', object_name)
+    return checksum

shared/data_store/put_model.py

@@ -0,0 +1,45 @@
+"""Definition of put_model function."""
+
+import logging
+import os
+from hashlib import md5
+from io import BytesIO
+
+import torch
+from minio import Minio
+from torch.nn import Module
+
+
+def put_model(
+    client: Minio,
+    model: Module,
+) -> str:
+    """Put model in model subfolder in bucket in Minio and return MD5 checksum
+    used as object name."""
+    assert isinstance(client, Minio)
+    assert isinstance(model, Module)
+    bucket_name = os.getenv('MINIO_BUCKET_NAME', default='')
+    assert len(bucket_name) > 0
+    subfolder = 'models'
+    # save image to buffer
+    buffer = BytesIO()
+    torch.save(model.state_dict(), buffer)
+    # get md5 of image
+    checksum = md5(buffer.getbuffer()).hexdigest()
+    # prepare for saving
+    num_bytes = buffer.tell()
+    buffer.seek(0)
+    # send data to bucket
+    object_name = f'{subfolder}/{checksum}'
+    try:
+        client.put_object(
+            bucket_name=bucket_name,
+            object_name=object_name,
+            length=num_bytes,
+            data=buffer,
+        )
+    except Exception as exc:
+        logging.error('failed saving data to MinIO')
+        raise exc
+    logging.debug('saved data to %s', object_name)
+    return checksum

shared/database/__init__.py

@@ -1,10 +1,11 @@
 """Database module content."""
+
 from __future__ import annotations
 
 from .classes.dataset import Dataset
 from .classes.exceptions import NoDocumentFoundException
 from .classes.visual_communication import VisualCommunication
-from .utils.connect import connect
+from .utils.connect_mongodb import connect_mongodb
 from .utils.count_documents import count_documents
 from .utils.get_visual_communication import get_visual_communication
 from .utils.list_names import list_names

shared/database/utils/__init__.py

@@ -1,4 +1,3 @@
 """Database utils module content."""
-from __future__ import annotations
 
-from .connect import connect
+from .connect_mongodb import connect_mongodb

shared/database/utils/connect_mongodb.py

@@ -1,8 +1,5 @@
-"""
+"""Definition of function to connect to database using environment
-Definition of function to connect to database
+variables."""
-using environment variables.
-"""
-from __future__ import annotations
 
 import logging
 import os
@@ -11,7 +8,7 @@ from dotenv import load_dotenv
 from pymongo import MongoClient
 
 
-def connect():
+def connect_mongodb():
     """Connect to MongoDB using env vars."""
     # load env vars
     load_dotenv()

shared/database/utils/get_dataset.py

@@ -1,4 +1,4 @@
-from __future__ import annotations
+"""Definition of get_dataset function."""
 
 from typing import Literal
 
@@ -10,4 +10,7 @@ def get_dataset(
     type: Literal['train', 'test', 'validation'],
 ) -> list[str]:
     """Get list of data names for the corresponding type."""
+    assert isinstance(collection, Collection)
+    assert isinstance(type, str)
+    assert type in ['train', 'test', 'validation']
     return []

shared/database/utils/save_dataset.py

@@ -20,10 +20,12 @@ def save_dataset(
 
 if __name__ == '__main__':
     from dotenv import load_dotenv
+
     load_dotenv('local.env')
-    from shared.database import list_names, connect
+    from shared.database import connect_mongodb, list_names
+
     # connect to database
-    collection, db, client = connect()
+    collection, db, client = connect_mongodb()
     print(client.server_info())
 
     name_list = list_names(collection=collection, only_with_annotation=True)

shared/dto/data_model.py

@@ -1,10 +1,9 @@
 """Definition of DataModel base class."""
-from __future__ import annotations
 
 import random
 
-from pydantic import BaseModel
+from pydantic import BaseModel, ValidationError
-from pydantic import ValidationError
+from torch import Tensor
 
 
 class DataModel(BaseModel):
@@ -33,26 +32,27 @@ class DataModel(BaseModel):
             raise ValidationError()
         assert isinstance(option, str), 'option is not a string'
         allowed_options_list = cls.list_fields()
-        assert option in allowed_options_list, \
+        assert (
-            f"{option} is not among allowed fields {allowed_options_list}"
+            option in allowed_options_list
+        ), f"{option} is not among allowed fields {allowed_options_list}"
         kwargs = {field: 0 for field in cls.list_fields()}
         kwargs[option] = 1
         return cls(**kwargs)
 
     @classmethod
-    def from_list(cls, data_list: list[float]):
+    def from_tensor(cls, tensor: Tensor):
         """Instantiate from list of values."""
-        kwargs = {key: val for key, val in zip(cls.list_fields(), data_list)}
+        assert tensor.size(dim=0) == 1, f'tensor batch larger than 1: {tensor}'
+        data_list = [float(t.item()) for t in tensor[0]]
+        kwargs = dict(zip(cls.list_fields(), data_list))
         return cls(**kwargs)
 
     def __repr__(self) -> str:
         model_dict = self.model_dump()
-        model_repr_str = f"{self.classname()}("
+        model_repr_str = f'{self.classname()}('
-        model_repr_str += ', '.join([
+        model_repr_str += ', '.join(
-            f"{field}={value:.3f}"
+            [f'{field}={value:.3f}' for field, value in model_dict.items()],
-            for field, value
+        )
-            in model_dict.items()
-        ])
         model_repr_str += ')'
         return model_repr_str
 

shared/dto/model_data.py

@@ -1,4 +1,5 @@
 """Definition of ModelData data model."""
+
 from __future__ import annotations
 
 from .angle import AngleData
@@ -17,6 +18,7 @@ from .visual_syntax import VisualSyntaxData
 
 class ModelData(DataModel):
     """ModelData model for data IO with combined ML model."""
+
     visual_syntax: VisualSyntaxData
     contact: ContactData
     angle: AngleData
@@ -34,8 +36,50 @@ class ModelData(DataModel):
         """Instantiate with random numbers."""
         kwargs = {
             field: field_info.annotation.from_random()  # type: ignore
-            for field, field_info
+            for field, field_info in cls.model_fields.items()
-            in cls.model_fields.items()
+        }
+        return cls(**kwargs)
+
+    @classmethod
+    def from_prediction_dict(
+        cls,
+        prediction_dict: dict,
+    ) -> ModelData:
+        """Instantiate from prediction dictionary."""
+        kwargs = {
+            'visual_syntax': VisualSyntaxData.from_tensor(
+                prediction_dict['visual_syntax'],
+            ),
+            'contact': ContactData.from_tensor(
+                prediction_dict['contact'],
+            ),
+            'angle': AngleData.from_tensor(
+                prediction_dict['angle'],
+            ),
+            'point_of_view': PointOfViewData.from_tensor(
+                prediction_dict['point_of_view'],
+            ),
+            'distance': DistanceData.from_tensor(
+                prediction_dict['distance'],
+            ),
+            'modality_lighting': ModalityLightingData.from_tensor(
+                prediction_dict['modality_lighting'],
+            ),
+            'modality_color': ModalityColorData.from_tensor(
+                prediction_dict['modality_color'],
+            ),
+            'modality_depth': ModalityDepthData.from_tensor(
+                prediction_dict['modality_depth'],
+            ),
+            'information_value': InformationValueData.from_tensor(
+                prediction_dict['information_value'],
+            ),
+            'framing': FramingData.from_tensor(
+                prediction_dict['framing'],
+            ),
+            'salience': SalienceData.from_tensor(
+                prediction_dict['salience'],
+            ),
         }
         return cls(**kwargs)
 
@@ -56,27 +100,16 @@ class ModelData(DataModel):
     ) -> ModelData:
         """Instantiate from annotation."""
         kwargs = {
-            'visual_syntax': VisualSyntaxData
+            'visual_syntax': VisualSyntaxData.from_choice(visual_syntax),
-            .from_choice(visual_syntax),
+            'contact': ContactData.from_choice(contact),
-            'contact': ContactData
+            'angle': AngleData.from_choice(angle),
-            .from_choice(contact),
+            'point_of_view': PointOfViewData.from_choice(point_of_view),
-            'angle': AngleData
+            'distance': DistanceData.from_choice(distance),
-            .from_choice(angle),
+            'modality_lighting': ModalityLightingData.from_choice(modality_lighting),
-            'point_of_view': PointOfViewData
+            'modality_color': ModalityColorData.from_choice(modality_color),
-            .from_choice(point_of_view),
+            'modality_depth': ModalityDepthData.from_choice(modality_depth),
-            'distance': DistanceData
+            'information_value': InformationValueData.from_choice(information_value),
-            .from_choice(distance),
+            'framing': FramingData.from_choice(framing),
-            'modality_lighting': ModalityLightingData
+            'salience': SalienceData.from_choice(salience),
-            .from_choice(modality_lighting),
-            'modality_color': ModalityColorData
-            .from_choice(modality_color),
-            'modality_depth': ModalityDepthData
-            .from_choice(modality_depth),
-            'information_value': InformationValueData
-            .from_choice(information_value),
-            'framing': FramingData
-            .from_choice(framing),
-            'salience': SalienceData
-            .from_choice(salience),
         }
         return cls(**kwargs)

shared/utils/check_env.py

@@ -1,4 +1,5 @@
 """Definition of check_env function."""
+
 from __future__ import annotations
 
 import os
@@ -18,11 +19,8 @@ def check_env() -> None:
         'MINIO_ACCESS_KEY',
         'MINIO_SECRET_KEY',
         'MINIO_BUCKET_NAME',
+        'MINIO_BUCKET_NAME_MODELS',
     }
     for env_var in necesasary_var_list:
         # ensure env var set
-        assert (
+        assert env_var in os.environ, f"environment variable not set: {env_var}"
-            env_var in os.environ
-        ), (
-            f"environment variable not set: {env_var}"
-        )

tests/generate_random_prediction_test.py

@@ -4,11 +4,10 @@ from pathlib import Path
 
 from dotenv import load_dotenv
 
-from shared.data_store import connect as connect_minio
+from shared.data_store import connect_minio
-from shared.database import connect as connect_mongo
+from shared.database import connect_mongodb
 from shared.database.classes import VisualCommunication
-from shared.utils import check_env
+from shared.utils import check_env, setup_logging
-from shared.utils import setup_logging
 
 if __name__ == '__main__':
     # load in env file
@@ -22,7 +21,7 @@ if __name__ == '__main__':
     # connect to minIO
     minio_client = connect_minio()
     # connect to MongoDB
-    collection, db, client = connect_mongo()
+    collection, db, client = connect_mongodb()
     # get list of image paths
     test_dir = Path(__file__).parent
     img_dir = test_dir / 'imgs'
@@ -31,8 +30,7 @@ if __name__ == '__main__':
     # instantiate data object
     vis_com_list = [
         VisualCommunication.from_file(path, minio_client=minio_client)
-        for path
+        for path in img_path_list
-        in img_path_list
     ]
     # generate random predictions
     for vis_com in vis_com_list:

tests/get_visual_communication_test.py

@@ -4,13 +4,9 @@ from pathlib import Path
 
 from dotenv import load_dotenv
 
-from shared.data_store import (
+from shared.data_store import connect_minio
-    connect as connect_minio,
+from shared.database import connect_mongodb, get_visual_communication
-)
+from shared.utils import check_env, setup_logging
-from shared.database import connect
-from shared.database import get_visual_communication
-from shared.utils import check_env
-from shared.utils import setup_logging
 
 if __name__ == '__main__':
     # load in env file
@@ -24,7 +20,7 @@ if __name__ == '__main__':
     # connect to minIO
     minio_client = connect_minio()
     # connect to MongoDB
-    collection, db, client = connect()
+    collection, db, client = connect_mongodb()
     # get visual communication
     vis_com = get_visual_communication(collection)
     print(repr(vis_com))

tests/image_upload_test.py

@@ -5,11 +5,10 @@ from pathlib import Path
 from dotenv import load_dotenv
 from pymongo.errors import DuplicateKeyError
 
-from shared.data_store import connect as connect_minio
+from shared.data_store import connect_minio
-from shared.database import connect as connect_mongo
+from shared.database import connect_mongodb
 from shared.database.classes import VisualCommunication
-from shared.utils import check_env
+from shared.utils import check_env, setup_logging
-from shared.utils import setup_logging
 
 if __name__ == '__main__':
     # load in env file
@@ -23,7 +22,7 @@ if __name__ == '__main__':
     # connect to minIO
     minio_client = connect_minio()
     # connect to MongoDB
-    collection, db, client = connect_mongo()
+    collection, db, client = connect_mongodb()
     # get list of image paths
     test_dir = Path(__file__).parent
     img_dir = test_dir / 'imgs'
@@ -32,8 +31,7 @@ if __name__ == '__main__':
     # instantiate data object
     vis_com_list = [
         VisualCommunication.from_file(path, minio_client=minio_client)
-        for path
+        for path in img_path_list
-        in img_path_list
     ]
     for vis_com in vis_com_list:
         print(repr(vis_com))

tests/image_upload_to_server_test.py

@@ -5,11 +5,9 @@ from pathlib import Path
 from dotenv import load_dotenv
 from pymongo.errors import DuplicateKeyError
 
-from shared.data_store import connect as connect_minio
+from shared.data_store import connect_minio
-from shared.database import connect as connect_mongo
+from shared.database import VisualCommunication, connect_mongodb
-from shared.database import VisualCommunication
+from shared.utils import check_env, setup_logging
-from shared.utils import check_env
-from shared.utils import setup_logging
 
 if __name__ == '__main__':
     # load in env file
@@ -23,21 +21,22 @@ if __name__ == '__main__':
     # connect to minIO
     minio_client = connect_minio()
     # connect to MongoDB
-    collection, db, client = connect_mongo()
+    collection, db, client = connect_mongodb()
     # get list of image paths
     ext_img_dir = Path('/Volumes/BW-PSSD/Mixed Methods/')
     assert ext_img_dir.exists()
     img_path_list = [
-        path for path in ext_img_dir.glob(
+        path
+        for path in ext_img_dir.glob(
             '*.jpg',
-        ) if path.is_file()
+        )
+        if path.is_file()
     ]
     print(f"found {len(img_path_list)} images")
     # create visual communication objects
     vis_com_list = [
         VisualCommunication.from_file(path, minio_client=minio_client)
-        for path
+        for path in img_path_list
-        in img_path_list
     ]
     print(f"created {len(vis_com_list)} visual communication objects")
     # upload images

tests/model_io_test.py

@@ -0,0 +1,34 @@
+"""Test that a model can be saved and loaded again."""
+
+from pathlib import Path
+
+import torch
+from dotenv import load_dotenv
+from torchinfo import summary
+
+from model.src.models import VisualCommunicationModel
+from shared.data_store import connect_minio, put_model
+from shared.utils import setup_logging
+
+DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+if __name__ == '__main__':
+    # load in env file
+    env_path = Path(__file__).parent.parent / 'server.env'
+    assert env_path.exists()
+    load_dotenv(env_path)
+    # setup logging
+    setup_logging()
+    # connect to minio
+    client = connect_minio()
+    # instantiate model
+    model = VisualCommunicationModel().to(DEVICE)
+    # show model weights
+    summary(model)
+    # put buffer in minio bucket
+    hash_str = put_model(
+        client=client,
+        model=model,
+    )
+    print(f"hash string: {hash_str}")
+    print('saved to Minio')

tests/new_model_to_minio_test.py

@@ -0,0 +1,32 @@
+"""Definition of function to generate a new randomly initialized model and save
+it in Minio datastore."""
+
+from pathlib import Path
+
+from dotenv import load_dotenv
+from torchinfo import summary
+
+from model.src.models import VisualCommunicationModel
+from shared.data_store import connect_minio, put_model
+from shared.utils import setup_logging
+
+if __name__ == '__main__':
+    # load in env file
+    env_path = Path(__file__).parent.parent / 'server.env'
+    assert env_path.exists()
+    load_dotenv(env_path)
+    # setup logging
+    setup_logging()
+    # connect to minio
+    client = connect_minio()
+    # instantiate model
+    model = VisualCommunicationModel(download_resnet_weights=True)
+    # show model weights
+    summary(model)
+    # put buffer in minio bucket
+    hash_str = put_model(
+        client=client,
+        model=model,
+    )
+    print(f"hash string: {hash_str}")
+    print('saved to Minio')

tests/prediction_upload_test.py

@@ -4,12 +4,10 @@ from pathlib import Path
 
 from dotenv import load_dotenv
 
-from shared.data_store import connect as connect_minio
+from shared.data_store import connect_minio
-from shared.database import connect as connect_mongo
+from shared.database import connect_mongodb, upsert_prediction
-from shared.database import upsert_prediction
 from shared.database.classes import VisualCommunication
-from shared.utils import check_env
+from shared.utils import check_env, setup_logging
-from shared.utils import setup_logging
 
 if __name__ == '__main__':
     # load in env file
@@ -23,7 +21,7 @@ if __name__ == '__main__':
     # connect to minIO
     minio_client = connect_minio()
     # connect to MongoDB
-    collection, db, client = connect_mongo()
+    collection, db, client = connect_mongodb()
     # get list of image paths
     test_dir = Path(__file__).parent
     img_dir = test_dir / 'imgs'
@@ -31,8 +29,7 @@ if __name__ == '__main__':
     # instantiate data object
     vis_com_list = [
         VisualCommunication.from_file(path, minio_client=minio_client)
-        for path
+        for path in img_path_list
-        in img_path_list
     ]
     # generate random predictions
     for vis_com in vis_com_list:

tests/total_annotated_test.py

@@ -5,8 +5,7 @@ from pathlib import Path
 
 from dotenv import load_dotenv
 
-from shared.database import connect
+from shared.database import connect_mongodb, count_documents
-from shared.database import count_documents
 
 if __name__ == '__main__':
     # prepare env vars
@@ -15,7 +14,7 @@ if __name__ == '__main__':
     load_dotenv(env_path)
     os.environ['MONGO_HOST'] = 'localhost'
     # connect to database
-    collection, db, client = connect()
+    collection, db, client = connect_mongodb()
     # get visual communication
     num_docs = count_documents(
         collection=collection,

tests/total_documents_test.py

@@ -5,8 +5,7 @@ from pathlib import Path
 
 from dotenv import load_dotenv
 
-from shared.database import connect
+from shared.database import connect_mongodb, count_documents
-from shared.database import count_documents
 
 if __name__ == '__main__':
     # prepare env vars
@@ -15,7 +14,7 @@ if __name__ == '__main__':
     load_dotenv(env_path)
     os.environ['MONGO_HOST'] = 'localhost'
     # connect to database
-    collection, db, client = connect()
+    collection, db, client = connect_mongodb()
     # get visual communication
     num_docs = count_documents(
         collection=collection,

web_ui/src/main.py

@@ -1,13 +1,14 @@
 """Definition of web_ui main script."""
+
 from __future__ import annotations
 
 import os
 
+from shared.data_store import connect_minio
+from shared.database import connect_mongodb
+from shared.utils import check_env, setup_logging
+
 from .app import init_app
-from shared.data_store import connect as connect_minio
-from shared.database import connect
-from shared.utils import check_env
-from shared.utils import setup_logging
 
 # ensure env vars set
 check_env()
@@ -16,7 +17,7 @@ check_env()
 setup_logging()
 
 # connect to database
-collection, db, client = connect()
+collection, db, client = connect_mongodb()
 
 # connect to minio
 minio_client = connect_minio()