86 lines
2.6 KiB
Python
Executable File
86 lines
2.6 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 PIL import Image
|
|
from pydantic import BaseModel
|
|
from pydantic import field_serializer
|
|
from pydantic import field_validator
|
|
|
|
from shared.dto import ModelData
|
|
|
|
|
|
class VisualCommunication(BaseModel):
|
|
"""Visual communication model."""
|
|
|
|
name: str
|
|
image: Image.Image
|
|
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 from_file(cls, path: Path) -> VisualCommunication:
|
|
"""Instantiate from file."""
|
|
name = path.stem
|
|
image = Image.open(path)
|
|
image.load()
|
|
return VisualCommunication(name=name, image=image)
|
|
|
|
@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)))
|
|
|
|
@field_serializer('image')
|
|
@classmethod
|
|
def serialize_image(cls, image: Image.Image) -> bytes: # type: ignore
|
|
"""Convert image to bytes for storage in database."""
|
|
buffer = BytesIO()
|
|
image.save(buffer, format='JPEG')
|
|
return buffer.getvalue()
|
|
|
|
@field_validator('image', mode='before')
|
|
@classmethod
|
|
def convert_to_image(
|
|
cls,
|
|
image: Image.Image | BytesIO | bytes,
|
|
) -> Image.Image:
|
|
"""Convert bytes input from database into image."""
|
|
if isinstance(image, bytes):
|
|
image = BytesIO(image)
|
|
if isinstance(image, BytesIO):
|
|
image = Image.open(image)
|
|
return image
|
|
|
|
def __repr__(self) -> str:
|
|
return f"{self.classname()}(name='{self.name}')"
|
|
|
|
def webencoded_image(self) -> str:
|
|
"""Convert image to be displayed on webpage."""
|
|
# convert images to bytes string
|
|
buffer = BytesIO()
|
|
self.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()
|