pytorch tensor计算三通道均值

    今天用pytorch处理图像时，涉及到了计算均值的问题，整理一下解决思路。

第一种思路：tensor转换为numpy再进行处理。

import torch
import cv2
img = cv2.imread("image path")
tensor_img = torch.from_numpy((img[:, :, ::-1] / 255.0)[None, ...].transpose(0, 3, 1, 2)).cuda()

...

numpy_img = (tensor_img.detach().cpu().numpy().transpose(2, 3, 1, 0).squeeze() * 255)[:, :, ::-1]

ave_color = np.mean(numpy_img , axis=(0, 1))

如果图像里有0值，不想计入运算：

numpy_img[numpy_img == 0] = np.nan
ave_color = np.nanmean(numpy_img, axis=(0, 1))

由于tensor和numpy来回转换会消耗资源、性能。又查了一番，直接在tensor中计算（非零均值计算）

reshape_tensor_img = tensor_img.view(tensor_img.size(0), tensor_img.size(1), -1)
ave_color = reshape_tensor_img.mean(2)

# mean value without 0
non_zero_img = reshape_tensor_img[reshape_tensor_img.nonzero(as_tuple=True)]
ave_color = non_zero_img.view(reshape_tensor_img.size(0), reshape_tensor_img.size(1), -1).mean(2)

计算完均值，想要加法运算的时候也会碰上一点维度上的麻烦，需要维度转换一下。

reshape_ave_color = ave_color.view(ave_color.size(0), ave_color.size(1), 1, 1)
add_img = tensor_img + reshape_ave_color