_base_ = '../yolov8/yolov8_l_syncbn_fast_8xb16-500e_coco.py'

max_epochs = 100
data_root = '/DATA2/ltb/mmyolo/roboflow_human_head/head_dataset/'
work_dir = './work_dirs/it1_worker_helmet_2'

load_from = 'yolov8_l_syncbn_fast_8xb16-500e_coco_20230217_182526-189611b6.pth'
train_batch_size_per_gpu = 8
train_num_workers = 0
save_epoch_intervals = 1

base_lr = _base_.base_lr / 4

class_name = ('head', 'background')
num_classes = len(class_name)
metainfo = dict(
    classes=class_name,
    palette=[(220, 20, 60), (90, 10, 90)]
)

train_cfg = dict(
    max_epochs=max_epochs,
    val_begin=1,
    val_interval=save_epoch_intervals
)

model = dict(
    bbox_head=dict(
        head_module=dict(num_classes=num_classes),
        # 아래는 예시로 bbox 로스 가중치를 높인 코드
        loss_cls=dict(loss_weight=0.5 * (num_classes / 2 * 3 / _base_.num_det_layers)),
        loss_bbox=dict(
            type='IoULoss',
            iou_mode='ciou',
            bbox_format='xywh',
            reduction='sum',
            loss_weight=10.0,  # 소형 객체 강조
            return_iou=False
        ),
    ),
    test_cfg=dict(
        # nms iou_threshold를 약간 낮춰 소형 객체 누락 감소
        nms=dict(type='nms', iou_threshold=0.4)
    )
)

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,
            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=[
                dict(type='LoadImageFromFile', backend_args=_base_.backend_args),
                dict(type='LoadAnnotations', with_bbox=True),

                # RandomCrop -> mmdet.RandomCrop
                dict(
                    type='mmdet.RandomCrop',
                    crop_size=(1024, 1024),
                    allow_negative_crop=True
                ),
                # RandomScale -> mmdet.RandomResize (ratio_range)
                dict(
                    type='mmdet.RandomResize',
                    # 원래 scale=None -> (640, 640) 등의 기준값 명시
                    scale=(640, 640),            
                    ratio_range=(0.8, 1.2),
                    keep_ratio=True
                ),
                dict(
                    type='Mosaic',
                    img_scale=(1280, 1280),
                    pad_val=114.0,
                    pre_transform=[
                        dict(type='LoadImageFromFile', backend_args=_base_.backend_args),
                        dict(type='LoadAnnotations', with_bbox=True)
                    ]
                ),
                dict(
                    type='YOLOv5RandomAffine',
                    max_rotate_degree=0.0,
                    max_shear_degree=0.0,
                    scaling_ratio_range=(0.8, 1.2),
                    max_aspect_ratio=100,
                    border=(-640, -640),
                    border_val=(114, 114, 114)
                ),
                dict(
                    type='mmdet.Albu',
                    transforms=[
                        dict(type='Blur', p=0.01),
                        dict(type='MedianBlur', p=0.01),
                        dict(type='ToGray', p=0.01),
                        dict(type='CLAHE', p=0.01)
                    ],
                    bbox_params=dict(
                        type='BboxParams',
                        format='pascal_voc',
                        label_fields=['gt_bboxes_labels', 'gt_ignore_flags']
                    ),
                    keymap={'img': 'image', 'gt_bboxes': 'bboxes'}
                ),
                dict(type='YOLOv5HSVRandomAug'),
                dict(type='mmdet.RandomFlip', prob=0.5),
                dict(
                    type='mmdet.PackDetInputs',
                    meta_keys=('img_id', 'img_path', 'ori_shape',
                               'img_shape', 'flip', 'flip_direction')
                )
            ]
        )
    )
)

val_dataloader = dict(
    dataset=dict(
        metainfo=metainfo,
        data_root=data_root,
        ann_file='val.json',
        data_prefix=dict(img='val/images/')
    )
)
test_dataloader = val_dataloader

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

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