import torch


def base_detection_collate_fn(batch):
    """
    Batch Processing helper function for detection training/testing.
    stacks the lists of images and targets into tensors and adds the image index to each target (so the targets could
    later be associated to the correct images)
         :param batch:   Input batch from the Dataset __get_item__ method
         :return:        batch with the transformed values
     """

    images_batch, labels_batch = list(zip(*batch))
    for i, labels in enumerate(labels_batch):
        # ADD TARGET IMAGE INDEX
        labels[:, 0] = i

    return torch.stack(images_batch, 0), torch.cat(labels_batch, 0)