resize_opencv_vs_pil_vs_pytorch.py

#%%
import cv2
import torch
import numpy as np
from PIL import Image
from torchvision.transforms.v2.functional import resize as tv_resize, InterpolationMode

#%%
Hin, Win = 256, 276

def make_uint8_images(Hin, Win):
    tensor_img = torch.randint(0, 256, size=(3, Hin, Win), dtype=torch.uint8)
    # Note: cv2 arrays are BGR, not RGB, but we don't care
    cv2_img = tensor_img.numpy().transpose(1, 2, 0)  # CHW -> HWC
    pil_img = Image.fromarray(cv2_img)

    return tensor_img, cv2_img, pil_img

def to_tensor(cv2_or_pil_img):
    if isinstance(cv2_or_pil_img, Image.Image):
        cv2_or_pil_img = np.asarray(cv2_or_pil_img)
    return torch.from_numpy(cv2_or_pil_img).permute(2, 0, 1)


def assert_close(a, b, atol=1, allowed_percent=0):
    # Assert that a and b differ by more than atol on no more than allowed_percentb % entries
    # Allowed_percent is expected in [0, 100]
    # When allowed_precent is 0 (default), we assert that all pixels differ by at most atol

    if allowed_percent == 0:
        torch.testing.assert_close(a, b, rtol=0, atol=atol)
        return

    abs_diff = (a.float() - b.float()).abs()
    actual_percent = (abs_diff > atol).float().mean() 
    if actual_percent > allowed_percent / 100:
        raise AssertionError(f"{actual_percent:.1%} pixels differ by more than {atol=}!")

def assert_exactly_equal(*args, **kwargs):
    torch.testing.assert_close(*args, **kwargs, rtol=0, atol=0)

def tv_resize_on_floats(uint8_tensor_img, *args, **kwargs):
    # Convert uint8 to float, resize, then round back to uint8.
    # This is what the v1 version of resize (torchvision.transforms.functional.resize) does
    # This is mathematically more correct for bicubic mode, but slower.
    # It doesn't matter in practice for models, you can just use the faster version on uint8
    return tv_resize(uint8_tensor_img.float(), *args, **kwargs).round().clamp(0, 255).to(torch.uint8)
# %%
# Bilinear interpolation, no antialiasing
# ---------------------------------------
for Hout, Wout in [(120, 128), (500, 480)]:  # Downsampling, upsampling
    tensor_img, cv2_img, _ = make_uint8_images(Hin, Win)

    out_cv2_linear = cv2.resize(cv2_img, (Wout, Hout), interpolation=cv2.INTER_LINEAR)
    out_cv2_linear_exact = cv2.resize(cv2_img, (Wout, Hout), interpolation=cv2.INTER_LINEAR_EXACT)
    out_tv = tv_resize(tensor_img, size=(Hout, Wout), interpolation=InterpolationMode.BILINEAR, antialias=False)
    out_tv_float = tv_resize_on_floats(tensor_img, size=(Hout, Wout), interpolation=InterpolationMode.BILINEAR, antialias=False)

    assert_close(to_tensor(out_cv2_linear), to_tensor(out_cv2_linear_exact), atol=1)
    assert_close(to_tensor(out_cv2_linear), to_tensor(out_cv2_linear_exact), atol=0, allowed_percent=20)
    assert_close(out_tv, to_tensor(out_cv2_linear), atol=1) 
    assert_close(out_tv, to_tensor(out_cv2_linear), atol=0, allowed_percent=15)
    assert_close(out_tv, to_tensor(out_cv2_linear_exact), atol=1) 
    assert_close(out_tv, to_tensor(out_cv2_linear_exact), atol=0, allowed_percent=15)
    assert_exactly_equal(out_tv, out_tv_float) 

# %%
# Bilinear interpolation, with antialiasing
# -----------------------------------------
for Hout, Wout in [(120, 128), (500, 480)]:  # Downsampling, upsampling
    tensor_img, _, pil_img = make_uint8_images(Hin, Win)

    out_pil = pil_img.resize((Wout, Hout), Image.BILINEAR)
    out_tv = tv_resize(tensor_img, size=(Hout, Wout), interpolation=InterpolationMode.BILINEAR, antialias=True)
    out_tv_float = tv_resize_on_floats(tensor_img, size=(Hout, Wout), interpolation=InterpolationMode.BILINEAR, antialias=True)

    assert_close(out_tv, to_tensor(out_pil), atol=1) 
    assert_close(out_tv, to_tensor(out_pil), atol=0, allowed_percent=20)
    assert_exactly_equal(out_tv, out_tv_float)
