import torch


def _convert_boxes_to_roi_format(boxes):
    """
    Convert rois into the torchvision format.

    :param boxes: The roi boxes as a native tensor[B, K, 4].
    :return: The roi boxes in the format that roi pooling and roi align in torchvision require. Native tensor[B*K, 5].
    """
    concat_boxes = boxes.view((-1, 4))
    ids = torch.full_like(boxes[:, :, :1], 0)
    for i in range(boxes.shape[0]):
        ids[i, :, :] = i
    ids = ids.view((-1, 1))
    rois = torch.cat([ids, concat_boxes], dim=1)
    return rois