torchvision transforms自动clip图像数值范围坑

用 [1] 在自己数据集跑实验，需要读 .exr 后缀的 depth map（参照 [2] 配置好 opencv），期望是保持 depth map 原始的值不变。由于 torchvison.transforms 接受 PIL.Image 或 torch.Tensor 做输入，而 [1] 原本的代码是用 PIL.Image 读的，于是一开头用 cv2 读完之后转成 PIL.Image 套用原代码。但发现这样操作，depth map 的范围会自动被其 transforms clip 成 [0, 255]，导致训练不了。

对比：

import cv2
import numpy as np
from PIL import Image
import torch
from torchvision import transforms
import traceback

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

depth = cv2.imread("027.jpg.geometric.exr", cv2.IMREAD_UNCHANGED)
# 原始范围
print("original:", depth.shape, depth.min(), depth.max(), depth.dtype)
# 转 PIL.Image -> 值没变
depth_pil = Image.fromarray(depth)
print("PIL:", depth_pil.size, np.min(depth_pil), np.max(depth_pil))#, depth_pil.dtype)
# 转 torch.Tensor -> 值没变
depth_t = torch.from_numpy(depth)
print("torch:", depth_t.size(), depth_t.min(), depth_t.max(), depth_t.dtype)

# 情况 1：对 PIL.Image 用，会（先 clip 到 [0, 255] 再）归一化到 [0, 1]
trfm1 = transforms.Compose([
    transforms.Grayscale(num_output_channels=1),  # 会 clip
    transforms.ToTensor()  # 会归一化到 [0, 1]
])
dmap1 = trfm1(depth_pil)
print("dmap1:", dmap1.min(), dmap1.max())

# 情况 2：对 torch.Tensor 用，报错
try:
    dmap2 = trfm1(depth_t.unsqueeze(0).tile((3, 1, 1)))
    print("dmap2:", dmap2.min(), dmap2.max())
except:
    print('\n' + '<' * 7)
    print(traceback.format_exc(), end="")
    print('>' * 7 + '\n')

# 情况 3：对 PIL.Image 用，去掉 `ToTensor`，只 clip 没归一化
trfm2 = transforms.Grayscale(num_output_channels=1)
dmap3 = trfm2(depth_pil)
print("dmap3:", np.min(dmap3), np.max(dmap3))

# 情况 4：对 torch.Tensor 用（要变换成 [3, H, W] 先），保持原值
dmap4 = trfm2(depth_t.unsqueeze(0).tile((3, 1, 1)))
print("dmap4:", dmap4.min(), dmap4.max())

• 输出

original: (1276, 717) 0.0 5640.0 float32
PIL: (717, 1276) 0.0 5640.0
torch: torch.Size([1276, 717]) tensor(0.) tensor(5640.) torch.float32
dmap1: tensor(0.) tensor(1.)

<<<<<<<
Traceback (most recent call last):
  File "test.py", line 69, in <module>
    dmap2 = trfm1(depth_t.unsqueeze(0).tile((3, 1, 1)))
  File "/home/tyliang/miniconda3/envs/pt110/lib/python3.8/site-packages/torchvision/transforms/transforms.py", line 61, in __call__
    img = t(img)
  File "/home/tyliang/miniconda3/envs/pt110/lib/python3.8/site-packages/torchvision/transforms/transforms.py", line 98, in __call__
    return F.to_tensor(pic)
  File "/home/tyliang/miniconda3/envs/pt110/lib/python3.8/site-packages/torchvision/transforms/functional.py", line 114, in to_tensor
    raise TypeError('pic should be PIL Image or ndarray. Got {}'.format(type(pic)))
TypeError: pic should be PIL Image or ndarray. Got <class 'torch.Tensor'>
>>>>>>>

dmap3: 0 255
dmap4: tensor(0.) tensor(5639.4360)

所以如果是 depth map 这种非 RGB 图像（数值范围不是 [0, 255]），但当成图像读、transform 时，考虑转成 torch.Tensor 再传进 torchvision transforms，而且要验证在 data loader 一套操作之后、传给模型 inference 之前，数值范围对不对。

References

1. antocad/FocusOnDepth
2. Open EXR files, how to enable? #21928