Last active
January 23, 2022 18:37
-
-
Save rsomani95/de0045c45f76b4f2ceb3a8a40fb186a4 to your computer and use it in GitHub Desktop.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
import os | |
import shutil | |
import sys | |
from pathlib import Path | |
import cv2 | |
import numpy as np | |
import nvidia.dali.fn as fn | |
import nvidia.dali.types as types | |
import torch, torchvision, nvidia.dali, numpy, cv2 | |
ver = lambda x: f"{x.__name__.ljust(14)}: {x.__version__.rjust(3)}" | |
print("\n".join(list(map(ver, [torch, torchvision, nvidia.dali, numpy, cv2])))) | |
from nvidia.dali import pipeline_def | |
from PIL import Image | |
from PIL.ImageDraw import ImageDraw | |
from torchvision.transforms.functional import InterpolationMode | |
from torchvision.transforms.functional import resize as torch_resize | |
from torchvision.transforms.functional import to_pil_image, to_tensor | |
mkdir = lambda x: Path(x).mkdir() | |
def create_test_img(size=(128, 128), outline_width: int = 3): | |
assert size[0] == size[1], f"Square image dimensions expected" | |
im = Image.new("RGB", size, color=(255, 255, 255)) | |
d = ImageDraw(im) | |
qtr = int(size[0] / 4) | |
pad = int(0.015625 * size[0]) | |
d.ellipse((qtr, qtr, qtr * 3, qtr * 3), fill=None, outline=0, width=outline_width) | |
d.rectangle( | |
((pad, pad), (size[0] - pad, size[1] - pad)), fill=None, outline=0, width=outline_width | |
) | |
return im | |
def resize_opencv(img: Image.Image, size=(128, 128), interpolation=cv2.INTER_LINEAR): | |
# im = cv2.cvtColor(np.array(img), cv2.COLOR_BGR2RGB) | |
im = np.array(img) | |
im = cv2.resize(im, size, interpolation) | |
return Image.fromarray(im) | |
def resize_torch( | |
img: Image.Image, | |
size=(128, 128), | |
interpolation=InterpolationMode.BICUBIC, | |
antialias: bool = None, | |
): | |
im = torch_resize(to_tensor(img), size, interpolation, antialias=antialias) | |
# return im | |
return to_pil_image(im) | |
def resize_dali(img: str, size=(128, 128), interpolation=types.INTERP_LINEAR, device="gpu"): | |
# Setup directory structure for DALI | |
if Path("tmp-test").exists(): | |
shutil.rmtree("tmp-test") | |
mkdir("tmp-test") | |
mkdir("tmp-test/0") | |
img.save("tmp-test/0/test-img.png", quality=100) | |
# Define pipeline | |
@pipeline_def | |
def test(): | |
imgs, _ = fn.readers.file(file_root="tmp-test") | |
imgs = fn.decoders.image(imgs, device="mixed" if device == "gpu" else "cpu") | |
imgs = fn.resize( | |
imgs, | |
device="gpu", | |
resize_x=size[0], | |
resize_y=size[0], | |
mode="stretch", | |
interp_type=interpolation, | |
) | |
return imgs, _ | |
# Build Pipeline | |
pipe = test(batch_size=1, num_threads=1, device_id=0) | |
pipe.build() | |
# Run Pipeline and get outputs | |
img = pipe.run()[0] | |
img = img.as_cpu().as_array()[0] # 0th element of the batch | |
assert img.ndim == 3 | |
return Image.fromarray(img) | |
img = create_test_img((2048, 2048), outline_width=6) | |
size = (128, 128) # Output size | |
# -- DALI | |
# --- Correct outputs | |
resize_dali(img, size, interpolation=types.INTERP_GAUSSIAN) | |
resize_dali(img, size, interpolation=types.INTERP_LANCZOS3) | |
resize_dali(img, size, interpolation=types.INTERP_TRIANGULAR) | |
# --- Wrong outputs | |
resize_dali(img, size, interpolation=types.INTERP_CUBIC) | |
resize_dali(img, size, interpolation=types.INTERP_LINEAR) | |
resize_dali(img, size, interpolation=types.INTERP_NN) | |
# -- PIL | |
# --- Correct outputs | |
img.resize(size, Image.BILINEAR) | |
img.resize(size, Image.BICUBIC) | |
img.resize(size, Image.LANCZOS) | |
img.resize(size, Image.BOX) | |
img.resize(size, Image.HAMMING) | |
# --- Wrong outputs | |
# None! | |
# -- Torch | |
# --- Correct outputs | |
resize_torch(img, size, interpolation=InterpolationMode.BILINEAR, antialias=True) | |
# --- Wrong outputs | |
resize_torch(img, size, interpolation=InterpolationMode.BILINEAR) | |
resize_torch(img, size, interpolation=InterpolationMode.BICUBIC) | |
resize_torch(img, size, interpolation=InterpolationMode.BICUBIC, antialias=True) | |
# -- OpenCV | |
# --- Correct outputs | |
# None! | |
# --- Wrong outputs | |
resize_opencv(img, size, cv2.INTER_NEAREST) | |
resize_opencv(img, size, cv2.INTER_LINEAR) | |
resize_opencv(img, size, cv2.INTER_AREA) | |
resize_opencv(img, size, cv2.INTER_CUBIC) | |
resize_opencv(img, size, cv2.INTER_LANCZOS4) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment