# ================================================================
#  YOLOv8-L  ‖  AI-Hub 공사현장 안전장비 45-class
#  • 입력  960×960 (멀티스케일 + 모자이크)
#  • epoch 100, batch 32 × (GPU 수)
# ================================================================

# ---------- ① 경로 ------------------------------------------------
DATA_ROOT = '/'
TRAIN_JSON = '/data2/ai-hub/105.공사현장_안전장비_인식_데이터/train_safety_dataset_coco_fix_broken_fix_2.json'
VAL_JSON   = '/data2/ai-hub/105.공사현장_안전장비_인식_데이터/val_safety_dataset_coco_fix.json'
LOAD_FROM = '/home/anonymous/ltb/mmyolo/mmyolo관련/yolov8_l_syncbn_fast_8xb16-500e_coco_20230217_182526-189611b6.pth'

# ---------- ② 클래스 정의 -----------------------------------------
CLASS_NAMES = (
    # 01–08  안전보호구
    'belt_on', 'belt_off', 'hook_on', 'hook_off',
    'shoe_on', 'shoe_off', 'helmet_on', 'helmet_off',
    # 09–19  중장비
    'excavator', 'payloader', 'forklift', 'dump_truck', 'concrete_mixer',
    'pump_car', 'pile_driver', 'truck', 'aerial_platform',
    'tower_crane', 'sky_lift',
    # 20–21  구조물
    'gang_form', 'al_form',
    # 22–23  사다리
    'ladder_A', 'ladder_worktable',
    # 24–29  설치물
    'switch_box', 'cover_opening', 'hazard_storage',
    'elev_falling_protection', 'hoist', 'jack_support',
    # 30–31  비계
    'pipe_scaffold', 'system_scaffold',
    # 32–45  자재·공구
    'cement_brick', 'hammer', 'drill', 'remital', 'brick_finish',
    'mixer', 'H_beam', 'cutting_machine', 'vibrator',
    'fire_extinguisher', 'welder', 'grinder',
    'hand_cart', 'spark_protection',
)
NUM_CLASSES = len(CLASS_NAMES)
CLASS_WEIGHTS = [1.0] * NUM_CLASSES
PALETTE = [
    (108, 60, 213),
    (127, 52, 97),
    (20, 41, 86),
    (180, 8, 173),
    (146, 110, 40),
    (225, 199, 203),
    (74, 47, 14),
    (212, 27, 157),
    (94, 154, 73),
    (230, 212, 62),
    (114, 131, 10),
    (45, 235, 253),
    (113, 65, 58),
    (118, 20, 133),
    (171, 134, 253),
    (14, 61, 58),
    (193, 185, 229),
    (27, 49, 35),
    (108, 236, 111),
    (92, 103, 75),
    (174, 145, 208),
    (232, 150, 1),
    (254, 58, 200),
    (12, 195, 158),
    (58, 156, 147),
    (94, 208, 44),
    (82, 80, 199),
    (152, 5, 187),
    (65, 29, 99),
    (132, 67, 171),
    (108, 22, 250),
    (40, 211, 158),
    (99, 128, 197),
    (63, 249, 37),
    (36, 185, 18),
    (249, 193, 187),
    (120, 94, 151),
    (222, 123, 74),
    (125, 139, 79),
    (137, 125, 134),
    (97, 28, 121),
    (84, 77, 13),
    (22, 136, 25),
    (5, 121, 254),
    (49, 88, 180),
    # (233, 56, 96),
    # (65, 144, 2),
    # (156, 235, 35),
    # (38, 80, 152),
    # (130, 76, 81),
    # (109, 9, 47),
    # (129, 245, 11),
    # (198, 206, 16),
    # (188, 196, 158),
    # (175, 200, 11),
    # (222, 148, 248),
    # (68, 81, 157),
    # (176, 172, 27),
    # (239, 107, 193),
    # (125, 220, 10),
    # (36, 158, 184),
    # (169, 61, 137),
    # (69, 231, 106),
    # (82, 120, 0),
    # (252, 209, 137),
    # (87, 89, 118),
    # (42, 38, 75),
    # (234, 55, 174),
    # (21, 58, 183),
    # (53, 140, 104),
    # (111, 100, 234),
    # (21, 164, 102),
    # (212, 54, 240),
    # (184, 229, 4),
    # (220, 214, 11),
    # (203, 179, 11),
    # (126, 120, 186),
    # (250, 178, 89),
    # (49, 153, 202),
    # (199, 144, 53),
    # (244, 246, 150),
    # (170, 14, 82),
    # (96, 96, 238),
    # (145, 120, 219),
    # (29, 133, 111),
    # (201, 157, 71),
    # (65, 47, 159),
    # (69, 196, 166),
    # (177, 239, 223),
    # (27, 218, 72),
    # (164, 142, 220),
    # (193, 118, 237),
    # (221, 143, 10),
    # (231, 241, 229),
    # (210, 119, 253),
    # (48, 164, 70),
    # (101, 223, 161),
    # (217, 205, 133),
    # (98, 42, 105),
    # (237, 79, 163),
    # (0, 180, 174),
    # (136, 5, 67),
    # (80, 197, 68),
    # (182, 124, 138),
    # (67, 130, 86),
]


