pytorch tensor计算三通道均值

最新推荐文章于 2024-05-01 01:09:22 发布

IFT_jason

最新推荐文章于 2024-05-01 01:09:22 发布

阅读量3.5k

点赞数 1

分类专栏：计算机视觉 Pytorch 文章标签：深度学习 pytorch 计算机视觉

本文链接：https://blog.csdn.net/infinite_jason/article/details/115242839

版权

计算机视觉同时被 2 个专栏收录

2 篇文章

订阅专栏

Pytorch

2 篇文章

订阅专栏

这篇博客介绍了在PyTorch中处理图像时计算均值的两种方法：通过numpy转换和直接在tensor上操作。针对含有0值的情况，提出了排除0值计算均值的解决方案，并详细展示了维度调整以进行加法运算的过程。内容涵盖了tensor与numpy之间的转换、非零元素均值计算以及维度匹配等技巧。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

今天用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