How to save an `exr` from a pytorch tensor in Python?

Previously, there was a function torchvision.utils.save_float_image with which it was possible to store an .exr file from a pytorch tensor. This function is gone in the current relase (0.17). Now, there is only the function torchvision.utils.save_image (https://pytorch.org/vision/stable/generated/torchvision.utils.save_image.html#torchvision.utils.save_image). But when I try to execute

with open(os.path.join('folder/', `foo.exr`), "wb") as fout:
    save_image(tensor, fout)

I'm receiving the error "unknown file extension *.exr". So, what can we do now?

Solution

The Python Imaging Library doesn't support exr files, ref. The only function available for saving images in Torchvision is save_image, it uses PIL to save the file. As far back as the first documented version (0.8), there was never a "save_float_image" function.

What you may use instead is imageio (pip install imageio): You can get a reference by looking at the unit test for saving an exr file. You can save to disk by first transposing and converting to numpy array:

arr = tensor.permute(1,2,0).numpy()
imageio.imwrite('out_path.exr', arr)

Where out_path has the .exr extension file, otherwise you must use the option: extension=".exr"

Here is a minimal reproducible example:

import os
os.environ["OPENCV_IO_ENABLE_OPENEXR"] = "1"

tensor = torch.rand(3,100,150)
arr = tensor.permute(1,2,0).numpy()
imageio.imwrite('out_path.exr', arr)