# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
from dynamic_rcnn.datasets.structures.image_list import to_image_list


class BatchCollator(object):
    """
    From a list of samples from the dataset,
    returns the batched images and targets.
    This should be passed to the DataLoader
    """

    def __init__(self, size_divisible=0, return_raw=False):
        self.size_divisible = size_divisible
        self.return_raw = return_raw

    def __call__(self, batch):
        transposed_batch = list(zip(*batch))
        if self.return_raw:
            ori_images = transposed_batch[0]
            ori_targets = transposed_batch[1]
            images = to_image_list(
                transposed_batch[2], self.size_divisible)
            targets = transposed_batch[3]
            img_ids = transposed_batch[4]
            return ori_images, ori_targets, images, targets, img_ids
        else:
            images = to_image_list(transposed_batch[0], self.size_divisible)
            targets = transposed_batch[1]
            img_ids = transposed_batch[2]
            return images, targets, img_ids


class BBoxAugCollator(object):
    """
    From a list of samples from the dataset,
    returns the images and targets.
    Images should be converted to batched images in `im_detect_bbox_aug`
    """

    def __call__(self, batch):
        return list(zip(*batch))

