Compare commits
145
Commits
b109ddb2b8
..
main
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
70492acafe | ||
|
|
06913064c3 | ||
|
|
7190b6c438 | ||
|
|
7ce7a368c7 | ||
|
|
a2b818bbcc | ||
|
|
08062ad87d | ||
|
|
7176dfdf90 | ||
|
|
6399b88ade | ||
|
|
c5bcd09bba | ||
|
|
7d8af1c582 | ||
|
|
a379270add | ||
|
|
0b54ccaadb | ||
|
|
ac8ca830e1 | ||
|
|
49d2c48d31 | ||
|
|
280dbb579e | ||
|
|
ab7e346939 | ||
|
|
ad82ad4591 | ||
|
|
400706e730 | ||
|
|
096ffd56e9 | ||
|
|
790204af6a | ||
|
|
14a4125688 | ||
|
|
62519c371e | ||
|
|
9c8089d091 | ||
|
|
77a3219f39 | ||
|
|
66370a33c2 | ||
|
|
5870c3b0a3 | ||
|
|
e9340f17c4 | ||
|
|
282fee0806 | ||
|
|
33d30941a6 | ||
|
|
cb8d17cbc3 | ||
|
|
df2c2b7e5f | ||
|
|
6756ae4871 | ||
|
|
fa30aae3a8 | ||
|
|
15f3944d27 | ||
|
|
e6b73fe4fc | ||
|
|
1f72a2db8d | ||
|
|
a610bf4d89 | ||
|
|
9af3d9bc54 | ||
|
|
bbbeb403a8 | ||
|
|
9be7ec9e2e | ||
|
|
a6c717a401 | ||
|
|
4db59d2801 | ||
|
|
3a0c1e9067 | ||
|
|
5a8350d0ad | ||
|
|
62b6224de7 | ||
|
|
68bfff87a7 | ||
|
|
ac21c7a111 | ||
|
|
a18b4d304c | ||
|
|
a2f892e35c | ||
|
|
4bcfbbed0a | ||
|
|
31be0d7aac | ||
|
|
d4c9738db4 | ||
|
|
813daa0c90 | ||
|
|
29bf41c7da | ||
|
|
d9bd5be46d | ||
|
|
c2a2c34967 | ||
|
|
9e66981d3e | ||
|
|
6840ac72f1 | ||
|
|
ee304f1a75 | ||
|
|
4827a3f7dd | ||
|
|
d92051f5c4 | ||
|
|
10ce0d9124 | ||
|
|
88a20c4b0b | ||
|
|
090a1cd8bb | ||
|
|
ac6a305da8 | ||
|
|
cbedc395bc | ||
|
|
0c762b55b2 | ||
|
|
99b5820ae6 | ||
|
|
8214bd2669 | ||
|
|
aff0ae8fc6 | ||
|
|
16ad2b80ee | ||
|
|
88698dd735 | ||
|
|
f28bff7e47 | ||
|
|
9ee27b65a8 | ||
|
|
de4f984110 | ||
|
|
3b677e5e97 | ||
|
|
226fb1a206 | ||
|
|
7a1e17c29f | ||
|
|
f52720685e | ||
|
|
a3da7b048d | ||
|
|
f61eb59afb | ||
|
|
a430366bbe | ||
|
|
06d9753c46 | ||
|
|
4e6499cee3 | ||
|
|
7da3482366 | ||
|
|
7b3374d451 | ||
|
|
b17c1745ab | ||
|
|
db4ce6b425 | ||
|
|
aa6a097e4f | ||
|
|
150a213ae2 | ||
|
|
d860b71b1a | ||
|
|
1a3ee96932 | ||
|
|
8e366f26bc | ||
|
|
8a602264e6 | ||
|
|
99d5c88c8d | ||
|
|
c9da7378a1 | ||
|
|
b2773ecf09 | ||
|
|
084a13426a | ||
|
|
499996329f | ||
|
|
ce08ee4ccd | ||
|
|
cbbb044177 | ||
|
|
c630498168 | ||
|
|
f9d23c5bd4 | ||
|
|
3fa79faa9e | ||
|
|
556e26f21d | ||
|
|
01ee4664ff | ||
|
|
169e0531d9 | ||
|
|
3c46275015 | ||
|
|
0b86acd1a9 | ||
|
|
cddfced177 | ||
|
|
53d179663b | ||
|
|
6be837025b | ||
|
|
2854f74db4 | ||
|
|
462f7c79b4 | ||
|
|
0c9537975c | ||
|
|
10bfac73bd | ||
|
|
bdf3e73ba9 | ||
|
|
42400d6f32 | ||
|
|
0676011f98 | ||
|
|
e189b7c1f7 | ||
|
|
dc535167fa | ||
|
|
b722ddcb83 | ||
|
|
b755a81ef9 | ||
|
|
ae18041d6e | ||
|
|
d8e3d077e3 | ||
|
|
b54e2ed541 | ||
|
|
30f37b6ca5 | ||
|
|
860150905d | ||
|
|
66c3f38dba | ||
|
|
1a6cf69346 | ||
|
|
6f06485776 | ||
|
|
c7fd78dcf8 | ||
|
|
7a8eb93c48 | ||
|
|
4ba68d056b | ||
|
|
ba7a7299c2 | ||
|
|
6dba06704c | ||
|
|
20831c9bc3 | ||
|
|
e2338ad710 | ||
|
|
d1ee43e135 | ||
|
|
909f8f220c | ||
|
|
7fa88166ee | ||
|
|
321defb7ff | ||
|
|
00f1c07cfa | ||
|
|
312f525a46 | ||
|
|
78363988f7 |
+29
-12
@@ -9,25 +9,42 @@ jobs:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
uses: actions/checkout@v3
|
||||
uses: actions/checkout@v4
|
||||
- name: Setup Environment
|
||||
uses: https://github.com/actions/setup-python@v3
|
||||
uses: https://github.com/actions/setup-python@v5
|
||||
with:
|
||||
python-verison: "3.12"
|
||||
python-version: "3.12"
|
||||
architecture: "x64"
|
||||
- name: Install Packages
|
||||
- name: Setup poetry
|
||||
env:
|
||||
PIP_INDEX_URL: http://192.168.1.2:5001/index/
|
||||
PIP_TRUSTED_HOST: 192.168.1.2
|
||||
POETRY_VERSION: 2.1.1
|
||||
POETRY_HOME: /opt/poetry
|
||||
POETRY_NO_INTERACTION: 1
|
||||
POETRY_NO_CACHE: 1
|
||||
run: |
|
||||
pip install poetry
|
||||
poetry install
|
||||
curl -sSL https://install.python-poetry.org | python3 -
|
||||
export PATH=$POETRY_HOME/bin:$PATH
|
||||
poetry --version
|
||||
- name: Install Dependencies
|
||||
env:
|
||||
PIP_INDEX_URL: ${{ vars.PIP_INDEX_URL }}
|
||||
PIP_TRUSTED_HOST: ${{ vars.PIP_TRUSTED_HOST }}
|
||||
run: |
|
||||
/opt/poetry/bin/poetry install
|
||||
- name: PEP8 Check
|
||||
run: |
|
||||
poetry run flake8 . --benchmark
|
||||
/opt/poetry/bin/poetry run flake8 . --benchmark
|
||||
- name: Type Check
|
||||
run: |
|
||||
poetry run mypy .
|
||||
- name: Pytest
|
||||
/opt/poetry/bin/poetry run mypy .
|
||||
- name: Pytest & Calculate Coverage
|
||||
env:
|
||||
MINIO_ENDPOINT: ${{ vars.MINIO_ENDPOINT }}
|
||||
MINIO_ACCESS_KEY: ${{ vars.MINIO_ACCESS_KEY }}
|
||||
MINIO_SECRET_KEY: ${{ secrets.MINIO_SECRET_KEY }}
|
||||
MONGO_ENDPOINT: ${{ vars.MONGO_ENDPOINT }}
|
||||
run: |
|
||||
poetry run pytest .
|
||||
/opt/poetry/bin/poetry run coverage run -m pytest .
|
||||
- name: Coverage Report
|
||||
run: |
|
||||
/opt/poetry/bin/poetry run coverage report -m
|
||||
|
||||
@@ -1,15 +1,16 @@
|
||||
repos:
|
||||
- repo: https://github.com/pre-commit/pre-commit-hooks
|
||||
rev: v4.5.0
|
||||
rev: v5.0.0
|
||||
hooks:
|
||||
- id: trailing-whitespace
|
||||
- id: end-of-file-fixer
|
||||
- id: check-yaml
|
||||
- id: check-added-large-files
|
||||
- id: debug-statements
|
||||
- id: double-quote-string-fixer
|
||||
- id: name-tests-test
|
||||
- repo: https://github.com/asottile/setup-cfg-fmt
|
||||
rev: v2.5.0
|
||||
rev: v2.7.0
|
||||
hooks:
|
||||
- id: setup-cfg-fmt
|
||||
- repo: https://github.com/pre-commit/mirrors-isort
|
||||
@@ -23,7 +24,7 @@ repos:
|
||||
hooks:
|
||||
- id: add-trailing-comma
|
||||
- repo: https://github.com/asottile/pyupgrade
|
||||
rev: v3.15.1
|
||||
rev: v3.19.1
|
||||
hooks:
|
||||
- id: pyupgrade
|
||||
args: [--py39-plus]
|
||||
@@ -35,12 +36,12 @@ repos:
|
||||
additional_dependencies:
|
||||
- "pyproject-flake8"
|
||||
- repo: https://github.com/pre-commit/mirrors-mypy
|
||||
rev: v1.8.0
|
||||
rev: v1.15.0
|
||||
hooks:
|
||||
- id: mypy
|
||||
exclude: ^testing/resources/
|
||||
- repo: https://github.com/psf/black
|
||||
rev: 24.4.2
|
||||
rev: 25.1.0
|
||||
hooks:
|
||||
- id: black
|
||||
language_version: python3.12
|
||||
@@ -55,7 +56,7 @@ repos:
|
||||
language: python
|
||||
types: [ python ]
|
||||
- repo: https://github.com/jendrikseipp/vulture
|
||||
rev: 'v2.6'
|
||||
rev: 'v2.14'
|
||||
hooks:
|
||||
- id: vulture
|
||||
entry: vulture . --min-confidence 90 --exclude */.venv/*.py
|
||||
entry: vulture . --min-confidence 90 --exclude */.venv/*.py,*/tests/*.py
|
||||
|
||||
@@ -1,39 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
from dotenv import load_dotenv
|
||||
|
||||
from shared.data_store import connect_minio
|
||||
from shared.database import connect_mongodb
|
||||
from shared.database.classes import VisualCommunication
|
||||
from shared.utils import check_env, 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)
|
||||
# ensure env vars set
|
||||
check_env()
|
||||
# setup logging
|
||||
setup_logging()
|
||||
# connect to minIO
|
||||
minio_client = connect_minio()
|
||||
# connect to MongoDB
|
||||
collection, db, client = connect_mongodb()
|
||||
# get list of image paths
|
||||
test_dir = Path(__file__).parent
|
||||
img_dir = test_dir / 'imgs'
|
||||
img_path_list = [path for path in img_dir.glob('*.jpeg') if path.is_file()]
|
||||
print(img_path_list)
|
||||
# instantiate data object
|
||||
vis_com_list = [
|
||||
VisualCommunication.from_file(path, minio_client=minio_client)
|
||||
for path in img_path_list
|
||||
]
|
||||
# generate random predictions
|
||||
for vis_com in vis_com_list:
|
||||
vis_com.generate_random_prediction()
|
||||
for vis_com in vis_com_list:
|
||||
print(vis_com)
|
||||
@@ -1,28 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
from dotenv import load_dotenv
|
||||
|
||||
from shared.data_store import connect_minio
|
||||
from shared.database import connect_mongodb, get_visual_communication
|
||||
from shared.utils import check_env, 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)
|
||||
# ensure env vars set
|
||||
check_env()
|
||||
# setup logging
|
||||
setup_logging()
|
||||
# connect to minIO
|
||||
minio_client = connect_minio()
|
||||
# connect to MongoDB
|
||||
collection, db, client = connect_mongodb()
|
||||
# get visual communication
|
||||
vis_com = get_visual_communication(collection)
|
||||
print(repr(vis_com))
|
||||
image = vis_com.get_image(minio_client=minio_client)
|
||||
image.show()
|
||||
@@ -5,8 +5,8 @@ from pathlib import Path
|
||||
|
||||
from dotenv import load_dotenv
|
||||
|
||||
from shared.database import connect_mongodb
|
||||
from shared.database.classes import VisualCommunication
|
||||
from shared.docstore import connect_mongodb
|
||||
from shared.docstore.classes import VisualCommunication
|
||||
|
||||
if __name__ == '__main__':
|
||||
# prepare env vars
|
||||
@@ -1,45 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
from dotenv import load_dotenv
|
||||
from pymongo.errors import DuplicateKeyError
|
||||
|
||||
from shared.data_store import connect_minio
|
||||
from shared.database import connect_mongodb
|
||||
from shared.database.classes import VisualCommunication
|
||||
from shared.utils import check_env, 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)
|
||||
# ensure env vars set
|
||||
check_env()
|
||||
# setup logging
|
||||
setup_logging()
|
||||
# connect to minIO
|
||||
minio_client = connect_minio()
|
||||
# connect to MongoDB
|
||||
collection, db, client = connect_mongodb()
|
||||
# get list of image paths
|
||||
test_dir = Path(__file__).parent
|
||||
img_dir = test_dir / 'imgs'
|
||||
img_path_list = [path for path in img_dir.glob('*.jpeg') if path.is_file()]
|
||||
print(img_path_list)
|
||||
# instantiate data object
|
||||
vis_com_list = [
|
||||
VisualCommunication.from_file(path, minio_client=minio_client)
|
||||
for path in img_path_list
|
||||
]
|
||||
for vis_com in vis_com_list:
|
||||
print(repr(vis_com))
|
||||
# upload images
|
||||
for vis_com in vis_com_list:
|
||||
try:
|
||||
result = collection.insert_one(vis_com.model_dump())
|
||||
except DuplicateKeyError as exc:
|
||||
print('ignoring:\n', exc)
|
||||
else:
|
||||
print(f"inserted document: {result}")
|
||||
@@ -1,49 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
from dotenv import load_dotenv
|
||||
from pymongo.errors import DuplicateKeyError
|
||||
|
||||
from shared.data_store import connect_minio
|
||||
from shared.database import VisualCommunication, connect_mongodb
|
||||
from shared.utils import check_env, 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)
|
||||
# ensure env vars set
|
||||
check_env()
|
||||
# setup logging
|
||||
setup_logging()
|
||||
# connect to minIO
|
||||
minio_client = connect_minio()
|
||||
# connect to MongoDB
|
||||
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(
|
||||
'*.jpg',
|
||||
)
|
||||
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 in img_path_list
|
||||
]
|
||||
print(f"created {len(vis_com_list)} visual communication objects")
|
||||
# upload images
|
||||
for vis_com in vis_com_list:
|
||||
try:
|
||||
result = collection.insert_one(vis_com.model_dump())
|
||||
except DuplicateKeyError as exc:
|
||||
print('ignoring:\n', exc)
|
||||
else:
|
||||
print(f"inserted document: {result}")
|
||||
@@ -1,34 +0,0 @@
|
||||
"""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')
|
||||
@@ -1,6 +1,4 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from shared.dto import ModelData
|
||||
from shared.docstore.classes import ModelData
|
||||
|
||||
if __name__ == '__main__':
|
||||
# instantiate data object
|
||||
@@ -1,32 +0,0 @@
|
||||
"""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')
|
||||
@@ -1,45 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
from dotenv import load_dotenv
|
||||
|
||||
from shared.data_store import connect_minio
|
||||
from shared.database import connect_mongodb, upsert_prediction
|
||||
from shared.database.classes import VisualCommunication
|
||||
from shared.utils import check_env, 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)
|
||||
# ensure env vars set
|
||||
check_env()
|
||||
# setup logging
|
||||
setup_logging()
|
||||
# connect to minIO
|
||||
minio_client = connect_minio()
|
||||
# connect to MongoDB
|
||||
collection, db, client = connect_mongodb()
|
||||
# get list of image paths
|
||||
test_dir = Path(__file__).parent
|
||||
img_dir = test_dir / 'imgs'
|
||||
img_path_list = [path for path in img_dir.glob('*.jpeg') if path.is_file()]
|
||||
# instantiate data object
|
||||
vis_com_list = [
|
||||
VisualCommunication.from_file(path, minio_client=minio_client)
|
||||
for path in img_path_list
|
||||
]
|
||||
# generate random predictions
|
||||
for vis_com in vis_com_list:
|
||||
vis_com.generate_random_prediction()
|
||||
# upload visual communication
|
||||
for vis_com in vis_com_list:
|
||||
if vis_com.prediction is None:
|
||||
continue
|
||||
upsert_prediction(
|
||||
collection=collection,
|
||||
vis_com_name=vis_com.name,
|
||||
predictions=vis_com.prediction,
|
||||
)
|
||||
@@ -5,7 +5,7 @@ from pathlib import Path
|
||||
|
||||
from dotenv import load_dotenv
|
||||
|
||||
from shared.database import connect_mongodb, count_documents
|
||||
from shared.docstore import connect_mongodb, count_documents
|
||||
|
||||
if __name__ == '__main__':
|
||||
# prepare env vars
|
||||
@@ -5,7 +5,7 @@ from pathlib import Path
|
||||
|
||||
from dotenv import load_dotenv
|
||||
|
||||
from shared.database import connect_mongodb, count_documents
|
||||
from shared.docstore import connect_mongodb, count_documents
|
||||
|
||||
if __name__ == '__main__':
|
||||
# prepare env vars
|
||||
+3
-8
@@ -9,21 +9,16 @@ from models import VisualCommunicationModel
|
||||
from tqdm import tqdm
|
||||
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
|
||||
|
||||
if __name__ == '__main__':
|
||||
# setup logging
|
||||
setup_logging()
|
||||
# connect to minio
|
||||
minio_client = connect_minio()
|
||||
# instantiate model
|
||||
model: VisualCommunicationModel = load_model(client=minio_client)
|
||||
model: VisualCommunicationModel = load_model()
|
||||
model.eval()
|
||||
# setup dataset
|
||||
dataset = VCDADataset(
|
||||
minio_client=minio_client,
|
||||
data_name_list=[
|
||||
'02dbaf48d713e4e6d3a6b98fd2dc866e',
|
||||
],
|
||||
@@ -39,10 +34,10 @@ if __name__ == '__main__':
|
||||
image = torch.unsqueeze(image, 0) # add artificial batch dimension
|
||||
image = image.to(DEVICE)
|
||||
# make prediction
|
||||
pred: ModelData = model(image)
|
||||
pred: dict = model(image)
|
||||
except Exception:
|
||||
print_exc()
|
||||
continue
|
||||
else:
|
||||
print(json.dumps(pred.model_dump(), indent=4))
|
||||
print(json.dumps(pred, indent=4))
|
||||
logging.debug('finished')
|
||||
|
||||
@@ -4,8 +4,6 @@ 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
|
||||
@@ -39,7 +37,7 @@ class VisualCommunicationModel(nn.Module):
|
||||
self.framing_tail = FramingTail()
|
||||
self.salience_tail = SalienceTail()
|
||||
|
||||
def forward(self, x) -> ModelData:
|
||||
def forward(self, x) -> dict:
|
||||
"""Calculate model output on data."""
|
||||
# generate visual representation
|
||||
features = self.resnet_head(x)
|
||||
@@ -57,6 +55,4 @@ class VisualCommunicationModel(nn.Module):
|
||||
'framing': self.framing_tail(features),
|
||||
'salience': self.salience_tail(features),
|
||||
}
|
||||
# convert to respective classes
|
||||
data = ModelData.from_prediction_dict(prediction_dict)
|
||||
return data
|
||||
return prediction_dict
|
||||
|
||||
@@ -4,33 +4,30 @@ 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 shared.repositories import ModelRepository
|
||||
|
||||
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:
|
||||
def load_model() -> 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,
|
||||
)
|
||||
# load model data
|
||||
with ModelRepository() as repo:
|
||||
model_data = repo.get_data(model_object_name)
|
||||
if model_data is None:
|
||||
raise FileNotFoundError(f'model {model_object_name} not found')
|
||||
model_checkpoint = torch.load(model_data.buffer)
|
||||
model.load_state_dict(model_checkpoint)
|
||||
# clean memory
|
||||
model_checkpoint.clear()
|
||||
|
||||
@@ -2,7 +2,6 @@
|
||||
|
||||
import random
|
||||
|
||||
from minio import Minio
|
||||
from PIL import Image
|
||||
from torch import Tensor
|
||||
from torch.utils.data import Dataset
|
||||
@@ -16,7 +15,7 @@ from torchvision.transforms.functional import (
|
||||
to_tensor,
|
||||
)
|
||||
|
||||
from shared.data_store import get_image
|
||||
from shared.repositories import ImageRepository
|
||||
|
||||
# resnet18 original normalization values
|
||||
RESNET_NORMALIZE_MEAN = [0.485, 0.456, 0.406]
|
||||
@@ -28,13 +27,11 @@ class VCDADataset(Dataset):
|
||||
|
||||
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
|
||||
@@ -56,12 +53,11 @@ class VCDADataset(Dataset):
|
||||
|
||||
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)
|
||||
image_name = self.data_name_list[idx]
|
||||
with ImageRepository() as repo:
|
||||
image_data = repo.get_data(image_name)
|
||||
assert image_data is not None
|
||||
tensor = self.image_to_tensor(image_data.image)
|
||||
if self.do_augment:
|
||||
tensor = self.augment(tensor)
|
||||
return tensor
|
||||
|
||||
+3
-6
@@ -18,7 +18,6 @@ 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():
|
||||
@@ -57,18 +56,16 @@ 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_dataset = VCDADataset(data_name_list=val_data_name_list)
|
||||
val_loader = DataLoader(dataset=val_dataset, num_workers=args.loader_workers)
|
||||
|
||||
# create trainer and evaluator
|
||||
@@ -140,7 +137,7 @@ to_save = {
|
||||
}
|
||||
checkpoint_handler = Checkpoint(
|
||||
to_save,
|
||||
f"runs/checkpoints/{run_name}",
|
||||
f'runs/checkpoints/{run_name}',
|
||||
n_saved=3,
|
||||
filename_prefix='best',
|
||||
score_function=lambda engine: -engine.state.metrics['loss'],
|
||||
@@ -154,7 +151,7 @@ if 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:
|
||||
with open(f'runs/configs/{run_name}.json', 'w', encoding='utf-8') as fh:
|
||||
json.dump(config, fh)
|
||||
|
||||
# start training
|
||||
|
||||
@@ -1,38 +0,0 @@
|
||||
"""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,
|
||||
)
|
||||
@@ -3,16 +3,17 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from io import BytesIO
|
||||
from pathlib import Path
|
||||
|
||||
from bson import ObjectId
|
||||
from dotenv import load_dotenv
|
||||
from pymongo.collection import Collection
|
||||
|
||||
from shared.data_store import connect_minio, put
|
||||
from shared.database import VisualCommunication, connect_mongodb
|
||||
from shared.datastore import Datastore
|
||||
from shared.docstore import connect_mongodb
|
||||
from shared.docstore.src.classes import VisualCommunication
|
||||
from shared.utils import check_env, setup_logging
|
||||
from web_ui.src.main import NECESSARY_ENV_VAR_LIST
|
||||
|
||||
|
||||
def list_mongo_document_ids(
|
||||
@@ -76,11 +77,13 @@ if __name__ == '__main__':
|
||||
assert env_path.exists()
|
||||
load_dotenv(env_path)
|
||||
# ensure env vars set
|
||||
check_env()
|
||||
check_env(NECESSARY_ENV_VAR_LIST)
|
||||
# setup logging
|
||||
setup_logging()
|
||||
# connect to minIO
|
||||
minio_client = connect_minio()
|
||||
datastore = Datastore()
|
||||
datastore.connect()
|
||||
assert datastore._client is not None
|
||||
# connect to MongoDB
|
||||
collection, db, client = connect_mongodb()
|
||||
# list documents in mongoDB
|
||||
@@ -94,13 +97,13 @@ if __name__ == '__main__':
|
||||
logging.error('failed getting image from document: %s', doc_id)
|
||||
continue
|
||||
try:
|
||||
# save image to buffer
|
||||
buffer = BytesIO()
|
||||
vis_com.image.save(buffer, 'png') # type: ignore
|
||||
# get image
|
||||
image = vis_com.get_image(
|
||||
minio_client=datastore._client,
|
||||
)
|
||||
# put buffer in minio
|
||||
object_name = put(
|
||||
client=minio_client,
|
||||
buffer=buffer,
|
||||
object_name = datastore.put_image(
|
||||
image=image,
|
||||
)
|
||||
except Exception as exc:
|
||||
logging.debug(exc)
|
||||
|
||||
Generated
+1692
-882
File diff suppressed because it is too large
Load Diff
+10
-2
@@ -22,6 +22,9 @@ flake8-pyproject = "^1.2.3"
|
||||
pandas-stubs = "^2.2.2.240603"
|
||||
types-tqdm = "^4.66.0.20240417"
|
||||
pytest = "^8.3.3"
|
||||
testcontainers = "^4.8.2"
|
||||
coverage = "^7.6.4"
|
||||
pytest-cov = "^6.0.0"
|
||||
|
||||
|
||||
[tool.poetry.group.dev.dependencies]
|
||||
@@ -29,10 +32,11 @@ pandas = "^2.2.1"
|
||||
selenium = "^4.18.1"
|
||||
webdriver-manager = "^4.0.1"
|
||||
retry = "^0.9.2"
|
||||
pre-commit = "^4.1.0"
|
||||
|
||||
|
||||
[tool.poetry.group.model.dependencies]
|
||||
torch = "^2.2.1"
|
||||
torch = "^2.0.0"
|
||||
torchvision = "^0.17.1"
|
||||
torchinfo = "^1.8.0"
|
||||
minio = "^7.2.7"
|
||||
@@ -108,9 +112,13 @@ module = "utils.*"
|
||||
ignore_missing_imports = true
|
||||
|
||||
[[tool.mypy.overrides]]
|
||||
module = "data_store.*"
|
||||
module = "shared.datastore.*"
|
||||
ignore_missing_imports = true
|
||||
|
||||
[[tool.mypy.overrides]]
|
||||
module = "models.*"
|
||||
ignore_missing_imports = true
|
||||
|
||||
[[tool.mypy.overrides]]
|
||||
module = "shared.docstore.*"
|
||||
ignore_missing_imports = true
|
||||
|
||||
@@ -1,8 +0,0 @@
|
||||
from .connect_minio import connect_minio
|
||||
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
|
||||
@@ -1,42 +0,0 @@
|
||||
"""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
|
||||
@@ -1,30 +0,0 @@
|
||||
"""Definition of delete function."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
|
||||
from minio import Minio
|
||||
|
||||
|
||||
def delete(
|
||||
client: Minio,
|
||||
object_name: str,
|
||||
) -> None:
|
||||
"""Delete object from MinIO."""
|
||||
assert isinstance(client, Minio)
|
||||
assert isinstance(object_name, str)
|
||||
bucket_name = os.getenv('MINIO_BUCKET_NAME', default='')
|
||||
assert len(bucket_name) > 0
|
||||
# remove object
|
||||
try:
|
||||
client.remove_object(
|
||||
bucket_name=bucket_name,
|
||||
object_name=object_name,
|
||||
)
|
||||
except Exception as exc:
|
||||
logging.debug(exc)
|
||||
logging.error('failed deleting %s', object_name)
|
||||
else:
|
||||
logging.debug('deleted %s', object_name)
|
||||
@@ -1,39 +0,0 @@
|
||||
"""Definition of get function."""
|
||||
|
||||
import logging
|
||||
import os
|
||||
from io import BytesIO
|
||||
from traceback import print_exc
|
||||
|
||||
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='')
|
||||
# get buffer
|
||||
try:
|
||||
response = client.get_object(
|
||||
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()
|
||||
@@ -1,41 +0,0 @@
|
||||
"""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
|
||||
@@ -1,30 +0,0 @@
|
||||
"""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
|
||||
@@ -1,39 +0,0 @@
|
||||
"""Definition of put function."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
from hashlib import md5
|
||||
from io import BytesIO
|
||||
|
||||
from minio import Minio
|
||||
|
||||
|
||||
def put(
|
||||
client: Minio,
|
||||
buffer: BytesIO,
|
||||
) -> str:
|
||||
"""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='')
|
||||
assert len(bucket_name) > 0
|
||||
# get md5 of image
|
||||
checksum = md5(buffer.getbuffer()).hexdigest()
|
||||
# prepare for saving
|
||||
num_bytes = buffer.tell()
|
||||
buffer.seek(0)
|
||||
# send data to bucket
|
||||
try:
|
||||
client.put_object(
|
||||
bucket_name=bucket_name,
|
||||
object_name=checksum,
|
||||
length=num_bytes,
|
||||
data=buffer,
|
||||
)
|
||||
except Exception as exc:
|
||||
logging.error('failed saving data to MinIO')
|
||||
raise exc
|
||||
logging.debug('saved data to %s', checksum)
|
||||
return checksum
|
||||
@@ -1,44 +0,0 @@
|
||||
"""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
|
||||
@@ -1,45 +0,0 @@
|
||||
"""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
|
||||
@@ -1,14 +0,0 @@
|
||||
"""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_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
|
||||
from .utils.upsert_annotation import upsert_annotation
|
||||
from .utils.upsert_prediction import upsert_prediction
|
||||
from .utils.upsert_visual_communication import upsert_visual_communication
|
||||
@@ -1,6 +0,0 @@
|
||||
"""Database classes module content."""
|
||||
from __future__ import annotations
|
||||
|
||||
from .dataset import Dataset
|
||||
from .exceptions import NoDocumentFoundException
|
||||
from .visual_communication import VisualCommunication
|
||||
@@ -1,71 +0,0 @@
|
||||
"""Definition of database Dataset class."""
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import random
|
||||
from datetime import datetime
|
||||
from datetime import UTC
|
||||
|
||||
from pydantic import BaseModel
|
||||
from pydantic import Field
|
||||
from pymongo.collection import Collection
|
||||
|
||||
|
||||
class Dataset(BaseModel):
|
||||
"""Database Dataset model."""
|
||||
create_time: datetime = Field(default_factory=lambda: datetime.now(UTC))
|
||||
train_names: list[str]
|
||||
test_names: list[str]
|
||||
validation_names: list[str]
|
||||
|
||||
@classmethod
|
||||
def fraction_map(cls) -> dict[str, float]:
|
||||
"""Dict with train, test and validation fractions."""
|
||||
# define map
|
||||
split_map = {
|
||||
'train': 0.7,
|
||||
'test': 0.2,
|
||||
'validation': 0.1,
|
||||
}
|
||||
# sanity check
|
||||
assert sum(split_map.values()) == 1.0
|
||||
return split_map
|
||||
|
||||
@classmethod
|
||||
def new_from_name_list(
|
||||
cls,
|
||||
name_list: list[str],
|
||||
) -> Dataset:
|
||||
"""Generate new dataset from list of filenames."""
|
||||
# calculate split fractions
|
||||
fraction_map = cls.fraction_map()
|
||||
num_total = len(name_list)
|
||||
num_validation = round(num_total * fraction_map['validation'])
|
||||
num_test = round(num_total * fraction_map['test'])
|
||||
# split data
|
||||
validation_name_list = random.choices(name_list, k=num_validation)
|
||||
name_list = [
|
||||
name for name in name_list if name not in validation_name_list
|
||||
]
|
||||
test_name_list = random.choices(name_list, k=num_test)
|
||||
train_name_list = [
|
||||
name for name in name_list if name not in test_name_list
|
||||
]
|
||||
# instantiate object
|
||||
dataset = Dataset(
|
||||
train_names=train_name_list,
|
||||
test_names=test_name_list,
|
||||
validation_names=validation_name_list,
|
||||
)
|
||||
logging.debug('finished')
|
||||
return dataset
|
||||
|
||||
def save(
|
||||
self,
|
||||
collection: Collection,
|
||||
) -> None:
|
||||
"""Save dataset to database."""
|
||||
res = collection.insert_one(
|
||||
document=self.model_dump(),
|
||||
)
|
||||
logging.debug('inserted document: %s', res)
|
||||
@@ -1,6 +0,0 @@
|
||||
"""Definition of database exception."""
|
||||
from __future__ import annotations
|
||||
|
||||
|
||||
class NoDocumentFoundException(Exception):
|
||||
"""Database exception for when no documents are found."""
|
||||
@@ -1,126 +0,0 @@
|
||||
"""Definition of VisualCommunication model."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from base64 import b64decode, b64encode
|
||||
from io import BytesIO
|
||||
from pathlib import Path
|
||||
|
||||
from minio import Minio
|
||||
from PIL import Image
|
||||
from pydantic import BaseModel, ConfigDict
|
||||
from pymongo.collection import Collection
|
||||
|
||||
from shared.data_store import get, put
|
||||
from shared.dto import ModelData
|
||||
|
||||
|
||||
class VisualCommunication(BaseModel):
|
||||
"""Visual communication model."""
|
||||
|
||||
name: str
|
||||
object_name: str
|
||||
annotation: ModelData | None = None
|
||||
prediction: ModelData | None = None
|
||||
model_config = ConfigDict(arbitrary_types_allowed=True)
|
||||
|
||||
@classmethod
|
||||
def classname(cls) -> str:
|
||||
"""Return classname."""
|
||||
return cls.__name__
|
||||
|
||||
@classmethod
|
||||
def upload_image_to_minio(
|
||||
cls,
|
||||
image: Image.Image,
|
||||
minio_client: Minio,
|
||||
) -> str:
|
||||
"""Upload image to MinIO and return MD5 checksum of hashed image."""
|
||||
assert isinstance(image, Image.Image)
|
||||
assert isinstance(minio_client, Minio)
|
||||
buffer = BytesIO()
|
||||
image.save(buffer, 'png')
|
||||
object_name = put(
|
||||
client=minio_client,
|
||||
buffer=buffer,
|
||||
)
|
||||
return object_name
|
||||
|
||||
@classmethod
|
||||
def from_name_and_image(
|
||||
cls,
|
||||
name: str,
|
||||
image: Image.Image,
|
||||
minio_client: Minio,
|
||||
) -> VisualCommunication:
|
||||
"""Instantiate from filename and image that is automatically uploaded
|
||||
to MinIO."""
|
||||
assert isinstance(name, str)
|
||||
assert isinstance(image, Image.Image)
|
||||
assert isinstance(minio_client, Minio)
|
||||
# upload file to minio
|
||||
object_name = VisualCommunication.upload_image_to_minio(
|
||||
image=image,
|
||||
minio_client=minio_client,
|
||||
)
|
||||
return VisualCommunication(name=name, object_name=object_name)
|
||||
|
||||
@classmethod
|
||||
def from_file(cls, path: Path, minio_client: Minio) -> VisualCommunication:
|
||||
"""Instantiate from file."""
|
||||
assert isinstance(path, Path)
|
||||
assert isinstance(minio_client, Minio)
|
||||
# determine name
|
||||
name = path.stem
|
||||
# open image
|
||||
image = Image.open(path)
|
||||
image.load()
|
||||
# instantiate object
|
||||
return VisualCommunication.from_name_and_image(
|
||||
name=name,
|
||||
image=image,
|
||||
minio_client=minio_client,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def decode_image(cls, content: str) -> Image.Image:
|
||||
"""Extract image from webencoded content."""
|
||||
_, content_data = content.split(',')
|
||||
return Image.open(BytesIO(b64decode(content_data)))
|
||||
|
||||
def get_image(self, minio_client: Minio) -> Image.Image:
|
||||
"""Load image data from minio."""
|
||||
assert isinstance(minio_client, Minio)
|
||||
# get buffer from minio
|
||||
buffer = get(
|
||||
client=minio_client,
|
||||
object_name=self.object_name,
|
||||
)
|
||||
# convert data to image
|
||||
im = Image.open(buffer)
|
||||
return im
|
||||
|
||||
def save_to_mongo(self, collection: Collection) -> None:
|
||||
"""Save self as document in MongoDB."""
|
||||
res = collection.insert_one(
|
||||
document=self.model_dump(),
|
||||
)
|
||||
assert res.acknowledged
|
||||
|
||||
def webencoded_image(self, minio_client: Minio) -> str:
|
||||
"""Convert image to be displayed on webpage."""
|
||||
assert isinstance(minio_client, Minio)
|
||||
# get image from minio
|
||||
image = self.get_image(minio_client)
|
||||
# convert images to bytes string
|
||||
buffer = BytesIO()
|
||||
image.save(buffer, format='png')
|
||||
img_enc = b64encode(buffer.getvalue()).decode('utf-8')
|
||||
return f"data:image/png;base64, {img_enc}"
|
||||
|
||||
def generate_random_prediction(self, force: bool = False) -> None:
|
||||
"""Generate random prediction values."""
|
||||
if not force and self.prediction is not None:
|
||||
logging.warning('set force=True to overwrite existing values.')
|
||||
self.prediction = ModelData.from_random()
|
||||
@@ -1,3 +0,0 @@
|
||||
"""Database utils module content."""
|
||||
|
||||
from .connect_mongodb import connect_mongodb
|
||||
@@ -1,30 +0,0 @@
|
||||
"""Definition of function to connect to database using environment
|
||||
variables."""
|
||||
|
||||
import logging
|
||||
import os
|
||||
|
||||
from dotenv import load_dotenv
|
||||
from pymongo import MongoClient
|
||||
|
||||
|
||||
def connect_mongodb():
|
||||
"""Connect to MongoDB using env vars."""
|
||||
# load env vars
|
||||
load_dotenv()
|
||||
necessary_env_vars = [
|
||||
'MONGO_HOST',
|
||||
'MONGO_DB',
|
||||
'MONGO_COLLECTION',
|
||||
]
|
||||
for env_var in necessary_env_vars:
|
||||
assert env_var in os.environ, f"{env_var} not found"
|
||||
# connect to database
|
||||
client = MongoClient(os.getenv('MONGO_HOST'))
|
||||
db = client[os.getenv('MONGO_DB')]
|
||||
# extract collection
|
||||
collection = db[os.getenv('MONGO_COLLECTION')]
|
||||
# set unique index on "name"
|
||||
collection.create_index('name', unique=True)
|
||||
logging.debug('finished')
|
||||
return collection, db, client
|
||||
@@ -1,21 +0,0 @@
|
||||
"""Definition of function to count documents in database."""
|
||||
from __future__ import annotations
|
||||
|
||||
from pymongo.collection import Collection
|
||||
|
||||
|
||||
def count_documents(
|
||||
collection: Collection,
|
||||
only_with_annotation: bool = False,
|
||||
) -> int:
|
||||
"""
|
||||
Get the total number of documents
|
||||
in database that matches the filters.
|
||||
"""
|
||||
assert isinstance(collection, Collection)
|
||||
assert isinstance(only_with_annotation, bool)
|
||||
# build query
|
||||
query = {}
|
||||
if only_with_annotation:
|
||||
query['annotation'] = {'$ne': None}
|
||||
return collection.count_documents(filter=query)
|
||||
@@ -1,16 +0,0 @@
|
||||
"""Definition of get_dataset function."""
|
||||
|
||||
from typing import Literal
|
||||
|
||||
from pymongo.collection import Collection
|
||||
|
||||
|
||||
def get_dataset(
|
||||
collection: Collection,
|
||||
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 []
|
||||
@@ -1,39 +0,0 @@
|
||||
"""Definition of function to get visual communication from database."""
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
|
||||
from pymongo.collection import Collection
|
||||
|
||||
from shared.database import NoDocumentFoundException
|
||||
from shared.database import VisualCommunication
|
||||
|
||||
|
||||
def get_visual_communication(
|
||||
collection: Collection,
|
||||
with_annotation: bool = False,
|
||||
) -> VisualCommunication:
|
||||
"""Get a random visual communication from the database."""
|
||||
query = {}
|
||||
if with_annotation:
|
||||
query['annotation'] = {'$ne': None}
|
||||
else:
|
||||
query['annotation'] = {'$eq': None}
|
||||
data = collection.aggregate(
|
||||
pipeline=[
|
||||
{
|
||||
'$match': query, # find using filters
|
||||
},
|
||||
{
|
||||
'$sample': {
|
||||
'size': 1, # get one random
|
||||
},
|
||||
},
|
||||
],
|
||||
)
|
||||
data_list = list(data) # read data from cursor object
|
||||
if len(data_list) == 0:
|
||||
raise NoDocumentFoundException()
|
||||
vis_com = VisualCommunication.model_validate(data_list[0])
|
||||
logging.debug('finished')
|
||||
return vis_com
|
||||
@@ -1,31 +0,0 @@
|
||||
"""
|
||||
Definition of function to list names
|
||||
of all visual communication documents in database.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from pymongo.collection import Collection
|
||||
|
||||
|
||||
def list_names(
|
||||
collection: Collection,
|
||||
only_with_annotation: bool = True,
|
||||
) -> list[str]:
|
||||
"""List the names of entries that match the filters."""
|
||||
assert isinstance(collection, Collection)
|
||||
assert isinstance(only_with_annotation, bool)
|
||||
# build query
|
||||
query = {}
|
||||
if only_with_annotation:
|
||||
query['annotation'] = {'$ne': None}
|
||||
# execute query
|
||||
res_list = collection.find(
|
||||
filter=query,
|
||||
projection={
|
||||
'_id': False,
|
||||
'name': True,
|
||||
},
|
||||
)
|
||||
# extract information
|
||||
name_list = [elem['name'] for elem in res_list]
|
||||
return name_list
|
||||
@@ -1,38 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
|
||||
from pymongo.collection import Collection
|
||||
|
||||
from shared.database import Dataset
|
||||
|
||||
|
||||
def save_dataset(
|
||||
collection: Collection,
|
||||
dataset: Dataset,
|
||||
) -> None:
|
||||
"""Save dataset to database."""
|
||||
res = collection.insert_one(
|
||||
document=dataset.model_dump(),
|
||||
)
|
||||
logging.debug('inserted document: %s', res)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
from dotenv import load_dotenv
|
||||
|
||||
load_dotenv('local.env')
|
||||
from shared.database import connect_mongodb, list_names
|
||||
|
||||
# connect to database
|
||||
collection, db, client = connect_mongodb()
|
||||
print(client.server_info())
|
||||
|
||||
name_list = list_names(collection=collection, only_with_annotation=True)
|
||||
ds = Dataset.new_from_name_list(name_list=name_list)
|
||||
|
||||
print(ds)
|
||||
# save_dataset(
|
||||
# collection=collection,
|
||||
# dataset=ds
|
||||
# )
|
||||
@@ -1,30 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
|
||||
from pymongo.collection import Collection
|
||||
|
||||
from shared.dto import ModelData
|
||||
|
||||
|
||||
def upsert_annotation(
|
||||
collection: Collection,
|
||||
vis_com_name: str,
|
||||
annotations: ModelData,
|
||||
) -> None:
|
||||
"""Upserts annotation data in the database."""
|
||||
query = {
|
||||
'name': vis_com_name,
|
||||
}
|
||||
update = {
|
||||
'$set': {
|
||||
'annotation': annotations.model_dump(),
|
||||
},
|
||||
}
|
||||
res = collection.update_one(
|
||||
filter=query,
|
||||
update=update,
|
||||
upsert=True,
|
||||
)
|
||||
logging.info('upserted document: %s', res)
|
||||
logging.info('finished')
|
||||
@@ -1,30 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
|
||||
from pymongo.collection import Collection
|
||||
|
||||
from shared.dto import ModelData
|
||||
|
||||
|
||||
def upsert_prediction(
|
||||
collection: Collection,
|
||||
vis_com_name: str,
|
||||
predictions: ModelData,
|
||||
) -> None:
|
||||
"""Upsert prediction data in the database."""
|
||||
query = {
|
||||
'name': vis_com_name,
|
||||
}
|
||||
update = {
|
||||
'$set': {
|
||||
'prediction': predictions.model_dump(),
|
||||
},
|
||||
}
|
||||
res = collection.update_one(
|
||||
filter=query,
|
||||
update=update,
|
||||
upsert=True,
|
||||
)
|
||||
logging.debug('upserted document: %s', res)
|
||||
logging.info('finished')
|
||||
@@ -1,23 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from pymongo.collection import Collection
|
||||
|
||||
from shared.database import VisualCommunication
|
||||
|
||||
|
||||
def upsert_visual_communication(
|
||||
collection: Collection,
|
||||
visual_communication_list: list[VisualCommunication],
|
||||
) -> bool:
|
||||
"""
|
||||
Upsert VisualCommunication object in the database.
|
||||
Returns bool stating success.
|
||||
"""
|
||||
response = collection.insert_many(
|
||||
[
|
||||
vis_com.model_dump()
|
||||
for vis_com
|
||||
in visual_communication_list
|
||||
],
|
||||
)
|
||||
return response.acknowledged
|
||||
@@ -1,15 +0,0 @@
|
||||
"""Data transfer objects module content."""
|
||||
from __future__ import annotations
|
||||
|
||||
from .angle import AngleData
|
||||
from .contact import ContactData
|
||||
from .distance import DistanceData
|
||||
from .framing import FramingData
|
||||
from .information_value import InformationValueData
|
||||
from .modality_color import ModalityColorData
|
||||
from .modality_depth import ModalityDepthData
|
||||
from .modality_lighting import ModalityLightingData
|
||||
from .model_data import ModelData
|
||||
from .point_of_view import PointOfViewData
|
||||
from .salience import SalienceData
|
||||
from .visual_syntax import VisualSyntaxData
|
||||
@@ -1,11 +0,0 @@
|
||||
"""Definition of Angle data model."""
|
||||
from __future__ import annotations
|
||||
|
||||
from .data_model import DataModel
|
||||
|
||||
|
||||
class AngleData(DataModel):
|
||||
"""Angle data model."""
|
||||
high: float
|
||||
eye_level: float
|
||||
low: float
|
||||
@@ -1,11 +0,0 @@
|
||||
"""Definition of ContactData data model."""
|
||||
from __future__ import annotations
|
||||
|
||||
from .data_model import DataModel
|
||||
|
||||
|
||||
class ContactData(DataModel):
|
||||
"""ContactData data model."""
|
||||
|
||||
offer: float
|
||||
demand: float
|
||||
@@ -1,12 +0,0 @@
|
||||
"""Definition of DistanceData data model."""
|
||||
from __future__ import annotations
|
||||
|
||||
from .data_model import DataModel
|
||||
|
||||
|
||||
class DistanceData(DataModel):
|
||||
"""DistanceData data model."""
|
||||
|
||||
long: float
|
||||
medium: float
|
||||
close: float
|
||||
@@ -1,13 +0,0 @@
|
||||
"""Definition of FramingData data model."""
|
||||
from __future__ import annotations
|
||||
|
||||
from .data_model import DataModel
|
||||
|
||||
|
||||
class FramingData(DataModel):
|
||||
"""FramingData data model."""
|
||||
|
||||
frame_lines: float
|
||||
empty_space: float
|
||||
colour_contrast: float
|
||||
form_contrast: float
|
||||
@@ -1,12 +0,0 @@
|
||||
"""Definition of InformationValueData data model."""
|
||||
from __future__ import annotations
|
||||
|
||||
from .data_model import DataModel
|
||||
|
||||
|
||||
class InformationValueData(DataModel):
|
||||
"""InformationValueData data model."""
|
||||
|
||||
given_new: float
|
||||
ideal_real: float
|
||||
central_marginal: float
|
||||
@@ -1,12 +0,0 @@
|
||||
"""Definition of ModalityColorData data model."""
|
||||
from __future__ import annotations
|
||||
|
||||
from .data_model import DataModel
|
||||
|
||||
|
||||
class ModalityColorData(DataModel):
|
||||
"""ModalityColorData data model."""
|
||||
|
||||
high: float
|
||||
medium: float
|
||||
low: float
|
||||
@@ -1,12 +0,0 @@
|
||||
"""Definition of ModalityDepthData data model."""
|
||||
from __future__ import annotations
|
||||
|
||||
from .data_model import DataModel
|
||||
|
||||
|
||||
class ModalityDepthData(DataModel):
|
||||
"""ModalityDepthData data model."""
|
||||
|
||||
high: float
|
||||
medium: float
|
||||
low: float
|
||||
@@ -1,12 +0,0 @@
|
||||
"""Definition of ModalityLightingData data model."""
|
||||
from __future__ import annotations
|
||||
|
||||
from .data_model import DataModel
|
||||
|
||||
|
||||
class ModalityLightingData(DataModel):
|
||||
"""ModalityLightingData data model."""
|
||||
|
||||
high: float
|
||||
medium: float
|
||||
low: float
|
||||
@@ -1,115 +0,0 @@
|
||||
"""Definition of ModelData data model."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from .angle import AngleData
|
||||
from .contact import ContactData
|
||||
from .data_model import DataModel
|
||||
from .distance import DistanceData
|
||||
from .framing import FramingData
|
||||
from .information_value import InformationValueData
|
||||
from .modality_color import ModalityColorData
|
||||
from .modality_depth import ModalityDepthData
|
||||
from .modality_lighting import ModalityLightingData
|
||||
from .point_of_view import PointOfViewData
|
||||
from .salience import SalienceData
|
||||
from .visual_syntax import VisualSyntaxData
|
||||
|
||||
|
||||
class ModelData(DataModel):
|
||||
"""ModelData model for data IO with combined ML model."""
|
||||
|
||||
visual_syntax: VisualSyntaxData
|
||||
contact: ContactData
|
||||
angle: AngleData
|
||||
point_of_view: PointOfViewData
|
||||
distance: DistanceData
|
||||
modality_lighting: ModalityLightingData
|
||||
modality_color: ModalityColorData
|
||||
modality_depth: ModalityDepthData
|
||||
information_value: InformationValueData
|
||||
framing: FramingData
|
||||
salience: SalienceData
|
||||
|
||||
@classmethod
|
||||
def from_random(cls) -> ModelData:
|
||||
"""Instantiate with random numbers."""
|
||||
kwargs = {
|
||||
field: field_info.annotation.from_random() # type: ignore
|
||||
for field, field_info 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)
|
||||
|
||||
@classmethod
|
||||
def from_annotations(
|
||||
cls,
|
||||
visual_syntax: str,
|
||||
contact: str,
|
||||
angle: str,
|
||||
point_of_view: str,
|
||||
distance: str,
|
||||
modality_lighting: str,
|
||||
modality_color: str,
|
||||
modality_depth: str,
|
||||
information_value: str,
|
||||
framing: str,
|
||||
salience: str,
|
||||
) -> ModelData:
|
||||
"""Instantiate from annotation."""
|
||||
kwargs = {
|
||||
'visual_syntax': VisualSyntaxData.from_choice(visual_syntax),
|
||||
'contact': ContactData.from_choice(contact),
|
||||
'angle': AngleData.from_choice(angle),
|
||||
'point_of_view': PointOfViewData.from_choice(point_of_view),
|
||||
'distance': DistanceData.from_choice(distance),
|
||||
'modality_lighting': ModalityLightingData.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)
|
||||
@@ -1,11 +0,0 @@
|
||||
"""Definition of PointOfViewData data model."""
|
||||
from __future__ import annotations
|
||||
|
||||
from .data_model import DataModel
|
||||
|
||||
|
||||
class PointOfViewData(DataModel):
|
||||
"""PointOfViewData data model."""
|
||||
|
||||
frontal: float
|
||||
oblique: float
|
||||
@@ -1,14 +0,0 @@
|
||||
"""Definition of SalienceData data model."""
|
||||
from __future__ import annotations
|
||||
|
||||
from .data_model import DataModel
|
||||
|
||||
|
||||
class SalienceData(DataModel):
|
||||
"""SalienceData data model."""
|
||||
|
||||
size: float
|
||||
colour: float
|
||||
tone: float
|
||||
form: float
|
||||
positioning: float
|
||||
@@ -0,0 +1,8 @@
|
||||
from .src import ImageRepository, ModelRepository, VisualCommunicationRepository
|
||||
from .src.dto import (
|
||||
HexadecimalString,
|
||||
ImageData,
|
||||
ModelData,
|
||||
VisualCommunicationData,
|
||||
VisualCommunicationValues,
|
||||
)
|
||||
@@ -0,0 +1,3 @@
|
||||
from .image_repository import ImageRepository
|
||||
from .model_repository import ModelRepository
|
||||
from .visual_communication_repository import VisualCommunicationRepository
|
||||
@@ -0,0 +1,7 @@
|
||||
from .hexadecimal_string import HexadecimalString
|
||||
from .image_data import ImageData
|
||||
from .model_data import ModelData
|
||||
from .visual_communication_data import (
|
||||
VisualCommunicationData,
|
||||
VisualCommunicationValues,
|
||||
)
|
||||
@@ -0,0 +1,37 @@
|
||||
"""Definition of BytesIO Pydantic Annotation."""
|
||||
|
||||
from io import BytesIO
|
||||
from typing import Any
|
||||
|
||||
from pydantic.json_schema import JsonSchemaValue
|
||||
from pydantic_core import core_schema
|
||||
|
||||
|
||||
class BytesIOPydanticAnnotation:
|
||||
"""Pydantic annotation that defines input validation, as well as general
|
||||
and json serialization."""
|
||||
|
||||
@classmethod
|
||||
def validate_input(cls, v: Any, handler) -> BytesIO:
|
||||
"""Pydantic-related function to validate input on instantiation."""
|
||||
if isinstance(v, BytesIO):
|
||||
return v
|
||||
s = handler(v)
|
||||
return BytesIO(s)
|
||||
|
||||
@classmethod
|
||||
def __get_pydantic_core_schema__(
|
||||
cls,
|
||||
source_type,
|
||||
_handler,
|
||||
) -> core_schema.CoreSchema:
|
||||
assert source_type is BytesIO
|
||||
return core_schema.no_info_wrap_validator_function(
|
||||
function=cls.validate_input,
|
||||
schema=core_schema.str_schema(),
|
||||
serialization=core_schema.to_string_ser_schema(),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def __get_pydantic_json_schema__(cls, _core_schema, handler) -> JsonSchemaValue:
|
||||
return handler(core_schema.str_schema())
|
||||
@@ -0,0 +1,24 @@
|
||||
"""Definition of Checksum DTO."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
|
||||
|
||||
class HexadecimalString(str):
|
||||
"""Hexadecimal-string class."""
|
||||
|
||||
def __new__(cls, string):
|
||||
# ensure proper input format
|
||||
pattern = r'[0-9-a-fA-F]{32}'
|
||||
match = re.match(pattern, string)
|
||||
if match is None:
|
||||
raise ValueError(f'format does not match a hexadecimal-string: {string}')
|
||||
return super().__new__(cls, string)
|
||||
|
||||
def __repr__(self) -> str:
|
||||
class_name = self.__class__.__name__
|
||||
return f"{class_name}('{self}')"
|
||||
|
||||
def __reduce__(self):
|
||||
return self.__class__, (self,)
|
||||
@@ -0,0 +1,38 @@
|
||||
"""Definition of HexadecimalString Pydantic Annotation."""
|
||||
|
||||
from typing import Any
|
||||
|
||||
from pydantic.json_schema import JsonSchemaValue
|
||||
from pydantic_core import core_schema
|
||||
|
||||
from .hexadecimal_string import HexadecimalString
|
||||
|
||||
|
||||
class HexadecimalStringPydanticAnnotation:
|
||||
"""Pydantic annotation that defines input validation, as well as general
|
||||
and json serialization."""
|
||||
|
||||
@classmethod
|
||||
def validate_input(cls, v: Any, handler) -> HexadecimalString:
|
||||
"""Pydantic-related function to validate input on instantiation."""
|
||||
if isinstance(v, HexadecimalString):
|
||||
return v
|
||||
s = handler(v)
|
||||
return HexadecimalString(s)
|
||||
|
||||
@classmethod
|
||||
def __get_pydantic_core_schema__(
|
||||
cls,
|
||||
source_type,
|
||||
_handler,
|
||||
) -> core_schema.CoreSchema:
|
||||
assert source_type is HexadecimalString
|
||||
return core_schema.no_info_wrap_validator_function(
|
||||
function=cls.validate_input,
|
||||
schema=core_schema.str_schema(),
|
||||
serialization=core_schema.to_string_ser_schema(),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def __get_pydantic_json_schema__(cls, _core_schema, handler) -> JsonSchemaValue:
|
||||
return handler(core_schema.str_schema())
|
||||
@@ -0,0 +1,16 @@
|
||||
"""Definition of VisualData DTO."""
|
||||
|
||||
from typing import Annotated
|
||||
|
||||
from PIL import Image
|
||||
from pydantic import Field
|
||||
|
||||
from .image_pydantic_annotation import ImagePydanticAnnotation
|
||||
from .type_checking_base_model import TypeCheckingBaseModel
|
||||
|
||||
|
||||
class ImageData(TypeCheckingBaseModel):
|
||||
"""Visual data class."""
|
||||
|
||||
image: Annotated[Image.Image, ImagePydanticAnnotation]
|
||||
name: str = Field(min_length=1)
|
||||
@@ -0,0 +1,37 @@
|
||||
"""Definition of BytesIO Pydantic Annotation."""
|
||||
|
||||
from typing import Any
|
||||
|
||||
from PIL import Image
|
||||
from pydantic.json_schema import JsonSchemaValue
|
||||
from pydantic_core import core_schema
|
||||
|
||||
|
||||
class ImagePydanticAnnotation:
|
||||
"""Pydantic annotation that defines input validation, as well as general
|
||||
and json serialization."""
|
||||
|
||||
@classmethod
|
||||
def validate_input(cls, v: Any, handler) -> Image.Image:
|
||||
"""Pydantic-related function to validate input on instantiation."""
|
||||
if isinstance(v, Image.Image):
|
||||
return v
|
||||
s = handler(v)
|
||||
return Image.open(s)
|
||||
|
||||
@classmethod
|
||||
def __get_pydantic_core_schema__(
|
||||
cls,
|
||||
source_type,
|
||||
_handler,
|
||||
) -> core_schema.CoreSchema:
|
||||
assert source_type is Image.Image
|
||||
return core_schema.no_info_wrap_validator_function(
|
||||
function=cls.validate_input,
|
||||
schema=core_schema.str_schema(),
|
||||
serialization=core_schema.to_string_ser_schema(),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def __get_pydantic_json_schema__(cls, _core_schema, handler) -> JsonSchemaValue:
|
||||
return handler(core_schema.str_schema())
|
||||
@@ -0,0 +1,60 @@
|
||||
"""Definition of ModelData DTO."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from hashlib import md5
|
||||
from io import BytesIO
|
||||
from typing import Annotated
|
||||
|
||||
import torch
|
||||
from pydantic import Field
|
||||
|
||||
from .bytes_io_pydantic_annotation import BytesIOPydanticAnnotation
|
||||
from .hexadecimal_string import HexadecimalString
|
||||
from .hexadecimal_string_pydantic_annotation import HexadecimalStringPydanticAnnotation
|
||||
from .type_checking_base_model import TypeCheckingBaseModel
|
||||
|
||||
|
||||
class ModelData(TypeCheckingBaseModel):
|
||||
"""Model Data DTO."""
|
||||
|
||||
buffer: Annotated[BytesIO, BytesIOPydanticAnnotation]
|
||||
buffer_checksum: Annotated[HexadecimalString, HexadecimalStringPydanticAnnotation]
|
||||
class_name: str = Field(
|
||||
min_length=1,
|
||||
description='name of model class to generate data.',
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def calculate_checksum(buffer: BytesIO) -> HexadecimalString:
|
||||
"""Calculate buffer checksum."""
|
||||
checksum = md5(buffer.getbuffer()).hexdigest()
|
||||
return HexadecimalString(checksum)
|
||||
|
||||
@staticmethod
|
||||
def model_to_buffer(model: torch.nn.Module) -> BytesIO:
|
||||
"""Save model to buffer."""
|
||||
assert isinstance(model, torch.nn.Module)
|
||||
buffer = BytesIO()
|
||||
torch.save(model.state_dict(), buffer)
|
||||
return buffer
|
||||
|
||||
@classmethod
|
||||
def from_model(cls, model: torch.nn.Module) -> ModelData:
|
||||
"""Instantiate from torch module."""
|
||||
assert isinstance(model, torch.nn.Module)
|
||||
# get model name
|
||||
class_name = type(model).__name__
|
||||
# save data to buffer
|
||||
buffer = cls.model_to_buffer(model)
|
||||
buffer = BytesIO()
|
||||
torch.save(model.state_dict(), buffer)
|
||||
# calculate checksum
|
||||
buffer_checksum = cls.calculate_checksum(buffer)
|
||||
# instantiate from buffer
|
||||
data = cls(
|
||||
buffer=buffer,
|
||||
buffer_checksum=buffer_checksum,
|
||||
class_name=class_name,
|
||||
)
|
||||
return data
|
||||
@@ -0,0 +1,11 @@
|
||||
"""Definition of TypeCheckingBaseModel class."""
|
||||
|
||||
from pydantic import BaseModel, ConfigDict
|
||||
|
||||
|
||||
class TypeCheckingBaseModel(BaseModel):
|
||||
"""BaseModel with added type checking on input types."""
|
||||
|
||||
model_config = ConfigDict(
|
||||
frozen=True, # ensure data immutability
|
||||
)
|
||||
@@ -0,0 +1,2 @@
|
||||
from .visual_communication_data import VisualCommunicationData
|
||||
from .visual_communication_values import VisualCommunicationValues
|
||||
@@ -0,0 +1,11 @@
|
||||
"""Definition of AngleValues DTO."""
|
||||
|
||||
from .values_model import ValuesModel
|
||||
|
||||
|
||||
class AngleValues(ValuesModel):
|
||||
"""Angle values DTO."""
|
||||
|
||||
high: float
|
||||
eye_level: float
|
||||
low: float
|
||||
@@ -0,0 +1,10 @@
|
||||
"""Definition of ContactValues DTO."""
|
||||
|
||||
from .values_model import ValuesModel
|
||||
|
||||
|
||||
class ContactValues(ValuesModel):
|
||||
"""Contact values DTO."""
|
||||
|
||||
offer: float
|
||||
demand: float
|
||||
@@ -0,0 +1,11 @@
|
||||
"""Definition of DistanceValues DTO."""
|
||||
|
||||
from .values_model import ValuesModel
|
||||
|
||||
|
||||
class DistanceValues(ValuesModel):
|
||||
"""Distance values DTO."""
|
||||
|
||||
long: float
|
||||
medium: float
|
||||
close: float
|
||||
@@ -0,0 +1,12 @@
|
||||
"""Definition of FramingValues DTO."""
|
||||
|
||||
from .values_model import ValuesModel
|
||||
|
||||
|
||||
class FramingValues(ValuesModel):
|
||||
"""Framing values DTO."""
|
||||
|
||||
frame_lines: float
|
||||
empty_space: float
|
||||
colour_contrast: float
|
||||
form_contrast: float
|
||||
@@ -0,0 +1,11 @@
|
||||
"""Definition of InformationValueValues DTO."""
|
||||
|
||||
from .values_model import ValuesModel
|
||||
|
||||
|
||||
class InformationValueValues(ValuesModel):
|
||||
"""Information value values DTO."""
|
||||
|
||||
given_new: float
|
||||
ideal_real: float
|
||||
central_marginal: float
|
||||
@@ -0,0 +1,11 @@
|
||||
"""Definition of ModalityColorValues DTO."""
|
||||
|
||||
from .values_model import ValuesModel
|
||||
|
||||
|
||||
class ModalityColorValues(ValuesModel):
|
||||
"""Modality color values DTO."""
|
||||
|
||||
high: float
|
||||
medium: float
|
||||
low: float
|
||||
@@ -0,0 +1,11 @@
|
||||
"""Definition of ModalityDepthValues DTO."""
|
||||
|
||||
from .values_model import ValuesModel
|
||||
|
||||
|
||||
class ModalityDepthValues(ValuesModel):
|
||||
"""Modality depth values DTO."""
|
||||
|
||||
high: float
|
||||
medium: float
|
||||
low: float
|
||||
@@ -0,0 +1,11 @@
|
||||
"""Definition of ModalityLightingValues DTO."""
|
||||
|
||||
from .values_model import ValuesModel
|
||||
|
||||
|
||||
class ModalityLightingValues(ValuesModel):
|
||||
"""Modality lighting values DTO."""
|
||||
|
||||
high: float
|
||||
medium: float
|
||||
low: float
|
||||
@@ -0,0 +1,10 @@
|
||||
"""Definition of PointOfViewValues DTO."""
|
||||
|
||||
from .values_model import ValuesModel
|
||||
|
||||
|
||||
class PointOfViewValues(ValuesModel):
|
||||
"""Point-of-view values DTO."""
|
||||
|
||||
frontal: float
|
||||
oblique: float
|
||||
@@ -0,0 +1,13 @@
|
||||
"""Definition of SalienceValues DTO."""
|
||||
|
||||
from .values_model import ValuesModel
|
||||
|
||||
|
||||
class SalienceValues(ValuesModel):
|
||||
"""Salience values DTO."""
|
||||
|
||||
size: float
|
||||
colour: float
|
||||
tone: float
|
||||
form: float
|
||||
positioning: float
|
||||
+18
-25
@@ -1,18 +1,20 @@
|
||||
"""Definition of DataModel base class."""
|
||||
"""Definition of ValuesModel base class."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import random
|
||||
|
||||
from pydantic import BaseModel, ValidationError
|
||||
from pydantic import BaseModel, ConfigDict
|
||||
from torch import Tensor
|
||||
|
||||
|
||||
class DataModel(BaseModel):
|
||||
"""DataModel base class."""
|
||||
class ValuesModel(BaseModel):
|
||||
"""ValuesModel base class."""
|
||||
|
||||
@classmethod
|
||||
def classname(cls) -> str:
|
||||
"""Return classname."""
|
||||
return cls.__name__
|
||||
model_config = ConfigDict(
|
||||
validate_assignment=True, # argument type checking
|
||||
frozen=True, # ensure data immutability
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def list_fields(cls) -> list[str]:
|
||||
@@ -26,16 +28,16 @@ class DataModel(BaseModel):
|
||||
return cls(**kwargs)
|
||||
|
||||
@classmethod
|
||||
def from_choice(cls, option: str):
|
||||
def from_choice(cls, option: str) -> ValuesModel:
|
||||
"""Instantiate from choice."""
|
||||
if option is None:
|
||||
raise ValidationError()
|
||||
assert isinstance(option, str), 'option is not a string'
|
||||
assert isinstance(option, str)
|
||||
assert len(option) > 0
|
||||
allowed_options_list = cls.list_fields()
|
||||
assert (
|
||||
option in allowed_options_list
|
||||
), f"{option} is not among allowed fields {allowed_options_list}"
|
||||
kwargs = {field: 0 for field in cls.list_fields()}
|
||||
if option not in allowed_options_list:
|
||||
raise ValueError(f'option {option} must be in {allowed_options_list}')
|
||||
# generate field values
|
||||
kwargs = {field: 0 for field in allowed_options_list}
|
||||
# set chosen value to max probability
|
||||
kwargs[option] = 1
|
||||
return cls(**kwargs)
|
||||
|
||||
@@ -47,15 +49,6 @@ class DataModel(BaseModel):
|
||||
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 += ', '.join(
|
||||
[f'{field}={value:.3f}' for field, value in model_dict.items()],
|
||||
)
|
||||
model_repr_str += ')'
|
||||
return model_repr_str
|
||||
|
||||
def highest_score_field(self) -> str:
|
||||
"""Return name of field with highest score."""
|
||||
model_dict = self.model_dump()
|
||||
@@ -0,0 +1,14 @@
|
||||
"""Definition of VisualCommunicationData DTO."""
|
||||
|
||||
from pydantic import Field
|
||||
|
||||
from ..type_checking_base_model import TypeCheckingBaseModel
|
||||
from .visual_communication_values import VisualCommunicationValues
|
||||
|
||||
|
||||
class VisualCommunicationData(TypeCheckingBaseModel):
|
||||
"""Visual communication data class."""
|
||||
|
||||
name: str = Field(min_length=1)
|
||||
annotation: VisualCommunicationValues | None = None
|
||||
prediction: VisualCommunicationValues | None = None
|
||||
@@ -0,0 +1,49 @@
|
||||
"""Definition of VisualCommunicationValues DTO."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from ..type_checking_base_model import TypeCheckingBaseModel
|
||||
from .angle_values import AngleValues
|
||||
from .contact_values import ContactValues
|
||||
from .distance_values import DistanceValues
|
||||
from .framing_values import FramingValues
|
||||
from .information_value_values import InformationValueValues
|
||||
from .modality_color_values import ModalityColorValues
|
||||
from .modality_depth_values import ModalityDepthValues
|
||||
from .modality_lighting_values import ModalityLightingValues
|
||||
from .point_of_view_values import PointOfViewValues
|
||||
from .salience_values import SalienceValues
|
||||
from .visual_syntax_values import VisualSyntaxValues
|
||||
|
||||
|
||||
class VisualCommunicationValues(TypeCheckingBaseModel):
|
||||
"""Visual communication values class."""
|
||||
|
||||
visual_syntax: VisualSyntaxValues
|
||||
contact: ContactValues
|
||||
angle: AngleValues
|
||||
point_of_view: PointOfViewValues
|
||||
distance: DistanceValues
|
||||
modality_lighting: ModalityLightingValues
|
||||
modality_color: ModalityColorValues
|
||||
modality_depth: ModalityDepthValues
|
||||
information_value: InformationValueValues
|
||||
framing: FramingValues
|
||||
salience: SalienceValues
|
||||
|
||||
@classmethod
|
||||
def from_random(cls) -> VisualCommunicationValues:
|
||||
"""Create a random instance."""
|
||||
return cls(
|
||||
visual_syntax=VisualSyntaxValues.from_random(),
|
||||
contact=ContactValues.from_random(),
|
||||
angle=AngleValues.from_random(),
|
||||
point_of_view=PointOfViewValues.from_random(),
|
||||
distance=DistanceValues.from_random(),
|
||||
modality_lighting=ModalityLightingValues.from_random(),
|
||||
modality_color=ModalityColorValues.from_random(),
|
||||
modality_depth=ModalityDepthValues.from_random(),
|
||||
information_value=InformationValueValues.from_random(),
|
||||
framing=FramingValues.from_random(),
|
||||
salience=SalienceValues.from_random(),
|
||||
)
|
||||
+4
-5
@@ -1,11 +1,10 @@
|
||||
"""Definition of VisualSyntaxData data model."""
|
||||
from __future__ import annotations
|
||||
"""Definition of VisualSyntaxValues DTO."""
|
||||
|
||||
from .data_model import DataModel
|
||||
from .values_model import ValuesModel
|
||||
|
||||
|
||||
class VisualSyntaxData(DataModel):
|
||||
"""VisualSyntaxData data model."""
|
||||
class VisualSyntaxValues(ValuesModel):
|
||||
"""Visual syntax values DTO."""
|
||||
|
||||
non_transactional_action: float
|
||||
non_transactional_reaction: float
|
||||
@@ -0,0 +1,74 @@
|
||||
"""Definition of ImageRepository class."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from io import BytesIO
|
||||
from pathlib import Path
|
||||
|
||||
from PIL import Image
|
||||
|
||||
from .dto import ImageData
|
||||
from .implementations import MinioImplementation
|
||||
from .interfaces import ImageInterface
|
||||
|
||||
|
||||
class ImageRepository(ImageInterface, MinioImplementation):
|
||||
"""Image repository class that handles CRUD functionality for
|
||||
VisualData."""
|
||||
|
||||
def __enter__(self) -> ImageRepository:
|
||||
self.connect()
|
||||
return self
|
||||
|
||||
@staticmethod
|
||||
def _build_path(image_name: str) -> Path:
|
||||
"""Build object path."""
|
||||
assert isinstance(image_name, str)
|
||||
path = Path('images') / image_name
|
||||
return path
|
||||
|
||||
def get_data(self, image_name: str) -> ImageData | None:
|
||||
"""Get Visual data."""
|
||||
assert isinstance(image_name, str)
|
||||
assert len(image_name) > 0
|
||||
# build path
|
||||
path = self._build_path(image_name)
|
||||
# get object from bucket
|
||||
buffer = self._get(path)
|
||||
# handle if no data found
|
||||
if not buffer:
|
||||
return None
|
||||
# convert data
|
||||
image = Image.open(buffer)
|
||||
data = ImageData(image=image, name=image_name)
|
||||
return data
|
||||
|
||||
def put_data(self, data: ImageData) -> None:
|
||||
"""Put visual data."""
|
||||
assert isinstance(data, ImageData)
|
||||
# build path
|
||||
path = self._build_path(data.name)
|
||||
# save image to buffer
|
||||
buffer = BytesIO()
|
||||
data.image.save(buffer, 'png')
|
||||
# put object in bucket
|
||||
self._put(path, buffer)
|
||||
|
||||
def remove_data(self, image_name: str) -> None:
|
||||
"""Remove visual data."""
|
||||
assert isinstance(image_name, str)
|
||||
assert len(image_name) > 0
|
||||
# build path
|
||||
path = self._build_path(image_name)
|
||||
# remove object
|
||||
self._delete(path)
|
||||
|
||||
def list_names(self) -> list[str]:
|
||||
"""List names of all images."""
|
||||
# build path
|
||||
path = self._build_path('')
|
||||
# list object paths
|
||||
obj_path_list = self._list_objects(path)
|
||||
# strip prefix
|
||||
name_list = [obj_path.split('/')[-1] for obj_path in obj_path_list]
|
||||
return name_list
|
||||
@@ -0,0 +1,2 @@
|
||||
from .minio_implementation import MinioImplementation
|
||||
from .mongo_implementation import MongoImplementation
|
||||
@@ -0,0 +1,186 @@
|
||||
"""MinIO implementation."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
from io import BytesIO
|
||||
from pathlib import Path
|
||||
from traceback import format_exc
|
||||
|
||||
from minio import Minio
|
||||
|
||||
from shared.utils import check_env
|
||||
|
||||
from ..interfaces import DatabaseInterface
|
||||
|
||||
|
||||
class MinioImplementation(DatabaseInterface):
|
||||
"""MinIO basic CRUD implementation."""
|
||||
|
||||
def __init__(self):
|
||||
# ensure necessary env vars available
|
||||
var_list = {
|
||||
'MINIO_ENDPOINT',
|
||||
'MINIO_ACCESS_KEY',
|
||||
'MINIO_SECRET_KEY',
|
||||
'MINIO_BUCKET_NAME',
|
||||
}
|
||||
check_env(var_list)
|
||||
# prepare internal variables
|
||||
self._client: Minio | None = None
|
||||
self._bucket_name: str | None = None
|
||||
|
||||
def connect(self):
|
||||
"""Connect to MinIO server."""
|
||||
# prepare arguments
|
||||
minio_endpoint = str(os.getenv('MINIO_ENDPOINT'))
|
||||
minio_access_key = str(os.getenv('MINIO_ACCESS_KEY'))
|
||||
minio_secret_key = str(os.getenv('MINIO_SECRET_KEY'))
|
||||
minio_bucket_name = str(os.getenv('MINIO_BUCKET_NAME'))
|
||||
# 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.debug('creating bucket: %s', minio_bucket_name)
|
||||
client.make_bucket(bucket_name=minio_bucket_name)
|
||||
# persist state
|
||||
self._client = client
|
||||
self._bucket_name = minio_bucket_name
|
||||
|
||||
def close(self) -> None:
|
||||
"""Close connection to MinIO server.
|
||||
|
||||
N.B. MinIO connection cannot be closed manually.
|
||||
"""
|
||||
self._client = None
|
||||
self._bucket_name = None
|
||||
|
||||
def connected(self):
|
||||
"""Check connection to Minio."""
|
||||
res = isinstance(self._client, Minio)
|
||||
logging.debug(res)
|
||||
return res
|
||||
|
||||
def __enter__(self) -> MinioImplementation:
|
||||
self.connect()
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc_val, exc_tb) -> None:
|
||||
if any(
|
||||
(
|
||||
exc_type is not None,
|
||||
exc_val is not None,
|
||||
exc_tb is not None,
|
||||
),
|
||||
):
|
||||
logging.error('error while exiting context')
|
||||
self.close()
|
||||
|
||||
def _put(
|
||||
self,
|
||||
path: Path,
|
||||
buffer: BytesIO,
|
||||
) -> None:
|
||||
"""Save in-memory buffer as object in MinIO."""
|
||||
assert isinstance(path, Path)
|
||||
assert isinstance(buffer, BytesIO)
|
||||
assert isinstance(self._client, Minio)
|
||||
assert isinstance(self._bucket_name, str)
|
||||
# prepare for saving
|
||||
num_bytes = len(buffer.getvalue())
|
||||
buffer.seek(0)
|
||||
# send data to bucket
|
||||
try:
|
||||
self._client.put_object(
|
||||
bucket_name=self._bucket_name,
|
||||
object_name=path.as_posix(),
|
||||
length=num_bytes,
|
||||
data=buffer,
|
||||
)
|
||||
logging.debug('saved data to %s', path)
|
||||
except Exception as exc:
|
||||
logging.error('failed saving data to MinIO')
|
||||
raise exc
|
||||
|
||||
def _get(
|
||||
self,
|
||||
path: Path,
|
||||
) -> BytesIO | None:
|
||||
"""Get object from MinIO as in-memory buffer."""
|
||||
assert isinstance(path, Path)
|
||||
assert isinstance(self._client, Minio)
|
||||
assert isinstance(self._bucket_name, str)
|
||||
try:
|
||||
# make request
|
||||
response = self._client.get_object(
|
||||
bucket_name=self._bucket_name,
|
||||
object_name=path.as_posix(),
|
||||
)
|
||||
assert response.status == 200
|
||||
# get buffer
|
||||
buffer = BytesIO()
|
||||
chunk_size = 2**14
|
||||
while chunk := response.read(chunk_size):
|
||||
buffer.write(chunk)
|
||||
buffer.seek(0)
|
||||
logging.debug('got %s', path)
|
||||
return buffer
|
||||
except Exception:
|
||||
logging.error('failed getting data from MinIO')
|
||||
logging.debug(format_exc())
|
||||
return None
|
||||
finally:
|
||||
# close connection if established
|
||||
if 'response' in locals():
|
||||
response.close()
|
||||
response.release_conn()
|
||||
|
||||
def _delete(
|
||||
self,
|
||||
path: Path,
|
||||
) -> None:
|
||||
"""Delete object from MinIO."""
|
||||
assert isinstance(path, Path)
|
||||
assert isinstance(self._client, Minio)
|
||||
assert isinstance(self._bucket_name, str)
|
||||
# remove object
|
||||
try:
|
||||
self._client.remove_object(
|
||||
bucket_name=self._bucket_name,
|
||||
object_name=path.as_posix(),
|
||||
)
|
||||
logging.debug('deleted %s', path)
|
||||
except Exception as exc:
|
||||
logging.error('failed deleting %s', path)
|
||||
logging.debug(format_exc())
|
||||
raise exc
|
||||
|
||||
def _list_objects(
|
||||
self,
|
||||
path: Path,
|
||||
) -> list[str]:
|
||||
"""List objects in bucket under path."""
|
||||
assert isinstance(path, Path)
|
||||
assert isinstance(self._client, Minio)
|
||||
assert isinstance(self._bucket_name, str)
|
||||
try:
|
||||
# list objects
|
||||
obj_list = self._client.list_objects(
|
||||
bucket_name=self._bucket_name,
|
||||
prefix=path.as_posix(),
|
||||
recursive=True,
|
||||
)
|
||||
# extract info
|
||||
name_list = [obj.object_name for obj in obj_list]
|
||||
logging.debug('got %s objects matching %s', len(name_list), path)
|
||||
return name_list
|
||||
except Exception as exc:
|
||||
logging.error('failed listing objects under %s', path)
|
||||
logging.debug(format_exc())
|
||||
raise exc
|
||||
@@ -0,0 +1,145 @@
|
||||
"""Mongo implementation."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
import traceback
|
||||
|
||||
from pymongo import MongoClient
|
||||
from pymongo.collection import Collection
|
||||
from pymongo.database import Database
|
||||
from pymongo.errors import ServerSelectionTimeoutError
|
||||
|
||||
from shared.utils import check_env
|
||||
|
||||
from ..interfaces import DatabaseInterface
|
||||
|
||||
|
||||
class MongoImplementation(DatabaseInterface):
|
||||
"""MongoDB basic CRUD implementation."""
|
||||
|
||||
def __init__(self):
|
||||
# ensure necessary env vars available
|
||||
var_list = {
|
||||
'MONGO_ENDPOINT',
|
||||
'MONGO_DB',
|
||||
'MONGO_COLLECTION',
|
||||
}
|
||||
check_env(var_list)
|
||||
# prepare internal variables
|
||||
self.client: MongoClient | None = None
|
||||
self.db: Database | None = None
|
||||
self.collection: Collection | None = None
|
||||
|
||||
def connect(self) -> None:
|
||||
"""Connect to Mongo server."""
|
||||
# prepare arguments
|
||||
mongo_endpoint = str(os.getenv('MONGO_ENDPOINT'))
|
||||
mongo_database = str(os.getenv('MONGO_DB'))
|
||||
mongo_collection = str(os.getenv('MONGO_COLLECTION'))
|
||||
# connect client
|
||||
client: MongoClient = MongoClient(mongo_endpoint)
|
||||
database = client[mongo_database]
|
||||
collection = database[mongo_collection]
|
||||
# set unique index on 'name'
|
||||
collection.create_index(keys='name', unique=True)
|
||||
# persist state
|
||||
self._client = client
|
||||
self._database = database
|
||||
self._collection = collection
|
||||
|
||||
def close(self) -> None:
|
||||
"""Close connection to Mongo server."""
|
||||
self._client.close()
|
||||
self._client = None # type: ignore
|
||||
self._database = None # type: ignore
|
||||
self._collection = None # type: ignore
|
||||
|
||||
def connected(self) -> bool:
|
||||
"""Check connection to Mongo."""
|
||||
if self._client is None:
|
||||
return False
|
||||
try:
|
||||
# trigger fetch data
|
||||
_ = self._client.server_info()
|
||||
res = True
|
||||
except ServerSelectionTimeoutError:
|
||||
res = False
|
||||
logging.debug(res)
|
||||
return res
|
||||
|
||||
def __enter__(self) -> MongoImplementation:
|
||||
self.connect()
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc_val, exc_tb) -> None:
|
||||
if any(
|
||||
(
|
||||
exc_type is not None,
|
||||
exc_val is not None,
|
||||
exc_tb is not None,
|
||||
),
|
||||
):
|
||||
logging.error('error while exiting context')
|
||||
traceback.print_exception(exc_type, exc_val, exc_tb)
|
||||
self.close()
|
||||
|
||||
def _save(
|
||||
self,
|
||||
data: dict,
|
||||
query: dict,
|
||||
) -> None:
|
||||
"""Save document in Mongo."""
|
||||
assert isinstance(data, dict)
|
||||
assert isinstance(query, dict)
|
||||
assert 'name' in data
|
||||
assert self.connected()
|
||||
self._collection.update_one(
|
||||
filter=query,
|
||||
update={
|
||||
'$set': data.copy(),
|
||||
},
|
||||
upsert=True,
|
||||
)
|
||||
logging.debug('Save %s', data)
|
||||
|
||||
def _get(
|
||||
self,
|
||||
query: dict,
|
||||
) -> dict | None:
|
||||
"""Get document from Mongo."""
|
||||
assert isinstance(query, dict)
|
||||
assert self.connected()
|
||||
doc = self._collection.find_one(query, projection={'_id': False})
|
||||
logging.debug('Found %s', doc)
|
||||
return doc
|
||||
|
||||
def _delete(
|
||||
self,
|
||||
query: dict,
|
||||
) -> None:
|
||||
"""Remove document from Mongo."""
|
||||
assert isinstance(query, dict)
|
||||
assert self.connected()
|
||||
doc = self._collection.delete_one(query)
|
||||
logging.debug('Deleted %s', doc)
|
||||
|
||||
def _list_documents(self, key='name') -> list[str]:
|
||||
"""List documents in Mongo."""
|
||||
assert isinstance(key, str)
|
||||
assert len(key) > 0
|
||||
# build query
|
||||
doc_list = list(
|
||||
self._collection.find(
|
||||
filter={},
|
||||
projection={
|
||||
'_id': False,
|
||||
key: True,
|
||||
},
|
||||
),
|
||||
)
|
||||
# extract values
|
||||
value_list = [doc[key] for doc in doc_list]
|
||||
logging.debug('Got %s document(s)', len(doc_list))
|
||||
return value_list
|
||||
@@ -0,0 +1,4 @@
|
||||
from .database_interface import DatabaseInterface
|
||||
from .image_interface import ImageInterface
|
||||
from .model_interface import ModelInterface
|
||||
from .visual_communication_interface import VisualCommunicationInterface
|
||||
@@ -0,0 +1,28 @@
|
||||
"""Definition of DatabaseInterface class."""
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
|
||||
|
||||
class DatabaseInterface(ABC):
|
||||
"""Interface base class adding 'connect', 'close' and context
|
||||
functionalities."""
|
||||
|
||||
@abstractmethod
|
||||
def connect(self):
|
||||
raise NotImplementedError()
|
||||
|
||||
@abstractmethod
|
||||
def close(self):
|
||||
raise NotImplementedError()
|
||||
|
||||
@abstractmethod
|
||||
def connected(self) -> bool:
|
||||
raise NotImplementedError()
|
||||
|
||||
@abstractmethod
|
||||
def __enter__(self):
|
||||
raise NotImplementedError()
|
||||
|
||||
@abstractmethod
|
||||
def __exit__(self, exc_type, exc_val, exc_tb):
|
||||
raise NotImplementedError()
|
||||
@@ -0,0 +1,26 @@
|
||||
"""Definition of ImageInterface."""
|
||||
|
||||
from abc import abstractmethod
|
||||
|
||||
from ..dto import ImageData
|
||||
from .database_interface import DatabaseInterface
|
||||
|
||||
|
||||
class ImageInterface(DatabaseInterface):
|
||||
"""Image interface class."""
|
||||
|
||||
@abstractmethod
|
||||
def get_data(self, image_name: str) -> ImageData | None:
|
||||
raise NotImplementedError()
|
||||
|
||||
@abstractmethod
|
||||
def put_data(self, data: ImageData) -> None:
|
||||
raise NotImplementedError()
|
||||
|
||||
@abstractmethod
|
||||
def remove_data(self, image_name: str) -> None:
|
||||
raise NotImplementedError()
|
||||
|
||||
@abstractmethod
|
||||
def list_names(self) -> list[str]:
|
||||
raise NotImplementedError()
|
||||
@@ -0,0 +1,26 @@
|
||||
"""Definition of ModelInterface."""
|
||||
|
||||
from abc import abstractmethod
|
||||
|
||||
from ..dto import ModelData
|
||||
from .database_interface import DatabaseInterface
|
||||
|
||||
|
||||
class ModelInterface(DatabaseInterface):
|
||||
"""Model interface class."""
|
||||
|
||||
@abstractmethod
|
||||
def get_data(self, object_name: str) -> ModelData | None:
|
||||
raise NotImplementedError()
|
||||
|
||||
@abstractmethod
|
||||
def put_data(self, data: ModelData) -> None:
|
||||
raise NotImplementedError()
|
||||
|
||||
@abstractmethod
|
||||
def remove_data(self, object_name: str) -> None:
|
||||
raise NotImplementedError()
|
||||
|
||||
@abstractmethod
|
||||
def list_names(self) -> list[str]:
|
||||
raise NotImplementedError()
|
||||
@@ -0,0 +1,26 @@
|
||||
"""Definition of VisualCommunicationInterface."""
|
||||
|
||||
from abc import abstractmethod
|
||||
|
||||
from ..dto import VisualCommunicationData
|
||||
from .database_interface import DatabaseInterface
|
||||
|
||||
|
||||
class VisualCommunicationInterface(DatabaseInterface):
|
||||
"""Visual communication interface class."""
|
||||
|
||||
@abstractmethod
|
||||
def get_data(self, name: str) -> VisualCommunicationData | None:
|
||||
raise NotImplementedError()
|
||||
|
||||
@abstractmethod
|
||||
def put_data(self, data: VisualCommunicationData) -> None:
|
||||
raise NotImplementedError()
|
||||
|
||||
@abstractmethod
|
||||
def remove_data(self, name: str) -> None:
|
||||
raise NotImplementedError()
|
||||
|
||||
@abstractmethod
|
||||
def list_names(self) -> list[str]:
|
||||
raise NotImplementedError()
|
||||
@@ -0,0 +1,77 @@
|
||||
"""Definition of ModelRepository class."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
from .dto import HexadecimalString, ModelData
|
||||
from .implementations import MinioImplementation
|
||||
from .interfaces import ModelInterface
|
||||
|
||||
|
||||
class ModelRepository(ModelInterface, MinioImplementation):
|
||||
"""Model repository class that handles CRUD functionality for ModelData."""
|
||||
|
||||
def __enter__(self) -> ModelRepository:
|
||||
self.connect()
|
||||
return self
|
||||
|
||||
@staticmethod
|
||||
def _prefix() -> Path:
|
||||
"""Object name prefix."""
|
||||
return Path('models')
|
||||
|
||||
@classmethod
|
||||
def _build_object_name(cls, data: ModelData) -> str:
|
||||
"""Build object name from data."""
|
||||
return f'{data.class_name}-{data.buffer_checksum}'
|
||||
|
||||
def get_data(self, object_name: str) -> ModelData | None:
|
||||
"""Get model data."""
|
||||
assert isinstance(object_name, str)
|
||||
# build path
|
||||
path = self._prefix() / object_name
|
||||
# get object from bucket
|
||||
buffer = self._get(path)
|
||||
# handle if no data found
|
||||
if not buffer:
|
||||
return None
|
||||
# extract info
|
||||
class_name, buffer_checksum_str = object_name.split('-')
|
||||
# convert data
|
||||
buffer_checksum = HexadecimalString(buffer_checksum_str)
|
||||
# instantiate data
|
||||
data = ModelData(
|
||||
buffer=buffer,
|
||||
buffer_checksum=buffer_checksum,
|
||||
class_name=class_name,
|
||||
)
|
||||
return data
|
||||
|
||||
def put_data(self, data: ModelData) -> None:
|
||||
"""Put model data."""
|
||||
assert isinstance(data, ModelData)
|
||||
# build object name
|
||||
object_name = self._build_object_name(data)
|
||||
# build path
|
||||
path = self._prefix() / object_name
|
||||
# put object in bucket
|
||||
self._put(path, data.buffer)
|
||||
|
||||
def remove_data(self, object_name: str) -> None:
|
||||
"""Remove model data."""
|
||||
assert isinstance(object_name, str)
|
||||
# build path
|
||||
path = self._prefix() / object_name
|
||||
# remove object
|
||||
self._delete(path)
|
||||
|
||||
def list_names(self) -> list[str]:
|
||||
"""List names of all models."""
|
||||
# build path
|
||||
path = self._prefix()
|
||||
# list object paths
|
||||
obj_path_list = self._list_objects(path)
|
||||
# strip prefix
|
||||
name_list = [obj_path.split('/')[-1] for obj_path in obj_path_list]
|
||||
return name_list
|
||||
@@ -0,0 +1,56 @@
|
||||
"""Definition of VisualCommunicationRepository class."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from .dto import VisualCommunicationData
|
||||
from .implementations import MongoImplementation
|
||||
from .interfaces import VisualCommunicationInterface
|
||||
|
||||
|
||||
class VisualCommunicationRepository(VisualCommunicationInterface, MongoImplementation):
|
||||
"""Visual communication repository class that handles CRUD functionality
|
||||
for VisualCommunicationData."""
|
||||
|
||||
def __enter__(self) -> VisualCommunicationRepository:
|
||||
self.connect()
|
||||
return self
|
||||
|
||||
def get_data(self, name: str) -> VisualCommunicationData | None:
|
||||
"""Get visual communication data."""
|
||||
assert isinstance(name, str)
|
||||
assert len(name) > 0
|
||||
# build query
|
||||
query = {'name': name}
|
||||
# get document from mongo
|
||||
doc = self._get(query)
|
||||
# handle if no data found
|
||||
if doc is None:
|
||||
return None
|
||||
# instantiate object
|
||||
data = VisualCommunicationData(**doc)
|
||||
return data
|
||||
|
||||
def put_data(self, data: VisualCommunicationData) -> None:
|
||||
"""Put visual communication data."""
|
||||
assert isinstance(data, VisualCommunicationData)
|
||||
# convert to dict
|
||||
data_dict: dict = data.model_dump(mode='json')
|
||||
# build query
|
||||
query = {'name': data.name}
|
||||
# save to mongo
|
||||
self._save(data_dict, query)
|
||||
|
||||
def remove_data(self, name: str) -> None:
|
||||
"""Remove visual communication data."""
|
||||
assert isinstance(name, str)
|
||||
assert len(name) > 0
|
||||
# build query
|
||||
query = {'name': name}
|
||||
# remove document from mongo
|
||||
self._delete(query)
|
||||
|
||||
def list_names(self) -> list[str]:
|
||||
"""List names of all documents."""
|
||||
# list documents
|
||||
name_list = self._list_documents(key='name')
|
||||
return name_list
|
||||
@@ -0,0 +1,403 @@
|
||||
"""Integration tests configuration."""
|
||||
|
||||
import os
|
||||
import random
|
||||
from collections.abc import Iterator
|
||||
from io import BytesIO
|
||||
from pathlib import Path
|
||||
|
||||
import minio
|
||||
import pytest
|
||||
from dotenv import load_dotenv
|
||||
from PIL import Image
|
||||
from pymongo import MongoClient
|
||||
|
||||
from model.src.models import VisualCommunicationModel
|
||||
from shared.repositories import (
|
||||
HexadecimalString,
|
||||
ImageData,
|
||||
ImageRepository,
|
||||
ModelData,
|
||||
ModelRepository,
|
||||
VisualCommunicationData,
|
||||
VisualCommunicationRepository,
|
||||
VisualCommunicationValues,
|
||||
)
|
||||
from shared.repositories.src.implementations import (
|
||||
MinioImplementation,
|
||||
MongoImplementation,
|
||||
)
|
||||
|
||||
# set random seed for reproducibility
|
||||
random.seed(13)
|
||||
|
||||
# define test environment variables
|
||||
necessary_env_vars = {
|
||||
'MINIO_ENDPOINT',
|
||||
'MINIO_ACCESS_KEY',
|
||||
'MINIO_SECRET_KEY',
|
||||
'MONGO_ENDPOINT',
|
||||
}
|
||||
env_var_map = {
|
||||
'MINIO_BUCKET_NAME': 'test-bucket',
|
||||
'MINIO_OBJECT_NAME': 'test-object',
|
||||
'MINIO_IMAGE_NAME': 'test-image',
|
||||
'MONGO_DB': 'test-db',
|
||||
'MONGO_COLLECTION': 'test-collection',
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture(scope='session', autouse=True)
|
||||
def setup_env(
|
||||
request: pytest.FixtureRequest,
|
||||
) -> None:
|
||||
"""Populate environment with variables used for testing."""
|
||||
# load in optional local test environment variables
|
||||
test_env_path = Path(__file__).parent.parent.parent.parent.parent / 'test.env'
|
||||
load_dotenv(test_env_path)
|
||||
# check if necessary env vars are set
|
||||
for key in necessary_env_vars:
|
||||
assert key in os.environ, f'{key} not set'
|
||||
# set env vars unique to this test
|
||||
for key, val in env_var_map.items():
|
||||
# set env var
|
||||
os.environ[key] = val
|
||||
|
||||
# ensure cleanup
|
||||
def cleanup_env():
|
||||
for key in env_var_map:
|
||||
_ = os.environ.pop(key, default=None)
|
||||
|
||||
request.addfinalizer(cleanup_env)
|
||||
|
||||
|
||||
@pytest.fixture(scope='session')
|
||||
def raw_minio_client(
|
||||
setup_env,
|
||||
) -> Iterator[minio.Minio]:
|
||||
"""Raw Minio client fixture."""
|
||||
# prepare arguments
|
||||
minio_endpoint = str(os.getenv('MINIO_ENDPOINT'))
|
||||
minio_access_key = str(os.getenv('MINIO_ACCESS_KEY'))
|
||||
minio_secret_key = str(os.getenv('MINIO_SECRET_KEY'))
|
||||
minio_bucket_name = str(os.getenv('MINIO_BUCKET_NAME'))
|
||||
# connect client
|
||||
client = minio.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)
|
||||
# expose client
|
||||
yield client
|
||||
# cleanup
|
||||
object_list = client.list_objects(minio_bucket_name, recursive=True)
|
||||
for obj in object_list:
|
||||
client.remove_object(
|
||||
bucket_name=obj.bucket_name,
|
||||
object_name=obj.object_name,
|
||||
)
|
||||
client.remove_bucket(minio_bucket_name)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def minio_client(
|
||||
setup_env,
|
||||
) -> Iterator[MinioImplementation]:
|
||||
"""MinioImplementation fixture."""
|
||||
# instantiate and connect client
|
||||
minio_client = MinioImplementation()
|
||||
minio_client.connect()
|
||||
# expose client
|
||||
yield minio_client
|
||||
# cleanup
|
||||
object_name_list = minio_client._list_objects(Path('*'))
|
||||
for name in object_name_list:
|
||||
minio_client._delete(Path(name))
|
||||
minio_client.close()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def buffer() -> Iterator[BytesIO]:
|
||||
"""Bytes buffer fixture."""
|
||||
# generate reproducible random data
|
||||
data = random.randbytes(n=2**21) # 2 MB
|
||||
# convert data
|
||||
buffer = BytesIO(data)
|
||||
# expose buffer
|
||||
yield buffer
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def buffer_in_minio(
|
||||
raw_minio_client: minio.Minio,
|
||||
buffer: BytesIO,
|
||||
) -> Iterator[tuple[Path, BytesIO]]:
|
||||
"""Buffer in Minio fixture."""
|
||||
# prepare arguments
|
||||
minio_bucket_name = str(os.getenv('MINIO_BUCKET_NAME'))
|
||||
minio_object_name = str(os.getenv('MINIO_OBJECT_NAME'))
|
||||
# prepare for saving
|
||||
num_bytes = len(buffer.getvalue())
|
||||
buffer.seek(0)
|
||||
# put data in bucket
|
||||
raw_minio_client.put_object(
|
||||
bucket_name=minio_bucket_name,
|
||||
object_name=minio_object_name,
|
||||
length=num_bytes,
|
||||
data=buffer,
|
||||
)
|
||||
# expose data
|
||||
yield Path(minio_object_name), buffer
|
||||
# cleanup
|
||||
raw_minio_client.remove_object(
|
||||
bucket_name=minio_bucket_name,
|
||||
object_name=minio_object_name,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def image_data() -> Iterator[ImageData]:
|
||||
"""Image data fixture."""
|
||||
# prepare arguments
|
||||
name = str(os.getenv('MINIO_IMAGE_NAME'))
|
||||
image = Image.new(mode='RGB', size=(480, 480))
|
||||
# instantiate data
|
||||
image_data = ImageData(image=image, name=name)
|
||||
# expose data
|
||||
yield image_data
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def image_data_in_minio(
|
||||
raw_minio_client: minio.Minio,
|
||||
image_data: ImageData,
|
||||
) -> Iterator[ImageData]:
|
||||
"""Image data in Minio fixture."""
|
||||
# prepare arguments
|
||||
minio_bucket_name = str(os.getenv('MINIO_BUCKET_NAME'))
|
||||
image_name = image_data.name
|
||||
# build object path
|
||||
object_path = ImageRepository._build_path(image_name)
|
||||
# save image to buffer
|
||||
buffer = BytesIO()
|
||||
image_data.image.save(buffer, 'png')
|
||||
# prepare for saving
|
||||
num_bytes = len(buffer.getvalue())
|
||||
buffer.seek(0)
|
||||
# put data in bucket
|
||||
raw_minio_client.put_object(
|
||||
bucket_name=minio_bucket_name,
|
||||
object_name=object_path.as_posix(),
|
||||
length=num_bytes,
|
||||
data=buffer,
|
||||
)
|
||||
# expose data
|
||||
yield image_data
|
||||
# cleanup
|
||||
raw_minio_client.remove_object(
|
||||
bucket_name=minio_bucket_name,
|
||||
object_name=object_path.as_posix(),
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(scope='session')
|
||||
def image_repo(
|
||||
setup_env,
|
||||
) -> Iterator[ImageRepository]:
|
||||
"""Image repository fixture."""
|
||||
# define test environment variables
|
||||
assert 'MINIO_ENDPOINT' in os.environ, 'MINIO_ENDPOINT not set'
|
||||
assert 'MINIO_ACCESS_KEY' in os.environ, 'MINIO_ACCESS_KEY not set'
|
||||
assert 'MINIO_SECRET_KEY' in os.environ, 'MINIO_SECRET_KEY not set'
|
||||
assert 'MONGO_ENDPOINT' in os.environ, 'MONGO_ENDPOINT not set'
|
||||
repo = ImageRepository()
|
||||
repo.connect()
|
||||
yield repo
|
||||
repo.close()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def model_data() -> Iterator[ModelData]:
|
||||
"""Model data fixture."""
|
||||
# prepare arguments
|
||||
vis_com_model = VisualCommunicationModel().to('cpu')
|
||||
class_name = type(vis_com_model).__name__
|
||||
buffer = ModelData.model_to_buffer(vis_com_model)
|
||||
buffer_checksum = HexadecimalString('77dcab1769563654a6e24f92d40f29bd')
|
||||
# instantiate data
|
||||
model_data = ModelData(
|
||||
buffer=buffer,
|
||||
buffer_checksum=buffer_checksum,
|
||||
class_name=class_name,
|
||||
)
|
||||
# expose model
|
||||
yield model_data
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def model_data_in_minio(
|
||||
raw_minio_client: minio.Minio,
|
||||
model_data: ModelData,
|
||||
) -> Iterator[ModelData]:
|
||||
"""Model data in Minio fixture."""
|
||||
# prepare arguments
|
||||
minio_bucket_name = str(os.getenv('MINIO_BUCKET_NAME'))
|
||||
object_name = ModelRepository._build_object_name(model_data)
|
||||
buffer = model_data.buffer
|
||||
# build object path
|
||||
object_path = ModelRepository._prefix() / object_name
|
||||
# prepare for saving
|
||||
num_bytes = len(buffer.getvalue())
|
||||
buffer.seek(0)
|
||||
# put data in bucket
|
||||
raw_minio_client.put_object(
|
||||
bucket_name=minio_bucket_name,
|
||||
object_name=object_path.as_posix(),
|
||||
length=num_bytes,
|
||||
data=buffer,
|
||||
)
|
||||
# expose data
|
||||
yield model_data
|
||||
# cleanup
|
||||
raw_minio_client.remove_object(
|
||||
bucket_name=minio_bucket_name,
|
||||
object_name=object_path.as_posix(),
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(scope='session')
|
||||
def model_repo(
|
||||
setup_env,
|
||||
) -> Iterator[ModelRepository]:
|
||||
"""Model repository fixture."""
|
||||
repo = ModelRepository()
|
||||
repo.connect()
|
||||
yield repo
|
||||
repo.close()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def raw_mongo_client(
|
||||
setup_env,
|
||||
) -> Iterator[MongoClient]:
|
||||
"""Raw mongo client fixture."""
|
||||
# prepare arguments
|
||||
mongo_endpoint = str(os.getenv('MONGO_ENDPOINT'))
|
||||
mongo_database = str(os.getenv('MONGO_DB'))
|
||||
# connect client
|
||||
client: MongoClient = MongoClient(mongo_endpoint)
|
||||
_ = client[mongo_database]
|
||||
# expose client
|
||||
yield client
|
||||
# cleanup
|
||||
client.drop_database(mongo_database)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mongo_client(
|
||||
setup_env,
|
||||
) -> Iterator[MongoImplementation]:
|
||||
"""MongoImplementation fixture."""
|
||||
# instantiate and connect client
|
||||
mongo_client = MongoImplementation()
|
||||
mongo_client.connect()
|
||||
# expose client
|
||||
yield mongo_client
|
||||
# cleanup
|
||||
mongo_client._collection.drop()
|
||||
mongo_client.close()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def dictionary() -> Iterator[dict]:
|
||||
"""Dictionary fixture."""
|
||||
# prepare data
|
||||
data = {
|
||||
'name': 'test-dictionary',
|
||||
'str_key': 'value',
|
||||
'int_key': 100,
|
||||
'float_key': 3.14,
|
||||
'list_key': [1, 2, 3],
|
||||
'dict_key': {
|
||||
'nested_key': 'nested_value',
|
||||
},
|
||||
}
|
||||
# expose data
|
||||
yield data
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def dictionary_in_mongo(
|
||||
raw_mongo_client: MongoClient,
|
||||
dictionary: dict,
|
||||
) -> Iterator[dict]:
|
||||
"""Dictionary in Mongo fixture."""
|
||||
# prepare arguments
|
||||
database = str(os.getenv('MONGO_DB'))
|
||||
collection = str(os.getenv('MONGO_COLLECTION'))
|
||||
# save data
|
||||
_ = raw_mongo_client[database][collection].insert_one(dictionary.copy())
|
||||
# expose data
|
||||
yield dictionary
|
||||
# cleanup
|
||||
raw_mongo_client[database][collection].delete_one(dictionary)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def visual_communication_values() -> Iterator[VisualCommunicationValues]:
|
||||
"""Visual communication values fixture."""
|
||||
# instantiate with random values
|
||||
visual_communication_values = VisualCommunicationValues.from_random()
|
||||
# expose values
|
||||
yield visual_communication_values
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def visual_communication_data(
|
||||
visual_communication_values: VisualCommunicationValues,
|
||||
) -> Iterator[VisualCommunicationData]:
|
||||
"""Visual communication data fixture."""
|
||||
# prepare arguments
|
||||
name = 'test-visual-communication'
|
||||
annotation = visual_communication_values
|
||||
# instantiate data
|
||||
visual_communication_data = VisualCommunicationData(
|
||||
name=name,
|
||||
annotation=annotation,
|
||||
)
|
||||
# expose data
|
||||
yield visual_communication_data
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def visual_communication_data_in_mongo(
|
||||
raw_mongo_client: MongoClient,
|
||||
visual_communication_data: VisualCommunicationData,
|
||||
) -> Iterator[VisualCommunicationData]:
|
||||
"""Visual communication data in Mongo fixture."""
|
||||
# prepare arguments
|
||||
database = str(os.getenv('MONGO_DB'))
|
||||
collection = str(os.getenv('MONGO_COLLECTION'))
|
||||
# convert data
|
||||
dictionary = visual_communication_data.model_dump(mode='dict')
|
||||
# save data
|
||||
_ = raw_mongo_client[database][collection].insert_one(dictionary.copy())
|
||||
# expose data
|
||||
yield visual_communication_data
|
||||
# cleanup
|
||||
raw_mongo_client[database][collection].delete_one(dictionary)
|
||||
|
||||
|
||||
@pytest.fixture(scope='session')
|
||||
def visual_communication_repo(
|
||||
setup_env,
|
||||
) -> Iterator[VisualCommunicationRepository]:
|
||||
"""Visual communication repository fixture."""
|
||||
repo = VisualCommunicationRepository()
|
||||
repo.connect()
|
||||
yield repo
|
||||
repo.close()
|
||||
@@ -0,0 +1,156 @@
|
||||
"""Integration tests for ImageRepository class."""
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
from PIL import Image
|
||||
|
||||
from shared.repositories import ImageData, ImageRepository
|
||||
|
||||
|
||||
def same_image(
|
||||
img_a: Image.Image,
|
||||
img_b: Image.Image,
|
||||
) -> bool:
|
||||
"""Check if two images contain the same data."""
|
||||
assert isinstance(img_a, Image.Image)
|
||||
assert isinstance(img_b, Image.Image)
|
||||
# check if images have a comparable number of channels
|
||||
if img_a.getbands() != img_b.getbands():
|
||||
return False
|
||||
# calculate pixel difference between images
|
||||
img_a_arr = np.asarray(img_a)
|
||||
img_b_arr = np.asarray(img_b)
|
||||
diff = np.subtract(img_a_arr, img_b_arr)
|
||||
if np.sum(diff) != 0:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def same_image_data(
|
||||
data_a: ImageData,
|
||||
data_b: ImageData,
|
||||
) -> bool:
|
||||
"""Check if two ImageData-objects contain the same data."""
|
||||
assert isinstance(data_a, ImageData)
|
||||
assert isinstance(data_b, ImageData)
|
||||
# compare names
|
||||
if data_a.name != data_b.name:
|
||||
return False
|
||||
# compare images
|
||||
if not same_image(data_a.image, data_b.image):
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def test_should_have_context_handler():
|
||||
"""Test that class has a working context handler implemented."""
|
||||
# ACT
|
||||
with ImageRepository() as repo:
|
||||
# ASSERT
|
||||
assert repo.connected()
|
||||
|
||||
|
||||
def test_should_get_image_data(
|
||||
image_repo: ImageRepository,
|
||||
image_data_in_minio: ImageData,
|
||||
):
|
||||
"""Test getting image data."""
|
||||
# ARRANGE
|
||||
image_name = image_data_in_minio.name
|
||||
# ACT
|
||||
received_image_data = image_repo.get_data(image_name)
|
||||
# ASSERT
|
||||
assert received_image_data is not None
|
||||
assert isinstance(received_image_data, ImageData)
|
||||
assert same_image_data(image_data_in_minio, received_image_data)
|
||||
|
||||
|
||||
def test_should_get_none_when_no_image_data(
|
||||
image_repo: ImageRepository,
|
||||
image_data: ImageData,
|
||||
):
|
||||
"""Test getting None when no data is available."""
|
||||
# ARRANGE
|
||||
image_name = image_data.name
|
||||
# ACT
|
||||
received_image_data = image_repo.get_data(image_name)
|
||||
# ASSERT
|
||||
assert received_image_data is None
|
||||
|
||||
|
||||
def test_should_delete_image_data(
|
||||
image_repo: ImageRepository,
|
||||
image_data_in_minio: ImageData,
|
||||
):
|
||||
"""Test deleting image data."""
|
||||
# ARRANGE
|
||||
image_name = image_data_in_minio.name
|
||||
# ACT
|
||||
image_repo.remove_data(image_name)
|
||||
received_image_data = image_repo.get_data(image_name)
|
||||
# ASSERT
|
||||
assert received_image_data is None
|
||||
|
||||
|
||||
def test_should_put_image_data(
|
||||
image_repo: ImageRepository,
|
||||
image_data: ImageData,
|
||||
):
|
||||
"""Test putting image data."""
|
||||
# ARRANGE
|
||||
image_name = image_data.name
|
||||
# ACT
|
||||
image_repo.put_data(image_data)
|
||||
received_image_data = image_repo.get_data(image_name)
|
||||
# ASSERT
|
||||
assert received_image_data is not None
|
||||
assert same_image_data(image_data, received_image_data)
|
||||
|
||||
|
||||
def test_should_update_image_data(
|
||||
image_repo: ImageRepository,
|
||||
image_data_in_minio: ImageData,
|
||||
):
|
||||
"""Test updating image data."""
|
||||
# ARRANGE
|
||||
updated_image_data = image_data_in_minio.model_copy(
|
||||
update={
|
||||
'image': Image.new(mode='RGB', size=(480, 480), color='white'),
|
||||
},
|
||||
)
|
||||
image_name = updated_image_data.name
|
||||
# ACT
|
||||
image_repo.put_data(updated_image_data)
|
||||
received_image_data = image_repo.get_data(image_name)
|
||||
# ASSERT
|
||||
assert not same_image_data(image_data_in_minio, updated_image_data)
|
||||
assert received_image_data is not None
|
||||
assert same_image_data(updated_image_data, received_image_data)
|
||||
|
||||
|
||||
def test_should_list_names(
|
||||
image_repo: ImageRepository,
|
||||
image_data_in_minio: ImageData,
|
||||
):
|
||||
"""Test get all image names."""
|
||||
# ARRANGE
|
||||
updated_image_data = image_data_in_minio.model_copy(
|
||||
update={
|
||||
'name': 'updated-test-image',
|
||||
},
|
||||
)
|
||||
image_repo.put_data(updated_image_data)
|
||||
expected_name_list = [
|
||||
image_data_in_minio.name,
|
||||
updated_image_data.name,
|
||||
]
|
||||
# ACT
|
||||
name_list = image_repo.list_names()
|
||||
# ASSERT
|
||||
assert len(name_list) == 2
|
||||
for name in name_list:
|
||||
assert name in expected_name_list
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
pytest.main(['-s', '-v', __file__])
|
||||
@@ -0,0 +1,133 @@
|
||||
"""Integration tests related to Minio implementation."""
|
||||
|
||||
import logging
|
||||
import os
|
||||
from io import BytesIO
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from shared.repositories.src.implementations import MinioImplementation
|
||||
|
||||
|
||||
def same_data(
|
||||
data_a: BytesIO,
|
||||
data_b: BytesIO,
|
||||
) -> bool:
|
||||
"""Check if two BytesIO-objects contain the same data."""
|
||||
assert isinstance(data_a, BytesIO)
|
||||
assert isinstance(data_b, BytesIO)
|
||||
# prepare for being read
|
||||
data_a.seek(0)
|
||||
data_b.seek(0)
|
||||
# convert to bytes
|
||||
data_a_bytes = data_a.read()
|
||||
data_b_bytes = data_b.read()
|
||||
# compare size
|
||||
if len(data_a_bytes) != len(data_b_bytes):
|
||||
logging.error(
|
||||
'data has different length: %s and %s',
|
||||
len(data_a_bytes),
|
||||
len(data_b_bytes),
|
||||
)
|
||||
return False
|
||||
# compare content
|
||||
if data_a_bytes != data_b_bytes:
|
||||
logging.error('data has different bytes')
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def test_should_connect_to_minio():
|
||||
"""Test connection to Minio."""
|
||||
# ARRANGE
|
||||
client = MinioImplementation()
|
||||
# ACT
|
||||
client.connect()
|
||||
# ASSERT
|
||||
assert client.connected()
|
||||
client.close()
|
||||
|
||||
|
||||
def test_should_have_context_handler():
|
||||
"""Test that class has a working context handler implemented."""
|
||||
# ACT
|
||||
with MinioImplementation() as client:
|
||||
# ASSERT
|
||||
assert client.connected()
|
||||
|
||||
|
||||
def test_should_get_data(
|
||||
minio_client: MinioImplementation,
|
||||
buffer_in_minio: tuple[Path, BytesIO],
|
||||
):
|
||||
"""Test getting data from Minio."""
|
||||
# ARRANGE
|
||||
path, buffer = buffer_in_minio
|
||||
# ACT
|
||||
received_buffer = minio_client._get(path)
|
||||
# ASSERT
|
||||
assert received_buffer is not None
|
||||
assert same_data(received_buffer, buffer)
|
||||
|
||||
|
||||
def test_should_get_none_when_no_data(
|
||||
minio_client: MinioImplementation,
|
||||
):
|
||||
"""Test getting None when no data is available in Minio."""
|
||||
# ARRANGE
|
||||
nonexistent_path = Path('nonexistent-object-name')
|
||||
# ACT
|
||||
received_buffer = minio_client._get(nonexistent_path)
|
||||
# ASSERT
|
||||
assert received_buffer is None
|
||||
|
||||
|
||||
def test_should_delete_data(
|
||||
minio_client: MinioImplementation,
|
||||
buffer_in_minio: tuple[Path, BytesIO],
|
||||
):
|
||||
"""Test deleting data from Minio."""
|
||||
# ARRANGE
|
||||
path, _ = buffer_in_minio
|
||||
# ACT
|
||||
minio_client._delete(path)
|
||||
# ASSERT
|
||||
received_buffer = minio_client._get(path)
|
||||
assert received_buffer is None
|
||||
|
||||
|
||||
def test_should_put_data(
|
||||
minio_client: MinioImplementation,
|
||||
buffer: BytesIO,
|
||||
):
|
||||
"""Test putting data in Minio."""
|
||||
# ARRANGE
|
||||
path = Path(os.getenv('MINIO_OBJECT_NAME', default=''))
|
||||
# ACT
|
||||
minio_client._put(path, buffer)
|
||||
received_buffer = minio_client._get(path)
|
||||
# ASSERT
|
||||
assert received_buffer is not None
|
||||
assert same_data(received_buffer, buffer)
|
||||
|
||||
|
||||
def test_should_update_data(
|
||||
minio_client: MinioImplementation,
|
||||
buffer_in_minio: tuple[Path, BytesIO],
|
||||
):
|
||||
"""Test updating data in Minio."""
|
||||
# ARRANGE
|
||||
path, buffer = buffer_in_minio
|
||||
updated_buffer = BytesIO(buffer.getvalue() + b'extra data')
|
||||
# ACT
|
||||
minio_client._put(path, updated_buffer)
|
||||
received_buffer = minio_client._get(path)
|
||||
# ASSERT
|
||||
assert not same_data(updated_buffer, buffer)
|
||||
assert received_buffer is not None
|
||||
assert same_data(received_buffer, updated_buffer)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
pytest.main(['-s', '-v', __file__])
|
||||
@@ -0,0 +1,136 @@
|
||||
"""Integration tests for ModelRepository class."""
|
||||
|
||||
from io import BytesIO
|
||||
|
||||
import pytest
|
||||
|
||||
from shared.repositories import ModelData, ModelRepository
|
||||
|
||||
|
||||
def same_buffer(
|
||||
buffer_a: BytesIO,
|
||||
buffer_b: BytesIO,
|
||||
) -> bool:
|
||||
"""Check if 2 buffers contain the same data."""
|
||||
assert isinstance(buffer_a, BytesIO)
|
||||
assert isinstance(buffer_b, BytesIO)
|
||||
# read buffers
|
||||
a_values = buffer_a.getvalue()
|
||||
b_values = buffer_b.getvalue()
|
||||
# compare length of buffers
|
||||
if len(a_values) != len(b_values):
|
||||
return False
|
||||
# compare content of buffers
|
||||
if a_values != b_values:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def same_model_data(
|
||||
data_a: ModelData,
|
||||
data_b: ModelData,
|
||||
) -> bool:
|
||||
"""Check if to ModelData-objects contain the same data."""
|
||||
assert isinstance(data_a, ModelData)
|
||||
assert isinstance(data_b, ModelData)
|
||||
# compare names
|
||||
if data_a.buffer_checksum != data_b.buffer_checksum:
|
||||
return False
|
||||
# compare buffer
|
||||
if not same_buffer(data_a.buffer, data_b.buffer):
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def test_should_have_context_handler():
|
||||
"""Test that class has a working context handler implemented."""
|
||||
# ACT
|
||||
with ModelRepository() as repo:
|
||||
# ASSERT
|
||||
assert repo.connected()
|
||||
|
||||
|
||||
def test_should_get_model_data(
|
||||
model_repo: ModelRepository,
|
||||
model_data_in_minio: ModelData,
|
||||
):
|
||||
"""Test getting model data."""
|
||||
# ARRANGE
|
||||
object_name = ModelRepository._build_object_name(model_data_in_minio)
|
||||
# ACT
|
||||
received_model_data = model_repo.get_data(object_name)
|
||||
# ASSERT
|
||||
assert received_model_data is not None
|
||||
assert isinstance(received_model_data, ModelData)
|
||||
assert same_model_data(model_data_in_minio, received_model_data)
|
||||
|
||||
|
||||
def test_should_get_none_when_no_model_data(
|
||||
model_repo: ModelRepository,
|
||||
model_data: ModelData,
|
||||
):
|
||||
"""Test getting None whne no data is available."""
|
||||
# ARRANGE
|
||||
object_name = ModelRepository._build_object_name(model_data)
|
||||
# ACT
|
||||
received_model_data = model_repo.get_data(object_name)
|
||||
# ASSERT
|
||||
assert received_model_data is None
|
||||
|
||||
|
||||
def test_should_delete_model_data(
|
||||
model_repo: ModelRepository,
|
||||
model_data_in_minio: ModelData,
|
||||
):
|
||||
"""Test deleting model data."""
|
||||
# ARRANGE
|
||||
object_name = ModelRepository._build_object_name(model_data_in_minio)
|
||||
# ACT
|
||||
model_repo.remove_data(object_name)
|
||||
received_model_data = model_repo.get_data(object_name)
|
||||
# ASSERT
|
||||
assert received_model_data is None
|
||||
|
||||
|
||||
def test_should_put_model_data(
|
||||
model_repo: ModelRepository,
|
||||
model_data: ModelData,
|
||||
):
|
||||
"""Test putting model data."""
|
||||
# ARRANGE
|
||||
object_name = ModelRepository._build_object_name(model_data)
|
||||
# ACT
|
||||
model_repo.put_data(model_data)
|
||||
received_model_data = model_repo.get_data(object_name)
|
||||
# ASSERT
|
||||
assert received_model_data is not None
|
||||
assert same_model_data(model_data, received_model_data)
|
||||
|
||||
|
||||
def test_should_list_names(
|
||||
model_repo: ModelRepository,
|
||||
model_data_in_minio: ModelData,
|
||||
):
|
||||
"""Test get all model names."""
|
||||
# ARRANGE
|
||||
new_buffer_checksum = ModelData.calculate_checksum(model_data_in_minio.buffer)
|
||||
updated_model_data = model_data_in_minio.model_copy(
|
||||
update={
|
||||
'buffer_checksum': new_buffer_checksum,
|
||||
},
|
||||
)
|
||||
model_repo.put_data(updated_model_data)
|
||||
expected_name_list = [
|
||||
ModelRepository._build_object_name(model_data_in_minio),
|
||||
ModelRepository._build_object_name(updated_model_data),
|
||||
]
|
||||
# ACT
|
||||
name_list = model_repo.list_names()
|
||||
# ASSERT
|
||||
assert len(name_list) == 2
|
||||
for name in name_list:
|
||||
assert name in expected_name_list
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
pytest.main(['-s', '-v', __file__])
|
||||
@@ -0,0 +1,113 @@
|
||||
"""Integration tests related to Mongo implementation."""
|
||||
|
||||
import pytest
|
||||
|
||||
from shared.repositories.src.implementations import MongoImplementation
|
||||
|
||||
|
||||
def test_should_connect_mongo():
|
||||
"""Test connecting to Mongo."""
|
||||
# ARRANGE
|
||||
client = MongoImplementation()
|
||||
# ACT
|
||||
client.connect()
|
||||
# ASSERT
|
||||
assert client.connected()
|
||||
client.close()
|
||||
|
||||
|
||||
def test_should_have_context_handler():
|
||||
"""Test that class has a working context handler implemented."""
|
||||
# ACT
|
||||
with MongoImplementation() as client:
|
||||
# ASSERT
|
||||
assert client.connected()
|
||||
|
||||
|
||||
def test_should_get_data(
|
||||
mongo_client: MongoImplementation,
|
||||
dictionary_in_mongo: dict,
|
||||
):
|
||||
"""Test getting data from mongo."""
|
||||
# ARRANGE
|
||||
name = dictionary_in_mongo['name']
|
||||
query = {'name': name}
|
||||
# ACT
|
||||
received_dictionary = mongo_client._get(query)
|
||||
# ASSERT
|
||||
assert received_dictionary is not None
|
||||
assert isinstance(received_dictionary, dict)
|
||||
assert received_dictionary == dictionary_in_mongo
|
||||
|
||||
|
||||
def test_should_get_none_when_no_data(
|
||||
mongo_client: MongoImplementation,
|
||||
dictionary: dict,
|
||||
):
|
||||
"""Test getting None when no data is available in mongo."""
|
||||
# ARRANGE
|
||||
name = dictionary['name']
|
||||
query = {'name': name}
|
||||
# ACT
|
||||
received_dictionary = mongo_client._get(query)
|
||||
# ASSERT
|
||||
assert received_dictionary is None
|
||||
|
||||
|
||||
def test_should_delete_data(
|
||||
mongo_client: MongoImplementation,
|
||||
dictionary_in_mongo: dict,
|
||||
):
|
||||
"""Test deleting data from mongo."""
|
||||
# ARRANGE
|
||||
name = dictionary_in_mongo['name']
|
||||
query = {'name': name}
|
||||
# ACT
|
||||
mongo_client._delete(query)
|
||||
received_dictionary = mongo_client._get(query)
|
||||
# ASSERT
|
||||
assert received_dictionary is None
|
||||
|
||||
|
||||
def test_should_save_data(
|
||||
mongo_client: MongoImplementation,
|
||||
dictionary: dict,
|
||||
):
|
||||
"""Test saving data to mongo."""
|
||||
# ARRANGE
|
||||
name = dictionary['name']
|
||||
query = {'name': name}
|
||||
# ACT
|
||||
mongo_client._save(
|
||||
data=dictionary,
|
||||
query=query,
|
||||
)
|
||||
# ASSERT
|
||||
received_dictionary = mongo_client._get(query)
|
||||
assert received_dictionary is not None
|
||||
assert isinstance(received_dictionary, dict)
|
||||
assert received_dictionary == dictionary
|
||||
|
||||
|
||||
def test_should_update_data(
|
||||
mongo_client: MongoImplementation,
|
||||
dictionary_in_mongo: dict,
|
||||
):
|
||||
"""Test updating data in mongo."""
|
||||
# ARRANGE
|
||||
name = dictionary_in_mongo['name']
|
||||
query = {'name': name}
|
||||
# ACT
|
||||
dictionary_in_mongo['str-key'] = 'updated-value'
|
||||
mongo_client._save(
|
||||
data=dictionary_in_mongo,
|
||||
query=query,
|
||||
)
|
||||
# ASSERT
|
||||
received_dictionary = mongo_client._get(query)
|
||||
assert received_dictionary is not None
|
||||
assert received_dictionary == dictionary_in_mongo
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
pytest.main(['-s', '-v', __file__])
|
||||
@@ -0,0 +1,125 @@
|
||||
"""Integration tests for VisualCommunicationRepository class."""
|
||||
|
||||
import pytest
|
||||
|
||||
from shared.repositories import (
|
||||
VisualCommunicationData,
|
||||
VisualCommunicationRepository,
|
||||
VisualCommunicationValues,
|
||||
)
|
||||
|
||||
|
||||
def test_should_have_context_handler():
|
||||
"""Test that class has a working context handler implemented."""
|
||||
# ACT
|
||||
with VisualCommunicationRepository() as repo:
|
||||
# ASSERT
|
||||
assert repo.connected()
|
||||
|
||||
|
||||
def test_should_get_visual_communication_data(
|
||||
visual_communication_repo: VisualCommunicationRepository,
|
||||
visual_communication_data_in_mongo: VisualCommunicationData,
|
||||
):
|
||||
"""Test getting visual communication data."""
|
||||
# ARRANGE
|
||||
name = visual_communication_data_in_mongo.name
|
||||
# ACT
|
||||
received_visual_communication_data = visual_communication_repo.get_data(name)
|
||||
# ASSERT
|
||||
assert received_visual_communication_data is not None
|
||||
assert isinstance(received_visual_communication_data, VisualCommunicationData)
|
||||
assert received_visual_communication_data == visual_communication_data_in_mongo
|
||||
|
||||
|
||||
def test_should_get_none_when_no_visual_communication_data(
|
||||
visual_communication_repo: VisualCommunicationRepository,
|
||||
visual_communication_data: VisualCommunicationData,
|
||||
):
|
||||
"""Test getting None when no data is available."""
|
||||
# ARRANGE
|
||||
name = visual_communication_data.name
|
||||
# ACT
|
||||
received_visual_communication_data = visual_communication_repo.get_data(name)
|
||||
# ASSERT
|
||||
assert received_visual_communication_data is None
|
||||
|
||||
|
||||
def test_should_delete_image_data(
|
||||
visual_communication_repo: VisualCommunicationRepository,
|
||||
visual_communication_data_in_mongo: VisualCommunicationData,
|
||||
):
|
||||
"""Test deleting visual communication data."""
|
||||
# ARRANGE
|
||||
name = visual_communication_data_in_mongo.name
|
||||
# ACT
|
||||
visual_communication_repo.remove_data(name)
|
||||
received_visual_communication_data = visual_communication_repo.get_data(name)
|
||||
# ASSERT
|
||||
assert received_visual_communication_data is None
|
||||
|
||||
|
||||
def test_should_put_visual_communication_data(
|
||||
visual_communication_repo: VisualCommunicationRepository,
|
||||
visual_communication_data: VisualCommunicationData,
|
||||
):
|
||||
"""Test putting visual communication data."""
|
||||
# ARRANGE
|
||||
name = visual_communication_data.name
|
||||
# ACT
|
||||
visual_communication_repo.put_data(visual_communication_data)
|
||||
received_visual_communication_data = visual_communication_repo.get_data(name)
|
||||
# ASSERT
|
||||
assert received_visual_communication_data is not None
|
||||
assert isinstance(received_visual_communication_data, VisualCommunicationData)
|
||||
assert received_visual_communication_data == visual_communication_data
|
||||
|
||||
|
||||
def test_should_update_visual_communication_data(
|
||||
visual_communication_repo: VisualCommunicationRepository,
|
||||
visual_communication_data_in_mongo: VisualCommunicationData,
|
||||
):
|
||||
"""Test updating visual communication data."""
|
||||
# ARRANGE
|
||||
name = visual_communication_data_in_mongo.name
|
||||
updated_annotation = VisualCommunicationValues.from_random()
|
||||
updated_visual_communication = VisualCommunicationData(
|
||||
name=name,
|
||||
annotation=updated_annotation,
|
||||
)
|
||||
# ACT
|
||||
visual_communication_repo.put_data(updated_visual_communication)
|
||||
received_visual_communication_data = visual_communication_repo.get_data(name)
|
||||
# ASSERT
|
||||
assert received_visual_communication_data is not None
|
||||
assert isinstance(received_visual_communication_data, VisualCommunicationData)
|
||||
assert received_visual_communication_data == updated_visual_communication
|
||||
|
||||
|
||||
def test_should_list_names(
|
||||
visual_communication_repo: VisualCommunicationRepository,
|
||||
visual_communication_data_in_mongo: VisualCommunicationData,
|
||||
):
|
||||
"""Test listing names of all documents."""
|
||||
# ARRANGE
|
||||
name = 'test-visual-communication-2'
|
||||
annotation = VisualCommunicationValues.from_random()
|
||||
second_visual_communication = VisualCommunicationData(
|
||||
name=name,
|
||||
annotation=annotation,
|
||||
)
|
||||
visual_communication_repo.put_data(second_visual_communication)
|
||||
expected_name_list = [
|
||||
visual_communication_data_in_mongo.name,
|
||||
second_visual_communication.name,
|
||||
]
|
||||
# ACT
|
||||
name_list = visual_communication_repo.list_names()
|
||||
# ASSERT
|
||||
assert len(name_list) == 2
|
||||
for name in name_list:
|
||||
assert name in expected_name_list
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
pytest.main(['-s', '-v', __file__])
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user