##############################################################################
# File: yolov8_l_960_custom.py
# Author: 사용자
# Date: 2023-xx-xx
# Description:
#   - YOLOv8-L 기반(COCO 80클래스 체크포인트) + 실제 클래스는 2개(Worker, helmet)
#   - bbox_head 부분 파라미터 mismatch를 무시하고 백본만 로드
#   - 해상도 (640,640) → (960,960)으로 변경
#   - Mosaic 과정에서 'img' 키 누락 방지 위해 base config의
#     pre_transform, mosaic_affine_transform, last_transform를 그대로 사용
#   - 한글 주석 및 상세 설명 포함
#   - 진행 상황(퍼센트)와 학습 시간은 default_runtime.py의 ProgressBarHook, IterTimerHook 사용
##############################################################################

##############################
# 1) 베이스 config 지정
##############################
_base_ = '../yolov8/yolov8_l_syncbn_fast_8xb16-500e_coco.py'

##############################
# 2) 기본 수정 파라미터
##############################
max_epochs = 36  # 기존 그대로
data_root = '/DATA2/ltb/mmyolo/it1_1761_black_except/'
work_dir = './work_dirs/1761_black_960_01'  # 원하는 폴더명
resume = False   # 기존 체크포인트를 이어서 학습(X), 새로 시작(O)

# COCO 80클래스로 학습된 YOLOv8-L
# bbox_head 부분이 80->2로 달라 mismatch 경고가 나므로 아래에서 무시(ignor_keys) 설정
load_from = '/DATA2/ltb/mmyolo/work_dirs/dl_all_final_fusion_960_01/epoch_6.pth'

# base_lr: base config의 lr * (1/4) (기존 custom config와 동일)
base_lr = _base_.base_lr / 4

train_batch_size_per_gpu = 16
train_num_workers = 0
save_epoch_intervals = 1

##############################
# 3) 클래스 정보
##############################
class_name = ('Worker', 'helmet')
num_classes = len(class_name)
metainfo = dict(
    classes=class_name,
    palette=[(220, 20, 60), (90, 10, 90)]
)

##############################
# 4) 모델 정의 - head의 cls 채널 2개
#    bbox_head는 COCO 체크포인트와 달라 mismatch가 발생.
#    => load_config.ignore_keys로 무시
##############################
model = dict(
    bbox_head=dict(
        head_module=dict(num_classes=num_classes),
        loss_cls=dict(
            # 기존 로직: 0.5 * (num_classes /2 * 3 / _base_.num_det_layers)
            loss_weight=0.5 * (num_classes / 2 * 3 / _base_.num_det_layers)
        )
    )
)

##############################
# 5) 해상도 (960,960) 적용
#    base config의 (640,640)을 override
##############################
img_scale = (960, 960)

##############################
# 6) base config에서 가져올
#    pre_transform, mosaic_affine_transform, last_transform
#    -> KeyError('img') 방지 위해 그대로 사용
##############################
# base config에 정의된 변수 불러오기
# (yolov8_l_syncbn_fast_8xb16-500e_coco.py) 안에서 다시 _base_가 겹칠 수 있으니
# 본 예시에서는 "이미 최종 단계"에 있다고 가정.
# 혹시 "No attribute _base_.pre_transform" 에러가 나면,
#   같은 폴더의 yolov8_m_syncbn_fast_8xb16-500e_coco.py 등
#   더 깊은 base config를 참조해야 할 수 있음.
pre_transform = _base_.pre_transform
mosaic_affine_transform = _base_.mosaic_affine_transform
last_transform = _base_.last_transform

