MNIST-ResNet-Demo / src /constants.py
shpr
Organize file structure
c454738
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"