如何在 PyTorch 中将 RGB 图像编码为 n_class One 热张量

Posted 2023-03-29

【中文标题】如何在 PyTorch 中将 RGB 图像编码为 n_class One 热张量

【英文标题】：How to encode a RGB Image to an n_class One Hot Tensor in PyTorch

【发布时间】：2019-11-28 07:24:52

【问题描述】：

所以，我正在做一个分割任务，我需要做的是将一个 RGB 图像转换为一个 n 通道一个热矩阵，以用作 U-Net 模型的标签

我所做的是计算图像中的颜色。数量或颜色相当于类的数量。

我在 PerPixelClassMatrix 函数中尝试做的是遍历 Image 并制作一个 0 和 1 的 n 维矩阵，因为我有每个像素的颜色和类别。

import glob
from tqdm import tqdm
import numpy as np


class HotEncoder():
    def __init__(self, dir, extension, is_binary=True):
        self.dir = dir
        self.extension = extension
        self.is_binary = is_binary
        if is_binary:
            self.color = (0, 0, 0): 1, (255, 255, 255): 2
        else:
            self.color = dict()

    def gen_colors(self):
        """Iterates through the entire dataset and finds the total colours
            in the images so that they can be used to one hot the image matrix
            for the training data"""
        if self.is_binary:
            return self.color
        else:
            n_color=1
            images = glob.glob(self.dir + '/*.' + self.extension)
            for img in tqdm(images, desc="Generating Color Pallte to Hot Encode"):
                image = skimage.io.imread(img)
                shape_ = image.shape
                for x in range(shape_[0]):
                    for y in range(shape_[1]):
                        clr= tuple(image[x][y][:])
                        if clr not in self.color.keys():
                            self.color.update(n_color: clr)
                            n_color+=1
                        else:
                            pass
        return self.color

    def PerPixelClassMatrix(self, Image):
        """Takes an Image and returns a per pixel class
            identification map"""
        class_list= []
        class_mat= np.array([])
        shape_= Image.shape
        for x in range(shape_[0]):
            for y in range(shape_[1]):
                clr= tuple(Image[x][y][:])
                if clr in self.color.keys():
                    class_list.append(self.color[clr])
                else:
                    pass
        return class_list

我不想运行整个循环来生成一个 n 通道的热图像。有没有一种简单的方法来构建这样一个颜色已知的矩阵。

【问题讨论】：

您想转换输出/标签以适应损失吗？如果是这样，有一个更简单的方法来做到这一点。

是的，你能告诉我更好的方法吗？ ^u^

【参考方案1】：

如果你想计算图像分割损失，你可以这样做：

output = model(input)  # raw logit output in shape [1, 3, 512, 512]
loss = criterion(F.log_softmax(output,1), target)  # target in shape [1, 512, 512]

Target 将包含带有您的掩码索引的标签[0, N)。我假设您的输入图像是 3 通道 RGB。

Source of the answer，如果需要，可以在那里找到示例。