##############################
# 0) custom_imports
##############################
custom_imports = dict(
    imports=['my_custom_pipeline.custom_transforms'],  # 커스텀 변환 모듈 임포트
    allow_failed_imports=False  # 임포트 실패 시 오류 발생
)

import os

# 이미지의 스케일 (가로, 세로)
img_scale = (960, 960)

# 최대 학습 에폭 수
max_epochs = 30

# 모델 체크포인트를 저장할 에폭 간격
save_epoch_intervals = 10

# GPU당 학습 배치 사이즈
train_batch_size_per_gpu = 32

# 학습 데이터 로더의 워커 수
train_num_workers = 0

# 워커의 지속성 여부
persistent_workers = False

# 기본 학습률
base_lr = 0.002

# 배경 이미지 디렉토리 경로
bg_image_dir = '/DATA2/ltb/mmyolo/it1_1761_black_except/images'

# 데이터셋 타입
dataset_type = 'YOLOv5CocoDataset'

# 데이터 루트 디렉토리 경로
data_root = '/DATA2/ltb/mmyolo/robo_head_only_1561_250117/'

# 클래스 이름 및 클래스 수 정의
class_name = ('head', 'background')
num_classes = 2

# 학습 및 검증 어노테이션 파일 경로
train_ann_file = 'train.json'
val_ann_file = 'val.json'

# 메타 정보 (클래스 및 색상 팔레트)
metainfo = dict(
    classes=class_name,
    palette=[(220, 20, 60), (90, 10, 90)]
)

##############################
# Albumentations (pascal_voc + clip=True)
##############################

# 학습 데이터 파이프라인 정의
train_pipeline = [
    dict(type='LoadImageFromFile'),  # 이미지 파일 로드
    dict(type='LoadAnnotations', with_bbox=True),  # 어노테이션(바운딩 박스) 로드

    # 커스텀 변환: 배경에 head 이미지 합성
    dict(
        type='PasteHeadOnBg',
        bg_img_dir=bg_image_dir,  # 배경 이미지 디렉토리
        prob=1.0,  # 합성 확률 (1.0: 항상 합성)
        min_scale=0.5,  # 최소 스케일 비율
        max_scale=1.0   # 최대 스케일 비율
    ),

    dict(type='YOLOv5HSVRandomAug'),  # 색상 관련 랜덤 증강
    dict(type='mmdet.RandomFlip', prob=0.5),  # 랜덤 플립 (확률 0.5)

    dict(type='YOLOv5KeepRatioResize', scale=img_scale),  # 비율 유지하며 리사이즈
    dict(
        type='LetterResize',
        scale=img_scale,  # 리사이즈 후 스케일
        allow_scale_up=False,  # 스케일 업 허용 여부
        pad_val=dict(img=114)  # 패딩 값 (이미지)
    ),

    dict(
        type='mmdet.PackDetInputs',
        meta_keys=('img_id','img_path','ori_shape','img_shape',
                   'flip','flip_direction','scale_factor','pad_param')  # 메타 데이터 키
    )
]

# 테스트 데이터 파이프라인 정의
test_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),  # 어노테이션(바운딩 박스) 로드
    dict(
        type='mmdet.PackDetInputs',
        meta_keys=('img_id','img_path','ori_shape','img_shape',
                   'scale_factor','pad_param')  # 메타 데이터 키
    )
]

# 학습 데이터로더 설정
train_dataloader = dict(
    batch_size=train_batch_size_per_gpu,  # 배치 사이즈
    num_workers=train_num_workers,        # 워커 수
    persistent_workers=persistent_workers,  # 워커 지속성
    sampler=dict(type='DefaultSampler', shuffle=True),  # 샘플러 타입 및 셔플 여부
    dataset=dict(
        type=dataset_type,  # 데이터셋 타입
        data_root=data_root,  # 데이터 루트 경로
        ann_file=train_ann_file,  # 학습 어노테이션 파일
        data_prefix=dict(img='train/images/'),  # 이미지 데이터 경로 접두사
        filter_cfg=dict(filter_empty_gt=True, min_size=32),  # 필터 설정: 빈 GT 필터링 활성화
        pipeline=train_pipeline,  # 학습 파이프라인
        metainfo=metainfo  # 메타 정보
    )
)

# 검증 데이터로더 설정
val_dataloader = dict(
    batch_size=1,  # 배치 사이즈 (검증은 보통 1)
    num_workers=0,  # 워커 수
    persistent_workers=False,  # 워커 지속성
    sampler=dict(type='DefaultSampler', shuffle=False),  # 샘플러 타입 및 셔플 여부
    dataset=dict(
        type=dataset_type,  # 데이터셋 타입
        data_root=data_root,  # 데이터 루트 경로
        ann_file=val_ann_file,  # 검증 어노테이션 파일
        data_prefix=dict(img='val/images/'),  # 이미지 데이터 경로 접두사
        test_mode=True,  # 테스트 모드 활성화
        pipeline=test_pipeline,  # 테스트 파이프라인
        metainfo=metainfo  # 메타 정보
    )
)