# ---------- ③ 학습 하이퍼파라미터 -------------------------------

# [B] Training 관련
EPOCHS              = 100      # 총 학습 에포크 수
SWITCH_EPOCH        = 80      # 파이프라인 전환 기준 에포크 (예: stage2 전환)
CLOSE_MOSAIC_EPOCHS = 10       # Mosaic 기법 종료 에포크

# [C] Batch 사이즈 (DataLoader)
TRAIN_BATCH_SIZE = 16
VAL_BATCH_SIZE   = 16
TEST_BATCH_SIZE  = 16
# [E] 기타 하이퍼파라미터 및 설정
BASE_LR      = 0.00125
MIXUP_PROB   = 0.15
WEIGHT_DECAY = 0.0005
MOMENTUM     = 0.937
ALBU_PROB    = 0.01           # Albumentations 관련 transform 확률
NUM_WORKERS_TRAIN = 16
NUM_WORKERS_VAL   = 16
NUM_WORKERS_TEST  = 16

# =============================================================================
# 아래부터는 위 전역변수를 사용하여 전체 config를 구성합니다.
# =============================================================================

_backend_args = None

# ----- TTA (Test Time Augmentation) 위한 다양한 해상도 transform 설정 -----
_multiscale_resize_transforms = [
    dict(
        type='Compose',
        transforms=[
            dict(type='YOLOv5KeepRatioResize', scale=(640, 640)),
            dict(
                type='LetterResize',
                scale=(640, 640),
                allow_scale_up=False,
                pad_val=dict(img=114)
            ),
        ]
    ),
    dict(
        type='Compose',
        transforms=[
            dict(type='YOLOv5KeepRatioResize', scale=(320, 320)),
            dict(
                type='LetterResize',
                scale=(320, 320),
                allow_scale_up=False,
                pad_val=dict(img=114)
            ),
        ]
    ),
    dict(
        type='Compose',
        transforms=[
            dict(type='YOLOv5KeepRatioResize', scale=(960, 960)),
            dict(
                type='LetterResize',
                scale=(960, 960),
                allow_scale_up=False,
                pad_val=dict(img=114)
            ),
        ]
    ),
]

affine_scale = 0.9
albu_train_transforms = [
    dict(type='Blur', p=ALBU_PROB),
    dict(type='MedianBlur', p=ALBU_PROB),
    dict(type='ToGray', p=ALBU_PROB),
    dict(type='CLAHE', p=ALBU_PROB),
]
backend_args = None

base_lr = BASE_LR
batch_shapes_cfg = None

# ----- 클래스 이름과 팔레트 정의 (loss, 시각화 등에서 사용) -----
class_name = CLASS_NAMES
close_mosaic_epochs = CLOSE_MOSAIC_EPOCHS

