What's the Fastest Way to Convert a Tensor to an Image File?

When serving generative image models in a production environment, a tensor representation of an image needs to be converted to a common image format like PNG, JPEG, or WEBP. However this conversion can be costly and those interested in super speedy inference need to know what the fastest way to get a file back to their users.

When making choices about how to deliver these image files, there will be a trade off between inference speed, image quality, and the number of bytes in the file.

In this post, I'll lay out some options to choose from share some benchmarks. All of my benchmarks are taken on a c6a.4xlarge EC2 instance. I use only a single 702x1248 image for each benchmark:

See the appendix for list of versions of packages I'm using in my benchmarks.

Library Options:

1. Python Image Library
2. Torch Vision
3. OpenCV

Formats:

1. PNG
2. WEBP

I won't consider lossy formats like JPEG because I think diffusion models services shouldn't risk any quality degradation as a result of lossy image compression.

PNG Benchmarks

Here's a sample of the code I used for Python Image Library (PIL). I assume the reader can figure out how to modify to produce all the results I report in the table that comes after the code.

import io
import time
from PIL import Image
import torchvision.transforms.functional as F

path = "/home/user/00000.png"
pil_image = Image.open(path)
pil_image = pil_image.convert("RGB")

image_tensor = F.to_tensor(pil_image)

def pil_png(out):
    pil_image: Image = F.to_pil_image(image_tensor)
    pil_image.save(out, format="PNG", compress_level=4)

t0 = time.time()
for i in range(100):
    out = io.BytesIO()
    pil_png(out)

print(f"Bytes in file:{len(out.getvalue())}")
print(f"Average time: {(time.time() - t0) / 100}")

Results for Python Image Library:

I also tested Torch Vision, which has a png encoder. Code for Torch Vision:

import io
import time
from PIL import Image
import torch
import torchvision.transforms.functional as F
import torchvision.io

path = "/home/user/00000.png"
pil_image = Image.open(path)
pil_image = pil_image.convert("RGB")

image_tensor = F.to_tensor(pil_image) * 255.0
image_tensor = image_tensor.to(torch.uint8)

def tv_png():
    return torchvision.io.encode_png(image_tensor, compression_level = 1)

t0 = time.time()
for i in range(100):
    val = tv_png()

print(f"Bytes in file:{len(val)}")
print(f"Average time: {(time.time() - t0) / 100}")

Results for Torch Vision:

My 3rd library that I considered was OpenCV. Here's the code I used to test that option:

path = "/home/aalangfo/00000.png"
pil_image = Image.open(path)
pil_image = pil_image.convert("RGB")

image_tensor = F.to_tensor(pil_image) * 255.0
image_tensor = image_tensor.to(torch.uint8)
image_tensor = image_tensor.numpy()
image_tensor = image_tensor.transpose((1, 2, 0))

def cv_png():
    result, buffer = cv2.imencode('.png', image_tensor, [cv2.IMWRITE_PNG_COMPRESSION, 0])
    return buffer

t0 = time.time()
for i in range(100):
    out = cv_png()

print(f"Bytes in file:{len(out.tobytes())}")
print(f"Average time: {(time.time() - t0) / 100}")

Results for OpenCV:

WebP Benchmarks

WebP is an image format developed by Google that allows for both lossless and lossy compression for images. It was designed to offer more efficient compression than JPEG and PNG. It also supports animation like GIF does, but improves on GIFs compression.

The downside of WebP is that it may not be supported in older browsers and devices. There's also a few more levers in the WebP encoding spec, so it may take a bit more time to find the right settings for you. You should explore a list of those options here before reviewing the results.

I will stick with only lossless encoding for WebP for the same reasons mentioned above.

Here's my code for benchmarking WebP with Python Image Library (PIL):

import io
import time
from PIL import Image
import torchvision.transforms.functional as F

path = "/home/user/00000.png"
pil_image = Image.open(path)
pil_image = pil_image.convert("RGB")

image_tensor = F.to_tensor(pil_image) * 255.0
image_tensor = image_tensor.to(torch.uint8)

def pil_webp(out):
    pil_image: Image = F.to_pil_image(image_tensor)
    pil_image.save(out, format="WebP", lossless=True, quality=0, method=0)

t0 = time.time()
for i in range(100):
    out = io.BytesIO()
    pil_webp(out)

print(f"Bytes in file:{len(out.getvalue())}")
print(f"Average time: {(time.time() - t0) / 100}")

Results for PIL:

As of this post, torch vision does not support writing images using WebP.

While OpenCV does support WebP, the only way to do lossless image encoding is to set the quality level higher than 100; For the curious, I was able to get the following results at quality level 101:

• Average time: 0.462
• Bytes in file: 811340

Conclusion

Generally OpenCV is the fastest way to encode images. Because of lack of support for WebP encoding flags in OpenCV, OpenCV with PNG at compression level 0 or 1 seems like a great way to go.

If number of bytes are really important for the use case, WebP provides superior lossless compression to PNG, saving 100s of KB or 1s of MB depending on the configuration.