############################
# final_2304_config.py
#   - partial load backbone with 2304 last channel + 2 classes
############################

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

max_epochs = 100
data_root = '/DATA2/ltb/mmyolo/it1-worker-helmet-1721-detail/'
work_dir = './work_dirs/it1_worker_helmet_2'

# Using this COCO-pretrained checkpoint
load_from = None  # We'll partial load using init_cfg below
resume = False

# training batch etc.
train_batch_size_per_gpu = 64
train_num_workers = 0
save_epoch_intervals = 1

# base_lr => the default in _base_ is ~0.01 for 8xb16=128
base_lr = _base_.base_lr / 4   # or you can set e.g. 1e-3

class_name = ('Worker', 'helmet',)
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(
    type='YOLODetector',
    # partial load backbone only
    init_cfg=dict(
        type='Pretrained',
        checkpoint='yolov8_l_syncbn_fast_8xb16-500e_coco_20230217_182526-189611b6.pth',
        prefix='backbone.'
    ),
    # -----Backbone-----
    backbone=dict(
        type='YOLOv8CSPDarknet',
        arch='P5',
        # Possibly the widen_factor, deepen_factor, or last_stage_out_channels
        # is set so that final out has 2304 channels
        widen_factor=1.0,
        deepen_factor=1.0,
        norm_cfg=dict(type='BN', momentum=0.03, eps=0.001),
        act_cfg=dict(type='SiLU', inplace=True),
    ),
    # -----Neck-----
    neck=dict(
        type='YOLOv8PAFPN',
        # We guess the three scale channels might be [768, 1152, 2304], or
        # [512, 1024, 2304]. We must match the actual structure used by the checkpoint
        # Suppose it's [768, 1152, 2304]
        in_channels=[768, 1152, 2304],
        out_channels=[768, 1152, 2304],
        deepen_factor=1.0,
        widen_factor=1.0,
        num_csp_blocks=3,
        norm_cfg=dict(type='BN', momentum=0.03, eps=0.001),
        act_cfg=dict(type='SiLU', inplace=True),
    ),
    # -----Head-----
    bbox_head=dict(
        type='YOLOv8Head',
        head_module=dict(
            num_classes=num_classes,  # 2
            # must match neck out_channels
            in_channels=[768, 1152, 2304],
            widen_factor=1.0,
            reg_max=16,
            norm_cfg=dict(type='BN', momentum=0.03, eps=0.001),
            act_cfg=dict(type='SiLU', inplace=True),
            featmap_strides=[8, 16, 32]
        ),
        prior_generator=dict(
            type='mmdet.MlvlPointGenerator',
            offset=0.5,
            strides=[8,16,32]
        ),
        bbox_coder=dict(type='DistancePointBBoxCoder'),
        loss_cls=dict(
            type='mmdet.CrossEntropyLoss',
            use_sigmoid=True,
            reduction='none',
            loss_weight=1.0),
        loss_bbox=dict(
            type='IoULoss',
            iou_mode='ciou',
            bbox_format='xywh',
            reduction='sum',
            loss_weight=7.5,
            return_iou=False),
        loss_dfl=dict(
            type='mmdet.DistributionFocalLoss',
            reduction='mean',
            loss_weight=1.5/4)
    ),
    test_cfg=dict(
        multi_label=True,
        score_thr=0.01,
        nms=dict(type='nms', iou_threshold=0.5),
        max_per_img=300
    )
)

# ---------------- Pipeline from your base config ----------------
train_pipeline = _base_.train_pipeline
# or you can copy/paste from your base config, then add Albumentations Rotate
# If you want to override it, do:
# train_pipeline = [
#   ...
# ]

# Because you have an existing pipeline in _base_.train_pipeline,
# you can modify or extend it if needed. For example:
# 
# train_pipeline = _base_.train_pipeline.copy()
# # Insert your Albumentations rotate transform
# idx = [i for i, t in enumerate(train_pipeline) if t['type']=='Mosaic'][0] + 1
# train_pipeline.insert(idx, dict(...)) # etc.

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=_base_.train_pipeline  # or your custom pipeline
        )
    )
)

val_dataloader = dict(
    dataset=dict(
        type=_base_.dataset_type,
        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(
        type='SGD',
        lr=base_lr,
        momentum=0.937,
        weight_decay=0.0005
    )
)

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