##############################################################################
# File: yolov8_l_960_custom.py
# Author: 사용자
# Date: 2025-01-21
# Description:
#   - YOLOv8-L 기반(COCO 80클래스 체크포인트) + 실제 클래스는 3개(Worker, Helmet, Head)
#   - bbox_head 부분 파라미터 mismatch를 무시하고 백본만 로드
#   - 해상도 (640,640) → (960,960)으로 변경
#   - Mosaic 과정에서 'img' 키 누락 방지 위해 base config의
#     pre_transform, mosaic_affine_transform, last_transform를 그대로 사용
#   - YOLOv5RandomAffine 변환으로 회전, 이동, 스케일링 적용
#   - 진행 상황(퍼센트)와 학습 시간은 default_runtime.py의 ProgressBarHook, IterTimerHook 사용
##############################################################################

##############################
# 1) 베이스 config 지정
##############################
_base_ = '../yolov8/yolov8_l_syncbn_fast_8xb16-500e_coco.py'  # 기본 YOLOv8-L 설정 파일 경로

##############################
# 2) 기본 수정 파라미터
##############################
max_epochs = 20  # 전체 학습 에폭 수
data_root = '/DATA2/ltb/mmyolo/it1_worker_helmet_1819_add_head/'  # 새로운 데이터셋 경로
work_dir = './work_dirs/it1_1819_add_head_02'  # 학습 결과를 저장할 작업 디렉토리
resume = False   # 기존 체크포인트에서 이어서 학습하지 않고 새로 시작

# COCO 80클래스로 학습된 YOLOv8-L 체크포인트 경로
# bbox_head 부분이 80->3으로 달라 mismatch 경고가 나므로 아래에서 무시(ignore_keys) 설정
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  # GPU당 배치 사이즈
train_num_workers = 0  # 데이터 로더의 워커 수
save_epoch_intervals = 1  # 몇 에폭마다 체크포인트를 저장할지 설정

##############################
# 3) 클래스 정보
##############################
class_name = ('Worker', 'helmet', 'head')  # 총 3개의 클래스
num_classes = len(class_name)  # 클래스 수

metainfo = dict(
    classes=class_name,
    palette=[
        (220, 20, 60),    # Worker 클래스 색상
        (90, 10, 90),     # Helmet 클래스 색상
        (70, 130, 180)    # Head 클래스 색상
    ]
)

##############################
# 4) 모델 정의 - head의 cls 채널 3개
#    bbox_head는 COCO 체크포인트와 달라 mismatch가 발생.
#    => load_config.ignore_keys로 무시
##############################
model = dict(
    bbox_head=dict(
        head_module=dict(
            act_cfg=dict(inplace=True, type='SiLU'),
            featmap_strides=[
                8,
                16,
                32,
            ],
            in_channels=[
                256,
                512,
                512,
            ],
            norm_cfg=dict(eps=0.001, momentum=0.03, type='BN'),
            num_classes=num_classes,
            reg_max=16,
            type='YOLOv8HeadModule',
            widen_factor=1.0),
        loss_cls=dict(
            loss_weight=0.75,         # 클래스 수 변경에 따른 loss_weight 재설정
            num_classes=num_classes,   # 클래스 수 전달
            type='CrossEntropyCustomLoss_mmyolo',  # 커스텀 손실 함수 사용
            use_sigmoid=True,        # 소프트맥스 사용
            class_weight=[1.0, 2.0, 10.0]


        )
    ),
    train_cfg=dict(
        assigner=dict(
            alpha=0.5,
            beta=6.0,
            eps=1e-09,
            num_classes=num_classes,
            topk=10,
            type='BatchTaskAlignedAssigner',
            use_ciou=True))
)

##############################
# 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에 정의된 변수 불러오기
pre_transform = _base_.pre_transform
mosaic_affine_transform = _base_.mosaic_affine_transform
last_transform = _base_.last_transform

