<gh_stars>1-10
import torch
from torch.utils.data import DataLoader

from ssd.config import paths_catalog
from ssd.utils.comm import get_world_size
from . import datasets as D
from . import samplers
from .collate_batch import BatchCollator
from .transforms import build_target_transforms
from .transforms import build_transforms


def build_dataset(dataset_list, transforms, target_transforms, dataset_catalog, is_train=True):
    if not isinstance(dataset_list, (tuple, list)):
        raise RuntimeError(
            f'dataset_list should be a list of strings, got {dataset_list}'
        )
    datasets = []
    for dataset_name in dataset_list:
        data = dataset_catalog.get(dataset_name)
        factory = getattr(D, data['factory'])
        args = data['args']

        if data['factory'] == 'COCODataset':
            args['remove_images_without_annotations'] = is_train
        if data['factory'] == 'PascalVOCDataset':
            args['use_difficult'] = not is_train
        args['transforms'] = transforms
        args['target_transforms'] = target_transforms
        dataset = factory(**args)
        datasets.append(dataset)

    if not is_train:
        return datasets
    dataset = datasets[0]
    if len(datasets) > 1:
        dataset = D.ConcatDataset(datasets)
    return [dataset]


def make_data_sampler(dataset, shuffle, distributed):
    if distributed:
        return samplers.DistributedSampler(dataset, shuffle=shuffle)
    if shuffle:
        sampler = torch.utils.data.sampler.RandomSampler(dataset)
    else:
        sampler = torch.utils.data.sampler.SequentialSampler(dataset)
    return sampler


def make_batch_sampler(dataset, sampler, images_per_batch, num_iters=None, start_iter=0):
    batch_sampler = torch.utils.data.BatchSampler(
        sampler,
        images_per_batch,
        drop_last=False
    )
    if num_iters is not None:
        batch_sampler = samplers.IterationBasedBatchSampler(
            batch_sampler,
            num_iters,
            start_iter
        )
    return batch_sampler


def make_data_loader(cfg, is_train=True, is_distributed=False, start_iter=0):
    num_gpus = get_world_size()
    if is_train:
        images_per_batch = cfg.SOLVER.IMS_PER_BATCH
        if images_per_batch % num_gpus != 0:
            raise ValueError(
                f'SOLVER.IMS_PER_BATCH ({images_per_batch}) must be divisible by the number of GPUs ({num_gpus}) used.'
            )
        images_per_gpu = images_per_batch // num_gpus
        shuffle = True
        num_iters = cfg.SOLVER.MAX_ITER
    else:
        images_per_batch = cfg.TEST.IMS_PER_BATCH
        if images_per_batch % num_gpus != 0:
            raise ValueError(
                f'TEST.IMS_PER_BATCH ({images_per_batch}) must be divisible by the number of GPUs ({num_gpus}) used.'
            )
        images_per_gpu = images_per_batch // num_gpus
        shuffle = False if not is_distributed else True
        num_iters = None
        start_iter = 0

    DatasetCatalog = paths_catalog.DatasetCatalog
    dataset_list = cfg.DATASETS.TRAIN if is_train else cfg.DATASETS.TEST

    transforms = build_transforms(cfg, is_train)
    target_transforms = build_target_transforms(cfg)

    datasets = build_dataset(dataset_list, transforms, target_transforms, DatasetCatalog, is_train)
    data_loaders = []
    for dataset in datasets:
        sampler = make_data_sampler(dataset, shuffle, is_distributed)
        batch_sampler = make_batch_sampler(
            dataset,
            sampler,
            images_per_gpu,
            num_iters,
            start_iter
        )
        collator = BatchCollator()
        num_workers = 4
        data_loader = DataLoader(
            dataset,
            num_workers=num_workers,
            batch_sampler=batch_sampler,
            collate_fn=collator,
            pin_memory=True if is_train else False
        )
        data_loaders.append(data_loader)
    if is_train:
        assert len(data_loaders) == 1
        return data_loaders[0]
    return data_loaders

