src/constants.py

from enum import Enum
from pathlib import Path

NUMBERS_LIST = [str(i) for i in list(range(10))]
SKETCHPAD_SIZE = 50


class GradioTrainModelChoice(Enum):
    RESNET = "Original ResNet-18"
    MNIST_FINETUNED = "Resnet-18 Finetuned on MNIST subset"


class GradioPredictModelChoice(Enum):
    PRETRAINED = "Pretrained on predefined set"
    USER = "Trained on generated in session set"


class GradioTrainDataChoice(Enum):
    PREDEFINED = "Predefined set"
    USER_GENERATED = "Generated in session"


class ModelPath(Enum):
    PRETRAINED_MNIST = Path("../data/pretrained.pkl")
    # OUTPUT_MODEL_PATH = Path("../data/model.pkl")


class DataPath(Enum):
    PREGEN_DATA_PATH = Path("../data/generated")


class SessionFiles(Enum):
    USER_DATA_FOLDER = "data_finetune"
    USER_TRAINED_PKL = "user_model.pkl"