# mosaic_affine_transform 내부의 Mosaic, YOLOv5RandomAffine img_scale을 (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 재설정
#    - 기존 base config 파이프라인을 override
#    - Mosaic 등에서 (960,960) 사용
#    - YOLOv5RandomAffine 변환 추가
##############################
train_pipeline = [
    *pre_transform,              # LoadImageFromFile, LoadAnnotations 등 기본 전처리
    dict(
        type='YOLOv5RandomAffine',
        max_rotate_degree=10.0,  # 최대 10도 회전
        max_shear_degree=5.0,    # 최대 5도 기울임
        scaling_ratio_range=(0.8, 1.2),  # 스케일링 범위 조정 (데이터 증강)
        border_val=(114, 114, 114)  # 패딩 색상
    ),
    *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=10.0,
        max_shear_degree=5.0,
        scaling_ratio_range=(0.8, 1.2),
        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,  # GPU당 배치 사이즈
    num_workers=train_num_workers,        # 데이터 로더 워커 수
    dataset=dict(
        _delete_=True,   # base config에 정의된 내용을 완전히 덮어씀
        type='RepeatDataset',
        times=3,  # 데이터셋을 3번 반복하여 학습 데이터 양 증가
        dataset=dict(
            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(
    batch_size=1,
    dataset=dict(
        type='YOLOv5CocoDataset',  # 사용자 정의 데이터셋 클래스
        data_root=data_root,        # 데이터셋 루트 경로
        metainfo=metainfo,          # 클래스 정보
        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',
    proposal_nums=(
        100,
        200,
        500,
    )
)
test_evaluator = val_evaluator  # 테스트 평가자도 검증 평가자와 동일하게 설정

##############################
# 11) Optimizer 및 Hooks 설정
##############################
optim_wrapper = dict(
    optimizer=dict(
        lr=base_lr  # 옵티마이저의 학습률 설정
    )
)

default_hooks = dict(
    checkpoint=dict(
        type='CheckpointHook',
        interval=save_epoch_intervals,  # 몇 에폭마다 체크포인트 저장
        max_keep_ckpts=10,  # 최대 저장할 체크포인트 수
        save_best='auto'  # 최적의 모델 자동 저장
    ),
    param_scheduler=dict(max_epochs=max_epochs),  # 학습 스케줄러 설정
    logger=dict(type='LoggerHook', interval=50),  # 로깅 설정 (50 iteration마다 로그 출력)
)

##############################
# 12) 커스텀 훅 설정
#    mosaic 해제 시점(stage2)에도 (960,960) 유지
##############################
custom_hooks = [
    # Exponential Moving Average (EMA) 설정
    dict(
        type='EMAHook',
        ema_type='ExpMomentumEMA',
        momentum=0.0001,  # EMA 모멘텀 설정
        update_buffers=True,  # 버퍼 업데이트 여부
        strict_load=False,    # 엄격한 로드 여부
        priority=49
    ),
    # mosaic 끄는 시점 설정 -> 마지막 5 에폭에 파이프라인 변경
    dict(
        type='mmdet.PipelineSwitchHook',
        switch_epoch=max_epochs - 5,  # 마지막 5 에폭에 파이프라인 변경
        switch_pipeline=train_pipeline_stage2  # 변경할 파이프라인
    )
]

##############################
# 13) 학습 루프 설정
##############################
train_cfg = dict(
    type='EpochBasedTrainLoop',  # Epoch 기반 학습 루프 사용
    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 -> 현재 3 classes)
#    bbox_head.head_module.cls_preds.* 항목 불러오지 않고
#    나머지만 로드
##############################
load_config = dict(
    ignore_keys=[
        # 각 detection layer마다 클래스 예측 파라미터 무시
        'bbox_head.head_module.cls_preds.0.weight',  # 첫 번째 detection layer의 cls_preds 무시 (3 클래스)
        'bbox_head.head_module.cls_preds.0.bias',
        'bbox_head.head_module.cls_preds.1.weight',  # 두 번째 detection layer의 cls_preds 무시
        'bbox_head.head_module.cls_preds.1.bias',
        'bbox_head.head_module.cls_preds.2.weight',  # 세 번째 detection layer의 cls_preds 무시
        'bbox_head.head_module.cls_preds.2.bias',
        # 추가로 필요한 경우 더 많은 cls_preds 키를 추가
    ],
    strict=False  # 엄격한 로드 여부 설정
)

##############################################################################
# 최종 주석
#  - 위 load_config 설정으로 bbox_head의 cls_preds 레이어는 무시하고,
#    나머지 파라미터(backbone, neck 등)만 로드
#  - mosaic 사용 시, base config에서 정의된 pre_transform / mosaic_affine_transform
#    순서 그대로 유지해야 'img' 누락 에러가 발생하지 않음
#  - 해상도 (960,960)이라 GPU 메모리에 여유가 있는지 확인 필요
#  - 실시간(%) 진행, 시간 표시는 default_runtime.py 또는 여기 default_hooks.progress_bar 추가
#  - 새로운 클래스 'Head'는 데이터가 적으므로 데이터 증강을 통해 보완
#  - 데이터셋 경로가 '/DATA2/ltb/mmyolo/it1_worker_helmet_1819_add_head/'로 변경됨
##############################################################################