# 테스트 데이터로더는 검증 데이터로더와 동일하게 설정
test_dataloader = val_dataloader

# 검증 평가자 설정
val_evaluator = dict(
    type='mmdet.CocoMetric',  # 평가자 타입 (COCO 메트릭)
    ann_file=os.path.join(data_root, val_ann_file),  # 어노테이션 파일 경로
    metric='bbox',  # 평가 메트릭 (바운딩 박스)
    format_only=False  # 포맷만 수행 여부
)
test_evaluator = val_evaluator  # 테스트 평가자도 동일하게 설정

##############################
# 모델 설정 등
##############################
model = dict(
    type='YOLODetector',  # 모델 타입
    data_preprocessor=dict(
        type='YOLOv5DetDataPreprocessor',  # 데이터 전처리기 타입
        bgr_to_rgb=True,  # BGR을 RGB로 변환
        mean=[0.0, 0.0, 0.0],  # 이미지 평균값
        std=[255.0, 255.0, 255.0]  # 이미지 표준편차
    ),
    backbone=dict(
        type='YOLOv5CSPDarknet',  # 백본 타입
        deepen_factor=0.33,  # 깊이 팩터 (예: 0.33)
        widen_factor=0.50    # 너비 팩터 (예: 0.50)
    ),
    neck=dict(
        type='YOLOv5PAFPN',  # 넥 타입
        deepen_factor=0.33,  # 깊이 팩터
        widen_factor=0.50,   # 너비 팩터
        in_channels=[128, 256, 512],  # 입력 채널
        out_channels=[128, 256, 512],  # 출력 채널
        num_csp_blocks=3  # CSP 블록 수
    ),
    bbox_head=dict(
        type='YOLOv5Head',  # 바운딩 박스 헤드 타입
        head_module=dict(
            type='YOLOv5HeadModule',  # 헤드 모듈 타입
            num_classes=num_classes,  # 클래스 수
            in_channels=[128, 256, 512],  # 입력 채널
            widen_factor=0.50,  # 너비 팩터
            featmap_strides=[8, 16, 32],  # 피처맵 스트라이드
            num_base_priors=3  # 기본 앵커 수
        ),
        prior_generator=dict(
            type='mmdet.YOLOAnchorGenerator',  # 앵커 생성기 타입
            base_sizes=[
                [(10,13),(16,30),(33,23)],
                [(30,61),(62,45),(59,119)],
                [(116,90),(156,198),(373,326)]
            ],  # 각 스케일 별 기본 앵커 사이즈
            strides=[8,16,32]  # 앵커 스트라이드
        ),
        loss_cls=dict(
            type='mmdet.CrossEntropyLoss',  # 분류 손실 함수 타입
            use_sigmoid=True,  # 시그모이드 활성화 사용
            reduction='none',  # 손실 축소 방법
            loss_weight=0.5  # 손실 가중치
        ),
        loss_bbox=dict(
            type='IoULoss',  # IoU 기반 손실 함수 타입
            iou_mode='ciou',  # IoU 모드 (Complete IoU)
            bbox_format='xywh',  # 바운딩 박스 포맷
            reduction='sum',  # 손실 축소 방법
            loss_weight=7.5  # 손실 가중치
        )
    ),
    train_cfg=dict(
        assigner=dict(
            type='BatchTaskAlignedAssigner',  # 어사이너 타입
            num_classes=num_classes,  # 클래스 수
            alpha=0.5,  # 알파 값
            beta=6.0,  # 베타 값
            use_ciou=True,  # Complete IoU 사용 여부
            topk=10,  # Top-K 설정
            eps=1e-9  # 작은 값
        )
    ),
    test_cfg=dict(
        multi_label=True,  # 멀티 라벨 사용 여부
        nms_pre=30000,  # NMS 사전 처리 단계
        score_thr=0.01,  # 점수 임계값
        nms=dict(type='nms', iou_threshold=0.65),  # NMS 설정 (IoU 임계값 0.65)
        max_per_img=300  # 이미지당 최대 검출 수
    )
)

# 최적화 래퍼 설정 (AMP 사용)
optim_wrapper = dict(
    type='AmpOptimWrapper',  # AMP 최적화 래퍼 타입
    optimizer=dict(
        type='SGD',  # 옵티마이저 타입 (SGD)
        lr=base_lr,  # 학습률
        momentum=0.937,  # 모멘텀
        nesterov=True,  # 네스테로프 모멘텀 사용 여부
        weight_decay=0.0005  # 가중치 감쇠
    ),
    clip_grad=dict(max_norm=5.0),  # 그래디언트 클리핑 설정 (최대 노름 5.0)
    constructor='YOLOv5OptimizerConstructor'  # 옵티마이저 생성자
)