# ----- custom_hooks: EMA, PipelineSwitchHook 등 -----
custom_hooks = [
    dict(
        type='EMAHook',
        ema_type='ExpMomentumEMA',
        momentum=0.0001,
        update_buffers=True,
        priority=49,
        strict_load=False
    ),
    dict(
        type='mmdet.PipelineSwitchHook',
        switch_epoch=SWITCH_EPOCH,  # 전역변수 사용
        switch_pipeline=[
            # 19epoch 이후에 사용할 파이프라인 예시
            dict(type='LoadImageFromFile', backend_args=None),
            dict(type='LoadAnnotations', with_bbox=True),
            dict(type='YOLOv5KeepRatioResize', scale=(960, 960)),
            dict(
                type='LetterResize',
                scale=(960, 960),
                allow_scale_up=True,
                pad_val=dict(img=114.0)
            ),
            dict(
                type='YOLOv5RandomAffine',
                max_rotate_degree=10.0,
                max_shear_degree=2.0,
                scaling_ratio_range=(0.8, 1.2),
                border_val=(114, 114, 114),
                max_aspect_ratio=100
            ),
            dict(
                type='mmdet.Albu',
                transforms=[
                    dict(type='Blur', p=ALBU_PROB),
                    dict(type='MedianBlur', p=ALBU_PROB),
                    dict(type='ToGray', p=ALBU_PROB),
                    dict(type='CLAHE', p=ALBU_PROB),
                ],
                bbox_params=dict(
                    format='pascal_voc',
                    label_fields=['gt_bboxes_labels', 'gt_ignore_flags'],
                    type='BboxParams'
                ),
                keymap=dict(img='image', gt_bboxes='bboxes'),
            ),
            dict(type='YOLOv5HSVRandomAug'),
            dict(type='mmdet.RandomFlip', prob=0.5),
            dict(
                type='mmdet.PackDetInputs',
                meta_keys=(
                    'img_id', 'img_path', 'ori_shape', 'img_shape',
                    'flip', 'flip_direction'
                )
            ),
        ]
    ),
]

data_root = DATA_ROOT  # 전역변수 적용
dataset_type = 'YOLOv5CocoDataset'
deepen_factor = 1.0

default_hooks = dict(
    checkpoint=dict(
        type='CheckpointHook',
        interval=1,
        max_keep_ckpts=10,
        save_best='auto'
    ),
    logger=dict(type='LoggerHook', interval=50),
    param_scheduler=dict(
        type='YOLOv5ParamSchedulerHook',
        scheduler_type='linear',
        max_epochs=EPOCHS,  # 전역변수 사용
        lr_factor=0.01
    ),
    sampler_seed=dict(type='DistSamplerSeedHook'),
    timer=dict(type='IterTimerHook'),
    visualization=dict(type='mmdet.DetVisualizationHook')
)

default_scope = 'mmyolo'

env_cfg = dict(
    cudnn_benchmark=True,
    dist_cfg=dict(backend='nccl'),
    mp_cfg=dict(mp_start_method='fork', opencv_num_threads=0)
)

img_scale = (640, 640)
img_scales = [(640, 640), (320, 320), (960, 960)]
last_stage_out_channels = 512

last_transform = [
    dict(
        type='mmdet.Albu',
        transforms=[
            dict(type='Blur', p=ALBU_PROB),
            dict(type='MedianBlur', p=ALBU_PROB),
            dict(type='ToGray', p=ALBU_PROB),
            dict(type='CLAHE', p=ALBU_PROB),
        ],
        bbox_params=dict(
            format='pascal_voc',
            label_fields=['gt_bboxes_labels', 'gt_ignore_flags'],
            type='BboxParams'
        ),
        keymap=dict(img='image', gt_bboxes='bboxes'),
    ),
    dict(type='YOLOv5HSVRandomAug'),
    dict(type='mmdet.RandomFlip', prob=0.5),
    dict(
        type='mmdet.PackDetInputs',
        meta_keys=(
            'img_id', 'img_path', 'ori_shape', 'img_shape',
            'scale_factor', 'pad_param'
        )
    ),
]

launcher = 'none'

load_config = dict(
    ignore_keys=[
        'bbox_head.head_module.cls_preds.0.weight',
        'bbox_head.head_module.cls_preds.0.bias',
        'bbox_head.head_module.cls_preds.1.weight',
        'bbox_head.head_module.cls_preds.1.bias',
        'bbox_head.head_module.cls_preds.2.weight',
        'bbox_head.head_module.cls_preds.2.bias',
    ],
    strict=False
)