# %%
# Bicubic interpolation, no antialiasing
# --------------------------------------
for Hout, Wout in [(120, 128), (500, 480)]:  # Downsampling, upsampling
    tensor_img, cv2_img, _ = make_uint8_images(Hin, Win)

    out_cv2 = cv2.resize(cv2_img, (Wout, Hout), interpolation=cv2.INTER_CUBIC)
    out_tv = tv_resize(tensor_img, size=(Hout, Wout), interpolation=InterpolationMode.BICUBIC, antialias=False)
    out_tv_float = tv_resize_on_floats(tensor_img, size=(Hout, Wout), interpolation=InterpolationMode.BICUBIC, antialias=False)

    assert_close(out_tv, to_tensor(out_cv2), atol=1, allowed_percent=7)
    assert_close(out_tv, to_tensor(out_cv2), atol=0, allowed_percent=30)
    assert_close(out_tv, out_tv_float, atol=1, allowed_percent=7)
    assert_close(out_tv, out_tv_float, atol=0, allowed_percent=30)
    assert_close(out_tv_float, to_tensor(out_cv2), atol=1)
    assert_close(out_tv_float, to_tensor(out_cv2), atol=0, allowed_percent=1)

# %%
# Bicubic interpolation, with antialiasing
# ----------------------------------------
for Hout, Wout in [(120, 128), (500, 480)]:  # Downsampling, upsampling
    tensor_img, _, pil_img = make_uint8_images(Hin, Win)

    out_pil = pil_img.resize((Wout, Hout), Image.BICUBIC)
    out_tv = tv_resize(tensor_img, size=(Hout, Wout), interpolation=InterpolationMode.BICUBIC, antialias=True)
    out_tv_float = tv_resize_on_floats(tensor_img, size=(Hout, Wout), interpolation=InterpolationMode.BICUBIC, antialias=True)

    assert_close(out_tv, to_tensor(out_pil), atol=2)
    assert_close(out_tv, to_tensor(out_pil), atol=1, allowed_percent=0.5)
    assert_close(out_tv, to_tensor(out_pil), atol=0, allowed_percent=1)
    assert_close(out_tv, out_tv_float, atol=1, allowed_percent=2)
    assert_close(out_tv, out_tv_float, atol=0, allowed_percent=25)
    assert_close(out_tv_float, to_tensor(out_pil), atol=2, allowed_percent=1)
    assert_close(out_tv_float, to_tensor(out_pil), atol=1, allowed_percent=2)

# %%
# Nearest and nearest-exact interpolation
# ---------------------------------------
for Hout, Wout in [(120, 128), (500, 480)]:  # Downsampling, upsampling
    tensor_img, cv2_img, pil_img = make_uint8_images(Hin, Win)

    out_cv2_nearest = cv2.resize(cv2_img, (Wout, Hout), interpolation=cv2.INTER_NEAREST)
    out_cv2_nearest_exact = cv2.resize(cv2_img, (Wout, Hout), interpolation=cv2.INTER_NEAREST_EXACT)
    out_cv2_area = cv2.resize(cv2_img, (Wout, Hout), interpolation=cv2.INTER_AREA)

    out_pil = pil_img.resize((Wout, Hout), Image.NEAREST)

    out_tv_nearest = tv_resize(tensor_img, size=(Hout, Wout), interpolation=InterpolationMode.NEAREST)
    out_tv_float_nearest = tv_resize(tensor_img, size=(Hout, Wout), interpolation=InterpolationMode.NEAREST)
    out_tv_nearest_exact = tv_resize(tensor_img, size=(Hout, Wout), interpolation=InterpolationMode.NEAREST_EXACT)
    out_tv_float_nearest_exact = tv_resize(tensor_img, size=(Hout, Wout), interpolation=InterpolationMode.NEAREST_EXACT)

    assert_exactly_equal(out_tv_nearest, to_tensor(out_cv2_nearest))
    assert_exactly_equal(out_tv_nearest, out_tv_float_nearest)

    assert_close(to_tensor(out_cv2_nearest_exact), to_tensor(out_pil), atol=0, allowed_percent=5)
    assert_close(out_tv_nearest_exact, to_tensor(out_pil), atol=0, allowed_percent=5)
    assert_close(out_tv_nearest_exact, to_tensor(out_cv2_nearest_exact), atol=0, allowed_percent=7)
    assert_exactly_equal(out_tv_nearest_exact, out_tv_float_nearest_exact)
# %%