import torch


def to_bbox_in_original_image(bboxes, original_image_size, model_input_size):
    """Convert bounding box to the one in the original image

    This function uses an original image size and a model input size to
    transform bounding box coordinates from the ones in the preprocessed
    image to the ones in the original image, because YOLO takes padded
    and resized image as input.

    Args:
        bboxes (torch.tensor): bounding boxes YOLO outputs (shape: Nx4),
            each row is (x0, y0, w, h) where
                (x0, y0): coordinates of top-left corner of bounding box
                (w, h): size of bounding box
        original_image_size (torch.tensor): original image size (shape: 2)
        model_input_size (torch.tensor): model input size (shape: 1)

    Returns:
        (torch.tensor): COCO-style bounding boxes in the original image
            (shape: Nx4)

    """
    w, h = original_image_size
    embedded_image_size = (original_image_size * model_input_size /
                           original_image_size.max()).long()
    offset = (max(embedded_image_size) - min(embedded_image_size)) // 2
    if w < h:
        offset = torch.tensor([offset, 0],
                              dtype=torch.float32,
                              device=bboxes.device)
    else:
        offset = torch.tensor([0, offset],
                              dtype=torch.float32,
                              device=bboxes.device)
    bboxes[:, :2] -= offset
    bboxes *= original_image_size.max() / model_input_size
    return bboxes