base.py

from enum import Enum

class TrainingObjective(Enum):
    CLASSIFICATION = 0
    REGRESSION = 1

class ModelBackbone(Enum):
    INCEPTION_V3 = 0
    EFFICIENT_NET_B0 = 1
    MOBILE_NET_V3_LARGE = 2


class TransformationType(Enum):
    RESIZE = 0
    TO_TENSOR = 1
    RANDOM_HORIZONTAL_FLIP = 2
    NORMALIZE = 3
    RANDOM_ROTATION = 4
    RANDOM_CLIP = 5

model_input_dict = {
    ModelBackbone.INCEPTION_V3: (299, 299),
    ModelBackbone.EFFICIENT_NET_B0: (224, 224),
    ModelBackbone.MOBILE_NET_V3_LARGE: (224, 224)
}