import torch
from torch.utils.data import DataLoader

from ssd.config import paths_catalog
from ssd.utils.comm import get_world_size
from . import datasets as D
from . import samplers
from .collate_batch import BatchCollator
from .transforms import build_target_transforms
from .transforms import build_transforms


def build_dataset(dataset_list, transforms, target_transforms, dataset_catalog, is_train=True):
    if not isinstance(dataset_list, (tuple, list)):
        raise RuntimeError(
            f'dataset_list should be a list of strings, got {dataset_list}'
        )
    datasets = []
    for dataset_name in dataset_list:
        data = dataset_catalog.get(dataset_name)
        factory = getattr(D, data['factory'])
        args = data['args']

        if data['factory'] == 'COCODataset':
            args['remove_images_without_annotations'] = is_train
        if data['factory'] == 'PascalVOCDataset':
            args['use_difficult'] = not is_train
        args['transforms'] = transforms
        args['target_transforms'] = target_transforms
        dataset = factory(**args)
        datasets.append(dataset)

    if not is_train:
        return datasets
    dataset = datasets[0]
    if len(datasets) > 1:
        dataset = D.ConcatDataset(datasets)
    return [dataset]


def make_data_sampler(dataset, shuffle, distributed):
    if distributed:
        return samplers.DistributedSampler(dataset, shuffle=shuffle)
    if shuffle:
        sampler = torch.utils.data.sampler.RandomSampler(dataset)
    else:
        sampler = torch.utils.data.sampler.SequentialSampler(dataset)
    return sampler


def make_batch_sampler(dataset, sampler, images_per_batch, num_iters=None, start_iter=0):
    batch_sampler = torch.utils.data.BatchSampler(
        sampler,
        images_per_batch,
        drop_last=False
    )
    if num_iters is not None:
        batch_sampler = samplers.IterationBasedBatchSampler(
            batch_sampler,
            num_iters,
            start_iter
        )
    return batch_sampler


def make_data_loader(cfg, is_train=True, is_distributed=False, start_iter=0):
    num_gpus = get_world_size()
    if is_train:
        images_per_batch = cfg.SOLVER.IMS_PER_BATCH
        if images_per_batch % num_gpus != 0:
            raise ValueError(
                f'SOLVER.IMS_PER_BATCH ({images_per_batch}) must be divisible by the number of GPUs ({num_gpus}) used.'
            )
        images_per_gpu = images_per_batch // num_gpus
        shuffle = True
        num_iters = cfg.SOLVER.MAX_ITER
    else:
        images_per_batch = cfg.TEST.IMS_PER_BATCH
        if images_per_batch % num_gpus != 0:
            raise ValueError(
                f'TEST.IMS_PER_BATCH ({images_per_batch}) must be divisible by the number of GPUs ({num_gpus}) used.'
            )
        images_per_gpu = images_per_batch // num_gpus
        shuffle = False if not is_distributed else True
        num_iters = None
        start_iter = 0

    DatasetCatalog = paths_catalog.DatasetCatalog
    dataset_list = cfg.DATASETS.TRAIN if is_train else cfg.DATASETS.TEST

    transforms = build_transforms(cfg, is_train)
    target_transforms = build_target_transforms(cfg)

    datasets = build_dataset(dataset_list, transforms, target_transforms, DatasetCatalog, is_train)
    data_loaders = []
    for dataset in datasets:
        sampler = make_data_sampler(dataset, shuffle, is_distributed)
        batch_sampler = make_batch_sampler(
            dataset,
            sampler,
            images_per_gpu,
            num_iters,
            start_iter
        )
        collator = BatchCollator()
        num_workers = 4
        data_loader = DataLoader(
            dataset,
            num_workers=num_workers,
            batch_sampler=batch_sampler,
            collate_fn=collator,
            pin_memory=True if is_train else False
        )
        data_loaders.append(data_loader)
    if is_train:
        assert len(data_loaders) == 1
        return data_loaders[0]
    return data_loaders