load_from = LOAD_FROM
log_level = 'INFO'
log_processor = dict(type='LogProcessor', by_epoch=True, window_size=50)

loss_bbox_weight = 7.5
loss_cls_weight  = 1.0
loss_dfl_weight  = 0.375
lr_factor        = 0.01
max_aspect_ratio = 100
max_epochs       = EPOCHS   # 전역변수 사용
max_keep_ckpts   = 2

metainfo = dict(
    classes=CLASS_NAMES,  # 전역변수 적용
    palette=PALETTE       # 전역변수 적용
)

mixup_prob = MIXUP_PROB

# ----- 모델 구조 설정 -----
model = dict(
    type='YOLODetector',
    data_preprocessor=dict(
        type='YOLOv5DetDataPreprocessor',
        bgr_to_rgb=True,
        mean=[0.0, 0.0, 0.0],
        std=[255.0, 255.0, 255.0]
    ),
    backbone=dict(
        type='YOLOv8CSPDarknet',
        arch='P5',
        deepen_factor=deepen_factor,
        widen_factor=1.0,
        norm_cfg=dict(type='BN', momentum=0.03, eps=0.001),
        act_cfg=dict(type='SiLU', inplace=True),
        last_stage_out_channels=512
    ),
    neck=dict(
        type='YOLOv8PAFPN',
        deepen_factor=deepen_factor,
        widen_factor=1.0,
        in_channels=[256, 512, 512],
        out_channels=[256, 512, 512],
        num_csp_blocks=3,
        norm_cfg=dict(type='BN', momentum=0.03, eps=0.001),
        act_cfg=dict(type='SiLU', inplace=True)
    ),
    bbox_head=dict(
        type='YOLOv8Head',
        bbox_coder=dict(type='DistancePointBBoxCoder'),
        head_module=dict(
            type='YOLOv8HeadModule',
            num_classes=NUM_CLASSES,  # 전역변수 사용
            in_channels=[256, 512, 512],
            featmap_strides=[8, 16, 32],
            reg_max=16,
            act_cfg=dict(type='SiLU', inplace=True),
            norm_cfg=dict(type='BN', momentum=0.03, eps=0.001)
        ),
        loss_bbox=dict(
            type='IoULoss',
            bbox_format='xywh',
            iou_mode='ciou',
            loss_weight=loss_bbox_weight,
            return_iou=False,
            reduction='sum'
        ),
        loss_cls=dict(
            type='CrossEntropyCustomLoss_mmyolo',
            use_sigmoid=True,
            num_classes=NUM_CLASSES,
            reduction='none',
            loss_weight=0.75,
            class_weight=CLASS_WEIGHTS  # 전역변수 사용
        ),
        loss_dfl=dict(
            type='mmdet.DistributionFocalLoss',
            reduction='mean',
            loss_weight=loss_dfl_weight
        ),
        prior_generator=dict(
            type='mmdet.MlvlPointGenerator',
            strides=[8, 16, 32],
            offset=0.5
        )
    ),
    train_cfg=dict(
        assigner=dict(
            type='BatchTaskAlignedAssigner',
            num_classes=NUM_CLASSES,
            topk=10,
            alpha=0.5,
            beta=6.0,
            eps=1e-09,
            use_ciou=True
        )
    ),
    test_cfg=dict(
        nms=dict(type='nms', iou_threshold=0.5),
        nms_pre=30000,
        score_thr=0.01,
        max_per_img=300,
        multi_label=True
    )
)

model_test_cfg = dict(
    nms=dict(type='nms', iou_threshold=0.5),
    nms_pre=30000,
    score_thr=0.01,
    max_per_img=300,
    multi_label=True
)

# ----- Mosaic transform 설정 예시 -----
mosaic_affine_transform = [
    dict(
        type='Mosaic',
        img_scale=(640, 640),
        pad_val=114.0,
        pre_transform=[
            dict(type='LoadImageFromFile', backend_args=None),
            dict(type='LoadAnnotations', with_bbox=True)
        ]
    ),
    dict(
        type='YOLOv5RandomAffine',
        border=(-320, -320),
        border_val=(114, 114, 114),
        max_aspect_ratio=100,
        max_rotate_degree=0.0,
        max_shear_degree=0.0,
        scaling_ratio_range=(0.1, 1.9)
    ),
]

