##############################################################################
# File: yolov8_l_960_head_bg.py
# Date: 2025-01-17
# Description:
#   - 2클래스 ('head', 'background') 선언
#   - 기존 어노테이션(Worker/helmet)은 모두 제거
#   - Custom Transform(RemoveBboxesAndLabels) + RandomHelmetPasteWithRotation(label_id=0)
##############################################################################

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


##############################
# 1) Import & base config
##############################
_base_ = '../yolov8/yolov8_l_syncbn_fast_8xb16-500e_coco.py'
from mmyolo.registry import TRANSFORMS

max_epochs = 36
data_root = '/DATA2/ltb/mmyolo/it1_1761_black_except/'
work_dir = './work_dirs/1761_black_960_head_bg'
resume = False

load_from = 'yolov8_l_syncbn_fast_8xb16-500e_coco_20230217_182526-189611b6.pth'
base_lr = _base_.base_lr / 4

train_batch_size_per_gpu = 32
train_num_workers = 0
save_epoch_intervals = 1

##############################
# 2) 클래스: ('head', 'background')
##############################
class_name = ('head', 'background')
num_classes = len(class_name)
metainfo = dict(
    classes=class_name,
    palette=[(220, 20, 60), (0, 255, 0)]  # 예시 팔레트
)

##############################
# 3) 모델 정의
##############################
model = dict(
    bbox_head=dict(
        head_module=dict(num_classes=num_classes),
        loss_cls=dict(
            # base config에서 사용하는 공식에 맞춰서
            loss_weight=0.5 * (num_classes / 2 * 3 / _base_.num_det_layers)
        )
    )
)

##############################
# 4) 해상도 (960,960)
##############################
img_scale = (960, 960)

##############################
# 5) base config 파이프라인 가져오기
##############################
pre_transform = _base_.pre_transform
mosaic_affine_transform = _base_.mosaic_affine_transform
last_transform = _base_.last_transform

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)

##############################
# 6) 커스텀 트랜스폼 (RemoveBboxesAndLabels)
#    + RandomHelmetPasteWithRotation(label_id=0)
##############################
train_pipeline = [
    # (1) 현장 이미지 + (Worker/helmet) 바운딩박스 로드
    *pre_transform,

    # (2) 기존 바운딩박스 제거
    dict(type='RemoveBboxesAndLabels'),  # ← 우리가 만든 커스텀 (아래에 코드 예시)

    # (3) head 합성 (helmet_dir에 .jpg도 있을 경우 확장자 패턴 추가해야 함)
    dict(
        type='RandomHelmetPasteWithRotation',
        helmet_dir='/DATA2/ltb/mmyolo/robo_head_only_1561_250117/train/images',
        paste_prob=0.5,
        label_id=0,  # head=0
        angle_range=(-30, 30)
    ),

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

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

# stage2 (마지막 5epoch) -> Mosaic 해제
train_pipeline_stage2 = [
    *pre_transform,
    dict(type='RemoveBboxesAndLabels'),  # stage2에도 동일하게 바운딩박스 제거
    dict(
        type='RandomHelmetPasteWithRotation',
        helmet_dir='/DATA2/ltb/mmyolo/robo_head_only_1561_250117/train/images',
        paste_prob=0.5,
        label_id=0,
        angle_range=(-30, 30)
    ),
    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
]

##############################
# 7) DataLoader 설정
##############################
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',
            data_prefix=dict(img='images/'),
            filter_cfg=dict(filter_empty_gt=False, min_size=32),
            pipeline=train_pipeline
        )
    )
)



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

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

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

##############################
# 10) 커스텀 훅
##############################
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,
        switch_pipeline=train_pipeline_stage2
    )
]

##############################
# 11) 학습 루프 설정
##############################
train_cfg = dict(
    max_epochs=max_epochs,
    val_interval=save_epoch_intervals
)
val_cfg = dict(type='ValLoop')
test_cfg = dict(type='TestLoop')

##############################
# 12) load_from 시 특정 키 무시
##############################
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) class_name=('head', 'background') => num_classes=2
# 2) RemoveBboxesAndLabels: Worker/helmet 바운딩박스 전부 제거
# 3) RandomHelmetPasteWithRotation(label_id=0): 합성된 "head" bbox만 남음
# 4) 배경=1에 대한 bbox는 없음(실제론 그냥 negative 영역이 됨)
##############################################################################