Rescale bounding boxes relatively to original image size

I hope everyone’s doing great. Could you please give me a hand with the following problem. I believe there’s no such function in pytorch, opencv etc, so I am trying to implement it manually.

Context: I collect a batch of frames and push them to GPU, keep them there. Then, I have a couple of networks (YOLOs). In order for this to work, I need to resize images on GPU, to the size YOLO expects keeping aspect ratio (there's padding that fills the rest of the image with grey colour). Done as described here: https://discuss.pytorch.org/t/resize-images-torch-tensor-already-loaded-on-gpu-keeping-their-aspect-ratio/83172/9

What I need to do now: rescale bounding box coordinates so that they are relative to the original size of the frames, not the resized ones given to the network for detection. This is how I am doing it now: code

@staticmethod
def rescale_bounding_box(detections: dict, current_dim: int, original_shape: tuple) -> dict:
    """
    Rescale bounding boxes predicted for image(s) of the size current_dim to the original_shape
    :param detections:
    :param current_dim:
    :param original_shape:
    :return:
    """
    original_h, original_w = original_shape
    # Added padding (check if the image was tall of short)
    pad_x = max(original_h - original_w, 0) * (current_dim / max(original_shape))
    pad_y = max(original_w - original_h, 0) * (current_dim / max(original_shape))

    # Image size after padding's been removed
    unpad_h = current_dim - pad_y
    unpad_w = current_dim - pad_x

    output = dict()
    for img_batch_index, predictions in detections.items():
        output[img_batch_index] = list()
        # On each image in the batch there's a number of detections that need to be rescaled
        for prediction in predictions:
            new_left = int((prediction[0] - pad_x // 2) * (original_w / unpad_w))
            new_top = int((prediction[1] - pad_y // 2) * (original_h / unpad_h))
            new_right = int((prediction[2] - pad_x // 2) * (original_w / unpad_w))
            new_bot = int((prediction[3] - pad_y // 2) * (original_h / unpad_h))
            obj_score = round(prediction[4], 4)
            conf = round(prediction[5], 4)
            index = int(prediction[6])

            print("Rescaled:", new_left, new_top, new_right, new_bot)
            # Save modified results
            assert all((new_left < new_right, new_top < new_bot)), "Coordinates rescaled wrong"
            assert all((new_left > 0, new_top > 0, new_right > 0, new_bot > 0)), "Coordinates rescaled wrong"
            assert all((new_right <= original_w, new_bot <= original_h)), "Coordinates rescaled wrong"
            output[img_batch_index].append(
                [new_left, new_top, new_right, new_bot, obj_score, conf, index]
            )

    return output

Unfortunately, it doesn’t pass the assert statements. My new_top value is usually < 0, new_right can go beyond the right edge.

detections - is a dictionary, where key - index of the image in the batch (1 for images), and values - list of lists. Each nested list is a detection predicted by the network.

If anyone could please give me a hand with it. Spent a lot of time on it, no success.

Thanks in advance.