mosaic_affine_transform_960 = [
    dict(
        type='Mosaic',
        img_scale=(960, 960),
        pad_val=114.0,
        pre_transform=[
            dict(type='LoadImageFromFile', backend_args=None),
            dict(type='LoadAnnotations', with_bbox=True)
        ]
    ),
    dict(
        type='YOLOv5RandomAffine',
        border=(-480, -480),
        border_val=(114, 114, 114),
        max_aspect_ratio=100,
        max_rotate_degree=0.0,
        max_shear_degree=0.0,
        scaling_ratio_range=(0.1, 1.9)
    ),
]

norm_cfg = dict(type='BN', momentum=0.03, eps=0.001)
num_classes = NUM_CLASSES
num_det_layers = 3

# ----- Optimizer 설정 -----
optim_wrapper = dict(
    type='AmpOptimWrapper',
    constructor='YOLOv5OptimizerConstructor',
    loss_scale='dynamic',
    optimizer=dict(
        type='SGD',
        lr=BASE_LR,
        momentum=MOMENTUM,
        nesterov=True,
        weight_decay=WEIGHT_DECAY,
        batch_size_per_gpu=TRAIN_BATCH_SIZE
    ),
    clip_grad=dict(max_norm=5.0)
)

param_scheduler = None
persistent_workers = False

pre_transform = [
    dict(type='LoadImageFromFile', backend_args=None),
    dict(type='LoadAnnotations', with_bbox=True),
]

resume = False
save_epoch_intervals = 1
strides = [8, 16, 32]

tal_alpha = 0.5
tal_beta = 6.0
tal_topk = 10

test_cfg = dict(type='TestLoop')
test_dataloader = dict(
    batch_size=TEST_BATCH_SIZE,
    num_workers=NUM_WORKERS_TEST,
    persistent_workers=False,
    pin_memory=True,
    sampler=dict(type='DefaultSampler', shuffle=False),
    dataset=dict(
        type='YOLOv5CocoDataset',
        data_root=DATA_ROOT,
        ann_file=VAL_JSON,
        data_prefix=dict(img=''),
        metainfo=dict(
            classes=CLASS_NAMES,
            palette=PALETTE
        ),
        pipeline=[
            dict(type='LoadImageFromFile'),
            dict(type='YOLOv5KeepRatioResize', scale=(960, 960)),
            dict(
                type='LetterResize',
                scale=(960, 960),
                allow_scale_up=False,
                pad_val=dict(img=114)
            ),
            # dict(type='mmdet.LoadAnnotations', with_bbox=True),
            dict(
                type='mmdet.PackDetInputs',
                meta_keys=(
                    'img_id', 'img_path', 'ori_shape', 'img_shape',
                    'scale_factor', 'pad_param'
                )
            ),
        ],
        test_mode=True,
        batch_shapes_cfg=None
    )
)

test_evaluator = dict(
    type='mmdet.CocoMetric',
    ann_file= VAL_JSON,
    metric='bbox',
    proposal_nums=(100, 200, 500)
)

test_pipeline = [
    dict(type='LoadImageFromFile', backend_args=None),
    dict(type='YOLOv5KeepRatioResize', scale=(640, 640)),
    dict(
        type='LetterResize',
        scale=(640, 640),
        allow_scale_up=False,
        pad_val=dict(img=114)
    ),
    # dict(type='mmdet.LoadAnnotations', with_bbox=True),
    dict(
        type='mmdet.PackDetInputs',
        meta_keys=(
            'img_id', 'img_path', 'ori_shape', 'img_shape',
            'scale_factor', 'pad_param'
        )
    ),
]

train_ann_file = TRAIN_JSON
train_batch_size_per_gpu = TRAIN_BATCH_SIZE

train_cfg = dict(
    type='EpochBasedTrainLoop',
    max_epochs=EPOCHS,
    val_interval=1,
    dynamic_intervals=[(499, 1)]
)

