# Inherit and overwrite part of the config based on this config
_base_ = './faster-rcnn_r50_fpn_1x_coco.py'

data_root = 'data/' # dataset root

train_batch_size_per_gpu = 16
train_num_workers = 1

max_epochs = 105
base_lr = 0.00001


metainfo = {
    'classes': ('orgaquant', ),
    'palette': [
        (220, 20, 60),
    ]
}

train_dataloader = dict(
    batch_size=train_batch_size_per_gpu,
    num_workers=train_num_workers,
    dataset=dict(
        data_root=data_root,
        metainfo=metainfo,
        data_prefix=dict(img='train/'),
        ann_file='train.json'))

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

test_dataloader = val_dataloader

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

test_evaluator = val_evaluator

model = dict(
    roi_head=dict(
        bbox_head=dict(num_classes=1)))



train_pipeline = [
dict(type='LoadImageFromFile', backend_args=None),
    dict(type='LoadAnnotations', with_bbox=True),
    dict(type='RandomFlip', prob=0.5),
    dict(type = 'RandomShift', prob = 0.5),
    dict(type = 'RandomAffine'),
    dict(type='PhotoMetricDistortion'),
    dict(type='PackDetInputs')
    ]
    
    
# optimizer
optim_wrapper = dict(
    _delete_=True,
    type='OptimWrapper',
    optimizer=dict(type='AdamW', lr=base_lr, weight_decay=0.05),
    paramwise_cfg=dict(
        norm_decay_mult=0, bias_decay_mult=0, bypass_duplicate=True))

default_hooks = dict(
    checkpoint=dict(
        interval=5,
        max_keep_ckpts=2,  # only keep latest 2 checkpoints
        save_best='auto'
    ),
    logger=dict(type='LoggerHook', interval=5))


# load COCO pre-trained weight

# load_from = './work_dirs/faster-rcnn_r50_fpn_organoid/best_coco_bbox_mAP_epoch_12.pth'


train_cfg = dict(type='EpochBasedTrainLoop', max_epochs=max_epochs, val_interval=1)
visualizer = dict(vis_backends=[dict(type='LocalVisBackend'),dict(type='TensorboardVisBackend')])