import torch


def collate_fn(items):
        """
        Creates mini-batch tensors from the list of tuples (image, caption).
        Args:
            data: list of tuple (image, caption).
                - image: torch tensor of shape
                - caption: torch tensor of shape (?); variable length.

        Returns:
            images: torch tensor of images.
            targets: torch tensor of shape (batch_size, padded_length).
            lengths: list; valid length for each padded caption.
            image_ids: List; batch中每个image唯一的id
        """
        image_batch, label_batch, gts_batch, info_batch = zip(*items)
        # image_batch, caption_batch, imageid_batch = zip(*items)
        # Merge images (from tuple of 3D tensor to 4D tensor).
        image_batch = torch.stack(image_batch, 0)
        info_batch = list(info_batch)

        lengths = [label.size()[0] for label in label_batch]

        label_batch = torch.nn.utils.rnn.pad_sequence(label_batch, batch_first=True, padding_value=0)
        gts_batch = torch.nn.utils.rnn.pad_sequence(gts_batch, batch_first=True, padding_value=0)


        mask_batch = torch.zeros_like(label_batch)
        for i, len in enumerate(lengths):
            for j in range(len):
                mask_batch[i, j] = 1
        return image_batch, label_batch, mask_batch, gts_batch, info_batch