train_data_prefix = ''
train_dataloader = dict(
    batch_size=TRAIN_BATCH_SIZE,
    num_workers=NUM_WORKERS_TRAIN,
    persistent_workers=False,
    pin_memory=True,
    sampler=dict(type='DefaultSampler', shuffle=True),
    collate_fn=dict(type='yolov5_collate'),
    dataset=dict(
        type='RepeatDataset',
        times=1,  # 원본+회전캔버스 데이터 반복 횟수 (데이터 균형 조절)
        dataset=dict(
            type='YOLOv5CocoDataset',
            data_root=DATA_ROOT,
            ann_file=TRAIN_JSON,
            data_prefix=dict(img=''),
            filter_cfg=dict(filter_empty_gt=False, min_size=32),
            metainfo=dict(
                classes=CLASS_NAMES,
                palette=PALETTE
            ),
            pipeline=[
                dict(type='LoadImageFromFile', backend_args=None),
                dict(type='LoadAnnotations', with_bbox=True),
                dict(
                    type='YOLOv5RandomAffine',
                    max_rotate_degree=10.0,
                    max_shear_degree=2.0,
                    scaling_ratio_range=(0.8, 1.2),
                    border_val=(114, 114, 114)
                ),
                dict(
                    type='Mosaic',
                    img_scale=(960, 960),
                    pad_val=114.0,
                    pre_transform=[
                        dict(type='LoadImageFromFile', backend_args=None),
                        dict(type='LoadAnnotations', with_bbox=True),
                    ]
                ),
                dict(
                    type='YOLOv5RandomAffine',
                    border=(-480, -480),
                    border_val=(114, 114, 114),
                    max_aspect_ratio=100,
                    max_rotate_degree=0.0,
                    max_shear_degree=0.0,
                    scaling_ratio_range=(0.1, 1.9)
                ),
                dict(
                    type='mmdet.Albu',
                    transforms=[
                        dict(type='Blur', p=ALBU_PROB),
                        dict(type='MedianBlur', p=ALBU_PROB),
                        dict(type='ToGray', p=ALBU_PROB),
                        dict(type='CLAHE', p=ALBU_PROB),
                    ],
                    bbox_params=dict(
                        type='BboxParams',
                        format='pascal_voc',
                        label_fields=['gt_bboxes_labels', 'gt_ignore_flags']
                    ),
                    keymap=dict(img='image', gt_bboxes='bboxes')
                ),
                dict(type='YOLOv5HSVRandomAug'),
                dict(type='mmdet.RandomFlip', prob=0.5),
                dict(
                    type='mmdet.PackDetInputs',
                    meta_keys=(
                        'img_id', 'img_path', 'ori_shape', 'img_shape',
                        'flip', 'flip_direction'
                    )
                ),
            ],
        )
    ),
)

train_num_workers = NUM_WORKERS_TRAIN
train_pipeline = [
    # 학습 시 pipeline은 위와 동일하게 사용합니다.
]

train_pipeline_stage2 = [
    dict(type='LoadImageFromFile', backend_args=None),
    dict(type='LoadAnnotations', with_bbox=True),
    dict(type='YOLOv5KeepRatioResize', scale=(960, 960)),
    dict(
        type='LetterResize',
        scale=(960, 960),
        allow_scale_up=True,
        pad_val=dict(img=114.0)
    ),
    dict(
        type='YOLOv5RandomAffine',
        max_rotate_degree=10.0,
        max_shear_degree=2.0,
        scaling_ratio_range=(0.8, 1.2),
        border_val=(114, 114, 114),
        max_aspect_ratio=100
    ),
    dict(
        type='mmdet.Albu',
        transforms=[
            dict(type='Blur', p=ALBU_PROB),
            dict(type='MedianBlur', p=ALBU_PROB),
            dict(type='ToGray', p=ALBU_PROB),
            dict(type='CLAHE', p=ALBU_PROB),
        ],
        bbox_params=dict(
            type='BboxParams',
            format='pascal_voc',
            label_fields=['gt_bboxes_labels', 'gt_ignore_flags']
        ),
        keymap=dict(img='image', gt_bboxes='bboxes')
    ),
    dict(type='YOLOv5HSVRandomAug'),
    dict(type='mmdet.RandomFlip', prob=0.5),
    dict(
        type='mmdet.PackDetInputs',
        meta_keys=(
            'img_id', 'img_path', 'ori_shape', 'img_shape',
            'flip', 'flip_direction'
        )
    ),
]

