如何从PyTorch张量保存`exr`文件？

Question

之前有一个函数叫做 torchvision.utils.save_float_image，可以用来把 pytorch 的张量保存成 .exr 文件。但是在当前的版本（0.17）中，这个函数已经不见了。现在只剩下一个函数 torchvision.utils.save_image（https://pytorch.org/vision/stable/generated/torchvision.utils.save_image.html#torchvision.utils.save_image）。可是当我尝试执行

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

时，出现了错误提示“未知的文件扩展名 *.exr”。那么，我们现在该怎么办呢？

Answer 1

Python图像库不支持exr文件，参考链接。在Torchvision中，唯一可以用来保存图像的函数是save_image，它使用PIL来保存文件。从第一个有记录的版本开始（0.8），就没有“save_float_image”这个函数。

你可以使用imageio（通过pip install imageio安装）来代替：你可以通过查看单元测试来了解如何保存exr文件。你可以先转置并转换为numpy数组，然后保存到磁盘：

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

其中out_path是带有.exr扩展名的文件，否则你必须使用选项：extension=".exr"

---

下面是一个简单的可复现示例：

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)