# custom_imports = dict(imports=['dataset.missing_person'], allow_failed_imports=False)

dataset_type = 'CrowdhumanDataset'  # Dataset type, this will be used to define the dataset
# dataset_type = 'CocoDataset'  # Dataset type, this will be used to define the dataset
# dataset_type = 'CocoDataset'  # Dataset type, this will be used to define the dataset
# data_root = '/mnt/home2/annonymous/CrowdHuman'  # Root path of data
data_root = '/mnt/home/annonymous/iclr2026/CrowdHuman'  # Root path of data
backend_args = None # Arguments to instantiate the corresponding file backend

DEFAULT_TRAIN_DATASET_PIPELINE = [
    dict(type='LoadImageFromFile', backend_args=backend_args),
    dict(type='LoadAnnotations', with_bbox=True),
    dict(type='Resize', scale=(1333, 800), keep_ratio=True),
    dict(type='RandomFlip', prob=0.5),
    dict(type='PackDetInputs')
]

DEFAULT_VAL_DATASET_PIPELINE = [
    dict(type='LoadImageFromFile', backend_args=backend_args),
    dict(type='Resize', scale=(1333, 800), keep_ratio=True),
    # If you don't have a gt annotation, delete the pipeline
    dict(type='LoadAnnotations', with_bbox=True),
    dict(
        type='PackDetInputs',
        meta_keys=('img_id', 'img_path', 'ori_shape', 'img_shape',
                   'scale_factor'))
]

DEFAULT_TRAIN_DATALOADER_DICT = dict(   # Train dataloader config
    batch_size=8,  # Batch size of a single GPU => total batch size should be 16
    num_workers=2,  # Worker to pre-fetch data for each single GPU
    persistent_workers=True,  # If ``True``, the dataloader will not shut down the worker processes after an epoch end, which can accelerate training speed.
    sampler=dict(  # training data sampler
        type='DefaultSampler',  # DefaultSampler which supports both distributed and non-distributed training. Refer to https://mmengine.readthedocs.io/en/latest/api/generated/mmengine.dataset.DefaultSampler.html#mmengine.dataset.DefaultSampler
        shuffle=True),  # randomly shuffle the training data in each epoch
    batch_sampler=dict(type='AspectRatioBatchSampler'),  # Batch sampler for grouping images with similar aspect ratio into a same batch. It can reduce GPU memory cost.
    dataset=dict(  # Train dataset config
        type=dataset_type,
        data_root=data_root,
        ann_file='crowdhuman_train.json',  # Path of annotation file
        data_prefix=dict(img='train/'),  # Prefix of image path
        # filter_cfg=dict(filter_empty_gt=True, min_size=32),  # Config of filtering images and annotations
        filter_cfg=dict(min_size=32),
        pipeline=DEFAULT_TRAIN_DATASET_PIPELINE,
        backend_args=backend_args))

DEFAULT_VAL_DATALOADER_DICT = dict(  # Validation dataloader config
    batch_size=8,  # Batch size of a single GPU. If batch-size > 1, the extra padding area may influence the performance.
    num_workers=2,  # Worker to pre-fetch data for each single GPU
    persistent_workers=True,  # If ``True``, the dataloader will not shut down the worker processes after an epoch end, which can accelerate training speed.
    drop_last=False,  # Whether to drop the last incomplete batch, if the dataset size is not divisible by the batch size
    sampler=dict(
        type='DefaultSampler',
        shuffle=False),  # not shuffle during validation and testing
    dataset=dict(
        type=dataset_type,
        data_root=data_root,
        ann_file='crowdhuman_val.json',
        data_prefix=dict(img='val/'),
        test_mode=True,  # Turn on the test mode of the dataset to avoid filtering annotations or images
        pipeline=DEFAULT_VAL_DATASET_PIPELINE,
        backend_args=backend_args))

# DEFAULT_TEST_DATALOADER_DICT = dict(  # Validation dataloader config
#     batch_size=2,  # Batch size of a single GPU. If batch-size > 1, the extra padding area may influence the performance.
#     num_workers=2,  # Worker to pre-fetch data for each single GPU
#     persistent_workers=True,  # If ``True``, the dataloader will not shut down the worker processes after an epoch end, which can accelerate training speed.
#     drop_last=False,  # Whether to drop the last incomplete batch, if the dataset size is not divisible by the batch size
#     sampler=dict(
#         type='DefaultSampler',
#         shuffle=False),  # not shuffle during validation and testing
#     dataset=dict(
#         type=dataset_type,
#         data_root=data_root,
#         ann_file='test/_annotations.json',
#         data_prefix=dict(img='test/'),
#         test_mode=True,  # Turn on the test mode of the dataset to avoid filtering annotations or images
#         pipeline=DEFAULT_VAL_DATASET_PIPELINE,
#         backend_args=backend_args))


train_dataloader = DEFAULT_TRAIN_DATALOADER_DICT
val_dataloader = DEFAULT_VAL_DATALOADER_DICT # TMP
# test_dataloader = DEFAULT_TEST_DATALOADER_DICT # TMP
test_dataloader = DEFAULT_VAL_DATALOADER_DICT # TMP