val_cfg = dict(type='ValLoop')
val_data_prefix = ''
val_dataloader = dict(
    batch_size=VAL_BATCH_SIZE,
    num_workers=NUM_WORKERS_VAL,
    persistent_workers=False,
    pin_memory=True,
    sampler=dict(type='DefaultSampler', shuffle=False),
    dataset=dict(
        type='YOLOv5CocoDataset',
        data_root=DATA_ROOT,
        ann_file=VAL_JSON,
        data_prefix=dict(img=''),
        metainfo=dict(
            classes=CLASS_NAMES,
            palette=PALETTE
        ),
        pipeline=[
            dict(type='LoadImageFromFile'),
            dict(type='YOLOv5KeepRatioResize', scale=(960, 960)),
            dict(
                type='LetterResize',
                scale=(960, 960),
                allow_scale_up=False,
                pad_val=dict(img=114)
            ),
            # dict(type='mmdet.LoadAnnotations', with_bbox=True),
            dict(
                type='mmdet.PackDetInputs',
                meta_keys=(
                    'img_id', 'img_path', 'ori_shape', 'img_shape',
                    'scale_factor', 'pad_param'
                )
            ),
        ],
        test_mode=True,
        batch_shapes_cfg=None
    )
)

val_evaluator = dict(
    type='mmdet.CocoMetric',
    ann_file= VAL_JSON,
    metric='bbox',
    proposal_nums=(100, 200, 500)
)

val_interval_stage2 = 1
val_num_workers = NUM_WORKERS_VAL

vis_backends = [dict(type='LocalVisBackend')]
visualizer = dict(
    type='mmdet.DetLocalVisualizer',
    name='visualizer',
    vis_backends=[dict(type='LocalVisBackend')]
)

weight_decay = WEIGHT_DECAY
widen_factor = 1.0
work_dir = 'work_dirs/ff_base_cutmix_combined_1,2,5_01'

tta_model = dict(
    type='mmdet.DetTTAModel',
    tta_cfg=dict(
        nms=dict(type='nms', iou_threshold=0.65),
        max_per_img=300
    )
)

tta_pipeline = [
    dict(type='LoadImageFromFile', backend_args=None),
    dict(
        type='TestTimeAug',
        transforms=[
            [
                dict(
                    type='Compose',
                    transforms=[
                        dict(type='YOLOv5KeepRatioResize', scale=(640, 640)),
                        dict(
                            type='LetterResize',
                            scale=(640, 640),
                            allow_scale_up=False,
                            pad_val=dict(img=114)
                        )
                    ]
                ),
                dict(
                    type='Compose',
                    transforms=[
                        dict(type='YOLOv5KeepRatioResize', scale=(320, 320)),
                        dict(
                            type='LetterResize',
                            scale=(320, 320),
                            allow_scale_up=False,
                            pad_val=dict(img=114)
                        )
                    ]
                ),
                dict(
                    type='Compose',
                    transforms=[
                        dict(type='YOLOv5KeepRatioResize', scale=(960, 960)),
                        dict(
                            type='LetterResize',
                            scale=(960, 960),
                            allow_scale_up=False,
                            pad_val=dict(img=114)
                        )
                    ]
                ),
            ],
            [
                dict(type='mmdet.RandomFlip', prob=1.0),
                dict(type='mmdet.RandomFlip', prob=0.0),
            ],
            [
                dict(type='mmdet.LoadAnnotations', with_bbox=True),
            ],
            [
                dict(
                    type='mmdet.PackDetInputs',
                    meta_keys=(
                        'img_id', 'img_path', 'ori_shape', 'img_shape',
                        'scale_factor', 'pad_param', 'flip', 'flip_direction'
                    )
                ),
            ],
        ]
    ),
]