133 lines
4.0 KiB
Python
Executable File
133 lines
4.0 KiB
Python
Executable File
"""Definition of VisualCommunication model."""
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
from base64 import b64decode
|
|
from base64 import b64encode
|
|
from io import BytesIO
|
|
from pathlib import Path
|
|
|
|
from minio import Minio
|
|
from PIL import Image
|
|
from pydantic import BaseModel
|
|
from pymongo.collection import Collection
|
|
|
|
from shared.data_store import get
|
|
from shared.data_store import 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
|
|
|
|
class Config:
|
|
"""BaseModel configuration."""
|
|
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()
|