_base_ = '../yolov5/yolov5_s-v61_syncbn_fast_1xb4-300e_balloon.py'

# ========================modified parameters======================
data_root = '/data/si/darwin/si/construction_signalman/'  # Root path of data
# Path of train annotation file
train_ann_file = 'train_balance.json'
train_data_prefix = 'images/'  # Prefix of train image path
# Path of val annotation file
val_ann_file = 'test_balance.json'
val_data_prefix = 'images/'  # Prefix of val image path

metainfo = {
    'classes': ('worker', 'signalman'),
    'palette': [
        (220, 20, 60), (119, 11, 32)
    ]
}

num_classes = 2

train_batch_size_per_gpu = 4
train_num_workers = 2
log_interval = 1

# =======================Unmodified in most cases==================
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_data_prefix),
        ann_file=train_ann_file))
val_dataloader = dict(
    dataset=dict(
        data_root=data_root,
        metainfo=metainfo,
        data_prefix=dict(img=val_data_prefix),
        ann_file=val_ann_file))
test_dataloader = val_dataloader
val_evaluator = dict(ann_file=data_root + val_ann_file)
test_evaluator = val_evaluator
model = dict(bbox_head=dict(head_module=dict(num_classes=num_classes)))
default_hooks = dict(logger=dict(interval=log_interval))