# -*- coding: utf-8 -*-
_base_ = '../yolov8/yolov8_x_mask-refine_syncbn_fast_8xb16-500e_coco.py'

# -----data related-----
data_root = '/'  # Root path of data

class_list = ('ConcreteCrack', 'Efflorescene', 'Exposure', 'Spalling')
persistent_workers = True

BATCH_SIZE = 5
NUM_WORKER = 32
NUM_CLASS  = 4

train_dataloader = dict(
    batch_size=BATCH_SIZE,
    num_workers=NUM_WORKER,
    dataset=dict(
        data_root=data_root,

        ann_file="/data/si/TTA_도공/server_aihub_train_5k.json",
        data_prefix=dict(img='/'),
        metainfo=dict(classes=class_list)))
val_dataloader = dict(
    batch_size=BATCH_SIZE,
    num_workers=NUM_WORKER,
    dataset=dict(
        data_root=data_root,
        ann_file="/data/si/TTA_도공/server_aihub_train_5k.json",
        data_prefix=dict(img='/'),
        metainfo=dict(classes=class_list)))
test_dataloader = val_dataloader
model = dict(
    bbox_head=dict(
        head_module=dict(num_classes=NUM_CLASS)),
    train_cfg=dict(
            assigner=dict(
                num_classes=NUM_CLASS)))
val_evaluator = dict(
    ann_file="/data/si/TTA_도공/server_aihub_train_5k.json")
test_evaluator = dict(
    ann_file="/data/si/TTA_도공/server_aihub_train_5k.json")
find_unused_parameters=True
# Customize train process
default_hooks = dict(
    timer=dict(type='IterTimerHook'),
    logger=dict(type='LoggerHook', interval=10),
    param_scheduler=dict(
        type='YOLOv5ParamSchedulerHook',
        scheduler_type='linear',
        lr_factor=0.01,
        max_epochs=100),
    checkpoint=dict(
        type='CheckpointHook', interval=5, save_best='auto',
        max_keep_ckpts=2),
    sampler_seed=dict(type='DistSamplerSeedHook'),
    visualization=dict(type='mmdet.DetVisualizationHook'))

train_cfg = dict(
    type='EpochBasedTrainLoop',
    max_epochs=100,
    val_interval=1,
    dynamic_intervals=[(40, 1)])

# Load pretrained model
# load_from = '/home/daitranskku/code/SmartInsideAI/projects/CRACKS_DETECTIONS/submission/best_coco_bbox_mAP_epoch_422.pth'