src/session.py

import re
import uuid
from pathlib import Path

from data_manager import ImageData
from model import PredictionModel

HTML_CLEANER = re.compile("<.*?>")


class Session:
    def __init__(
        self,
        name: str,
        finetune_data: ImageData,
        predict_data: ImageData,
        pregen_data: ImageData,
        model: PredictionModel,
    ) -> None:
        self.name = name
        self.id = uuid.uuid4()
        self.path = Path("../data/sessions") / str(self.id)
        self.finetune_data = finetune_data
        self.predict_data = predict_data
        self.pregen_data = pregen_data
        self.model = model
        self._set_paths()

    def _set_paths(self):

        self.finetune_data.path = self.path
        self.predict_data.path = self.path
        self.model.session_path = self.path

    def load_local_storage(self, local_storage_id: str):
        session_id = re.sub(HTML_CLEANER, "", local_storage_id).rstrip()
        self.id = uuid.UUID(session_id)
        self.path = Path("../data/sessions") / session_id
        self._set_paths()