##############################################################################
# File: yolov8_l_960_custom.py
# Author: 사용자
# Date: 2023-xx-xx
# Description:
#   - YOLOv8-L 기반(COCO 80클래스 체크포인트) + 실제 클래스는 2개('helmet', 'background')
#   - bbox_head 부분 파라미터 mismatch를 무시하고 백본만 로드
#   - 해상도 (640,640) → (960,960)으로 변경
#   - Mosaic 과정에서 'img' 키 누락 방지 위해 base config의
#     pre_transform, mosaic_affine_transform, last_transform를 그대로 사용
#   - 커스텀 트랜스폼(RandomHelmetPasteWithRotation) 추가하여 현장 이미지에 합성
#   - epoch 36 (마지막 5epoch엔 mosaic 해제)
##############################################################################

##############################
# 0) 커스텀 임포트
##############################
custom_imports = dict(
    imports=[
        'my_custom_pipeline.random_helmet_paste'
    ],
    allow_failed_imports=False
)

##############################
# 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

# COCO 80클래스로 학습된 YOLOv8-L
load_from = 'yolov8_l_syncbn_fast_8xb16-500e_coco_20230217_182526-189611b6.pth'

# base_lr: base config의 lr * (1/4)
base_lr = _base_.base_lr / 4

train_batch_size_per_gpu = 32
train_num_workers = 0
save_epoch_intervals = 1

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

##############################
# 4) 모델 정의 (head의 cls 채널=2)
##############################
model = dict(
    bbox_head=dict(
        head_module=dict(num_classes=num_classes),
        loss_cls=dict(
            loss_weight=0.5 * (num_classes / 2 * 3 / _base_.num_det_layers)
        )
    )
)

##############################
# 5) 해상도 (960,960) 적용
##############################
img_scale = (960, 960)

##############################
# 6) base config에서 일부 파이프라인 가져오기
##############################
pre_transform = _base_.pre_transform
mosaic_affine_transform = _base_.mosaic_affine_transform
last_transform = _base_.last_transform

# mosaic_affine_transform -> (960,960)으로 수정
mosaic_affine_transform_960 = []
for t in mosaic_affine_transform:
    if t['type'] == 'Mosaic':
        new_t = t.copy()
        new_t['img_scale'] = img_scale
        mosaic_affine_transform_960.append(new_t)
    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
##############################
train_pipeline = [
    # (1) 현장 이미지 + 바운딩박스 로드
    *pre_transform,

    # (2) 커스텀 트랜스폼: 헬멧 합성
    dict(
        type='RandomHelmetPasteWithRotation',
        helmet_dir='/DATA2/ltb/mmyolo/robo_6403_helmet_only_0117/train/images',
        paste_prob=0.5,
        label_id=0,          # 'helmet'=0
        angle_range=(-30, 30)
    ),

    # (3) Mosaic(960,960) + RandomAffine
    *mosaic_affine_transform_960,

    # (4) HSVRandomAug, Flip, Pack 등
    *last_transform
]

# stage2 (마지막 5epoch) -> Mosaic 해제
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 (train)
##############################
train_dataloader = dict(
    batch_size=train_batch_size_per_gpu,
    num_workers=train_num_workers,
    dataset=dict(
        _delete_=True,
        type='RepeatDataset',
        times=3,
        dataset=dict(
            type=_base_.dataset_type,  # 'YOLOv5CocoDataset'
            data_root=data_root,
            metainfo=metainfo,
            ann_file='train.json',   # => /DATA2/ltb/mmyolo/it1_1761_black_except/train.json
            data_prefix=dict(img='images/'),  # => 현장 이미지 폴더
            filter_cfg=dict(filter_empty_gt=False, min_size=32),
            pipeline=train_pipeline
        )
    )
)

##############################
# 9) DataLoader (val/test)
##############################
val_dataloader = dict(
    dataset=dict(
        metainfo=metainfo,
        data_root=data_root,
        ann_file='val.json',
        data_prefix=dict(img='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)
)

##############################
# 12) 커스텀 훅
#     마지막 5 epoch에 Mosaic 해제
##############################
custom_hooks = [
    dict(
        type='EMAHook',
        ema_type='ExpMomentumEMA',
        momentum=0.0001,
        update_buffers=True,
        strict_load=False,
        priority=49
    ),
    dict(
        type='mmdet.PipelineSwitchHook',
        switch_epoch=max_epochs - 5,  # 마지막 5epoch에 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 -> 현재 2 classes)
##############################
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
)

##############################################################################
# 최종 요약:
# 1) 현장 이미지 + 어노테이션:
#    - data_root=/DATA2/ltb/mmyolo/it1_1761_black_except/
#    - train.json, val.json은 여기 위치
#    - 실제 이미지는 images/ 폴더 내부
# 2) 합성용 헬멧 PNG:
#    - /DATA2/ltb/mmyolo/roboflow_helmet_dataset/train/images
#    - RandomHelmetPasteWithRotation(helmet_dir=...)에서 불러와
#      현장 이미지에 알파블렌딩+회전
# 3) stage2 시점(마지막 5epoch)에 mosaic 해제
# 4) 클래스: ('helmet', 'background')
##############################################################################