2

所以，我正在做一个分割任务，我需要做的是将一个 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 通道的一个热图像。有没有一种简单的方法来构建这样一个颜色已知的矩阵。

4

1 回答 1

0

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

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]

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

答案的来源，如果需要，可以在那里找到示例。

于 2019-07-19T14:06:13.707 回答