# 파라미터 스케줄러 설정 (여기서는 사용하지 않음)
param_scheduler = None

# 학습 설정
train_cfg = dict(
    type='EpochBasedTrainLoop',  # 학습 루프 타입 (에폭 기반)
    max_epochs=max_epochs,  # 최대 에폭 수
    val_interval=save_epoch_intervals  # 검증 간격 (에폭 단위)
)
val_cfg = dict(type='ValLoop')  # 검증 루프 타입
test_cfg = dict(type='TestLoop')  # 테스트 루프 타입

# 기본 훅 설정
default_hooks = dict(
    checkpoint=dict(
        type='CheckpointHook',  # 체크포인트 훅 타입
        interval=save_epoch_intervals,  # 체크포인트 저장 간격 (에폭 단위)
        max_keep_ckpts=3,  # 최대 저장할 체크포인트 수
        save_best='auto'  # 최적의 체크포인트 자동 저장
    ),
    logger=dict(type='LoggerHook', interval=10),  # 로깅 훅 설정 (10 반복마다 로그 출력)
    param_scheduler=dict(
        type='YOLOv5ParamSchedulerHook',  # 파라미터 스케줄러 훅 타입
        scheduler_type='linear',  # 스케줄러 타입 (선형)
        lr_factor=0.01,  # 학습률 팩터
        max_epochs=max_epochs  # 최대 에폭 수
    ),
    sampler_seed=dict(type='DistSamplerSeedHook'),  # 샘플러 시드 훅
    timer=dict(type='IterTimerHook'),  # 타이머 훅
    visualization=dict(type='mmdet.DetVisualizationHook')  # 시각화 훅
)

# 커스텀 훅 설정
custom_hooks = [
    dict(
        type='EMAHook',  # EMA 훅 타입
        ema_type='ExpMomentumEMA',  # EMA 타입 (지수 모멘텀 EMA)
        momentum=0.0001,  # EMA 모멘텀 값
        update_buffers=True,  # 버퍼 업데이트 여부
        priority=49  # 우선순위
    )
]

# 기본 스코프 설정
default_scope = 'mmyolo'

# 환경 설정
env_cfg = dict(
    cudnn_benchmark=True,  # CuDNN 벤치마크 사용 여부
    mp_cfg=dict(mp_start_method='fork', opencv_num_threads=0),  # 멀티프로세싱 설정
    dist_cfg=dict(backend='nccl')  # 분산 학습 백엔드 설정
)

# 시각화 백엔드 설정
vis_backends = [dict(type='LocalVisBackend')]  # 로컬 시각화 백엔드 사용
visualizer = dict(
    type='mmdet.DetLocalVisualizer',  # 시각화기 타입
    vis_backends=vis_backends,  # 시각화 백엔드 설정
    name='visualizer'  # 시각화기 이름
)

# 로깅 레벨 설정
log_level = 'INFO'

# 사전 학습된 모델 로드 설정 (없음)
load_from = None

# 학습 재개 설정 (False: 새로 시작)
resume = False

# 작업 디렉토리 설정 (모델 체크포인트 및 로그 저장 위치)
work_dir = '/DATA2/ltb/mmyolo/work_dirs/head_only_with_bg'

##############################
# 실행 후 출력 정보 설정
##############################
# 실행 시간 및 경로 등의 정보는 기본적으로 mmengine이 처리하므로,
# 추가적인 설정 없이도 로그에 간결하게 출력됩니다.
# 필요한 경우 LoggerHook의 설정을 조정하여 로그 출력을 세부 조정할 수 있습니다.

# 예를 들어, logger의 interval을 변경하여 출력 빈도를 조절할 수 있습니다.
# 현재는 interval=10으로 설정되어 있어, 10번의 반복(iteration)마다 로그가 출력됩니다.

# 추가적으로, 작업 디렉토리(work_dir)와 관련된 정보는 CheckpointHook에서 자동으로 관리됩니다.
# 모델의 체크포인트는 '/DATA2/ltb/mmyolo/work_dirs/head_only_with_bg' 디렉토리에 저장됩니다.

# 실행 시간을 확인하고 싶다면, 학습 시작 시와 종료 시의 시간을 기록하는 커스텀 훅을 추가할 수 있습니다.
# 하지만, 기본 타이머 훅(IterTimerHook)이 이미 각 반복(iteration)의 소요 시간을 측정하여 로그로 출력합니다.

##############################
# 기타 설정
##############################
# 필요에 따라 추가적인 설정을 여기에 포함할 수 있습니다.
# 예: GPU 설정, 랜덤 시드 고정 등