# mosaic_affine_transform 내부의 Mosaic, RandomAffine img_scale을 (960,960)으로 재정의
# 아래는 기존 base에서 가져온 리스트를 deepcopy 후 수정할 수도 있음
mosaic_affine_transform_960 = []
for t in mosaic_affine_transform:
    # Mosaic 변환인 경우, img_scale 수정
    if t['type'] == 'Mosaic':
        new_t = t.copy()
        new_t['img_scale'] = img_scale
        mosaic_affine_transform_960.append(new_t)
    # YOLOv5RandomAffine인 경우, border 값도 (960//2) 로 변경
    elif t['type'] == 'YOLOv5RandomAffine':
        new_t = t.copy()
        new_t['border'] = (-img_scale[0] // 2, -img_scale[1] // 2)
        mosaic_affine_transform_960.append(new_t)
    else:
        # 그대로
        mosaic_affine_transform_960.append(t)

##############################
# 7) Train Pipeline 재설정
#    - 기존 base config 파이프라인을 override
#    - Mosaic 등에서 (960,960) 사용
##############################
train_pipeline = [
    *pre_transform,              # LoadImageFromFile, LoadAnnotations 등
    *mosaic_affine_transform_960,  # Mosaic+RandomAffine (960,960)
    *last_transform             # HSVRandomAug, Flip, PackDetInputs 등
]

# stage2(마지막 몇 에폭)에도 (960,960) 유지
train_pipeline_stage2 = [
    *pre_transform,
    dict(type='YOLOv5KeepRatioResize', scale=img_scale),
    dict(
        type='LetterResize',
        scale=img_scale,
        allow_scale_up=True,
        pad_val=dict(img=114.0)),
    dict(
        type='YOLOv5RandomAffine',
        max_rotate_degree=0.0,
        max_shear_degree=0.0,
        scaling_ratio_range=(1 - 0.5, 1 + 0.5),
        max_aspect_ratio=100,
        border_val=(114, 114, 114)),
    *last_transform
]

##############################
# 8) DataLoader 설정
#    base config에서 times=3 (RepeatDataset) 등을 유지
##############################
train_dataloader = dict(
    batch_size=train_batch_size_per_gpu,
    num_workers=train_num_workers,
    dataset=dict(
        _delete_=True,   # base config에 정의된 내용을 완전히 덮어씀
        type='RepeatDataset',
        times=3,
        dataset=dict(
            type=_base_.dataset_type,  # 예: 'YOLOv5CocoDataset'
            data_root=data_root,
            metainfo=metainfo,
            ann_file='train.json',
            data_prefix=dict(img='train/images/'),
            filter_cfg=dict(filter_empty_gt=False, min_size=32),
            pipeline=train_pipeline
        )
    )
)

##############################
# 9) Val / Test
#    마찬가지로 (960,960)
##############################
val_dataloader = dict(
    dataset=dict(
        metainfo=metainfo,
        data_root=data_root,
        ann_file='val.json',
        data_prefix=dict(img='val/images/'),
        pipeline=[
            dict(type='LoadImageFromFile'),
            dict(type='YOLOv5KeepRatioResize', scale=img_scale),
            dict(
                type='LetterResize',
                scale=img_scale,
                allow_scale_up=False,
                pad_val=dict(img=114)),
            dict(type='LoadAnnotations', with_bbox=True, _scope_='mmdet'),
            dict(
                type='mmdet.PackDetInputs',
                meta_keys=('img_id', 'img_path', 'ori_shape', 'img_shape',
                           'scale_factor', 'pad_param'))
        ]
    )
)
test_dataloader = val_dataloader

##############################
# 10) Evaluator
##############################
val_evaluator = dict(ann_file=data_root + 'val.json')
test_evaluator = val_evaluator

##############################
# 11) Optim / Hooks
##############################
optim_wrapper = dict(optimizer=dict(lr=base_lr))

default_hooks = dict(
    checkpoint=dict(
        type='CheckpointHook',
        interval=save_epoch_intervals,
        max_keep_ckpts=5,
        save_best='auto'),
    param_scheduler=dict(max_epochs=max_epochs),
    logger=dict(type='LoggerHook', interval=10),
    # 아래처럼 progress_bar / timer를 명시해도 됨
    # progress_bar=dict(type='ProgressBarHook', interval=1),
    # timer=dict(type='IterTimerHook')
)

##############################
# 12) 커스텀 훅
#    mosaic 해제 시점(stage2)에도 (960,960) 유지
##############################
custom_hooks = [
    # Exponential Moving Average
    dict(
        type='EMAHook',
        ema_type='ExpMomentumEMA',
        momentum=0.0001,
        update_buffers=True,
        strict_load=False,
        priority=49),
    # mosaic 끄는 시점 -> 마지막 n 에폭
    dict(
        type='mmdet.PipelineSwitchHook',
        switch_epoch=max_epochs - 5,  # 예: 마지막 5 epoch에 mosaic 끄기
        switch_pipeline=train_pipeline_stage2)
]

##############################
# 13) 학습 루프 설정
##############################
train_cfg = dict(
    max_epochs=max_epochs,
    val_interval=save_epoch_intervals,
)

val_cfg = dict(type='ValLoop')
test_cfg = dict(type='TestLoop')

##############################
# 14) load_from 시 특정 키 무시
#    (COCO 80 classes -> 현재 2 classes)
#    bbox_head.head_module.cls_preds.* 항목 불러오지 않고
#    나머지만 로드
##############################
load_config = dict(
    ignore_keys=[
        'bbox_head.head_module.cls_preds.0.2.weight',
        'bbox_head.head_module.cls_preds.0.2.bias',
        'bbox_head.head_module.cls_preds.1.2.weight',
        'bbox_head.head_module.cls_preds.1.2.bias',
        'bbox_head.head_module.cls_preds.2.2.weight',
        'bbox_head.head_module.cls_preds.2.2.bias'
    ],
    strict=False
)

##############################################################################
# 최종 주석
#  - 위 load_config 설정으로 bbox_head의 최종 cls 레이어는 무시하고,
#    나머지 파라미터(backbone, neck 등)만 로드
#  - mosaic 사용 시, base config에서 정의된 pre_transform / mosaic_affine_transform
#    순서 그대로 유지해야 'img' 누락 에러가 발생하지 않음
#  - 해상도 (960,960)이라 GPU 메모리에 여유가 있는지 확인 필요
#  - 실시간(%) 진행, 시간 표시는 default_runtime.py 또는 여기 default_hooks.progress_bar 추가
##############################################################################