from yacs.config import CfgNode as CN


cfg = CN()
cfg.SEED = -1
cfg.DEVICE = 'cpu'
cfg.MODE = 'hm'
cfg.BATCH_SIZE = 1
cfg.NUM_WORKERS = 2
cfg.PIN_MEMORY = False
cfg.NUM_EPOCHS = 1
cfg.PRINT_FREQ = 1

cfg.DATASET = CN()
cfg.DATASET.TRAIN = ''
cfg.DATASET.TEST = ''
cfg.DATASET.NUM_KEYPOINTS = 6
cfg.DATASET.NUM_LIMBS = 5
cfg.DATASET.MAX_NUM_DETECTIONS = 10
cfg.DATASET.SIGMA = 2.0
cfg.DATASET.FLIP = 0.5
cfg.DATASET.TRANSLATE = 0.4
cfg.DATASET.SCALE = [0.75, 1.25]
cfg.DATASET.ROTATION = 30
cfg.DATASET.INPUT_SIZE = [384, 384]
cfg.DATASET.OUTPUT_SIZE = [96, 96]
cfg.DATASET.MEAN = (0.485, 0.456, 0.406)
cfg.DATASET.STD = (0.229, 0.224, 0.225)

cfg.MODEL = CN()
cfg.MODEL.NUM_SCALES = 1
cfg.MODEL.FINAL_KERNEL = 1
cfg.MODEL.BACKBONE = 'resnet18'
cfg.MODEL.PRETRAINED = True

cfg.LOSS = CN()
cfg.LOSS.USE_WEIGHT = False

cfg.TRAIN = CN()
cfg.TRAIN.OPTIMIZER = 'adam'
cfg.TRAIN.LR = 1e-3
cfg.TRAIN.WD = 1e-4
cfg.TRAIN.MOMENTUM = 0.0
cfg.TRAIN.LR_FACTOR = 0.0
cfg.TRAIN.LR_STEPS = []

cfg.TEST = CN()
cfg.TEST.KEEP_RATIO = False
cfg.TEST.PROJECT = False
cfg.TEST.NMS_KERNEL = 3
cfg.TEST.DETECTION_THRESHOLD = 0.1
cfg.TEST.LIMB = CN()
cfg.TEST.LIMB.NUM_MIDPOINTS = 10
cfg.TEST.LIMB.THRESHOLD = 0.05
cfg.TEST.LIMB.IGNORE_FEW_PARTS = True
cfg.TEST.LIMB.CONNECTION_RATIO = 0.8
cfg.TEST.LIMB.LENGTH_RATE = 16
cfg.TEST.LIMB.CONNECTION_TOLERANCE = 0.7
cfg.TEST.LIMB.DELETE_SHARED_PARTS = False
cfg.TEST.LIMB.MIN_NUM_CONNECTED_PARTS = 3
cfg.TEST.LIMB.MIN_MEAN_SCORE = 0.2
