import numpy as np

import torch
import torchvision
from torchvision.prototype import features
from torchvision.prototype.transforms.functional import perspective_bounding_box, perspective_image_tensor, perspective_segmentation_mask

size = (64, 76)
# xyxy format
in_boxes = [
    [10, 15, 25, 35],
    [50, 5, 70, 22],
    [45, 46, 56, 62],
]
labels = [1, 2, 3]

im1 = 255 * np.ones(size + (3, ), dtype=np.uint8)
mask = np.zeros(size, dtype=np.int64)
for in_box, label in zip(in_boxes, labels):
    im1[in_box[1]:in_box[3], in_box[0]:in_box[2], :] = (127, 127, 127)
    mask[in_box[1]:in_box[3], in_box[0]:in_box[2]] = label
    
in_im = torch.tensor(im1).permute(2, 0, 1).view(1, 3, *size)

in_boxes = features.BoundingBox(
    in_boxes, format=features.BoundingBoxFormat.XYXY, image_size=size
)
in_mask = features.SegmentationMask(torch.tensor(mask)).view(1, *size)
    

coeffs = [3.7801878452301025,
 2.09562087059021,
 -59.95088195800781,
 0.9144350290298462,
 2.9016025066375732,
 -37.96672821044922,
 0.0310051292181015,
 0.01576310582458973]

out_boxes = perspective_bounding_box(
    in_boxes, 
    in_boxes.format,
    coeffs
)

out_mask = perspective_segmentation_mask(
    in_mask, 
    coeffs
)

out_im = perspective_image_tensor(in_im, coeffs, interpolation=torchvision.transforms.InterpolationMode.NEAREST)