一种无监督语义分割算法：Unsupervised Semantic Segmentation using Invariance and Equivariance in Clustering

论文题目：PiCIE: Unsupervised Semantic Segmentation using Invariance and Equivariance in Clustering

`1 摘要`

本文提出了一种新的无注释的语义分割框架。现成的聚类方法仅限于策划的、单标签的和以对象为中心的图像，而现实世界的数据主要是未策划的、多标签的和以场景为中心的。我们将聚类从图像扩展到像素，并为每个图像中的不同实例分配单独的聚类成员关系。然而，仅仅依赖于像素级的特征相似度并不能学习高级语义概念，以及对低级视觉线索的过度拟合。我们提出了一种将几何一致性作为归纳偏差的方法来学习光度和几何变化的不变性和等方差。有了我们新颖的学习目标，我们的框架可以学习高级语义概念。我们的方法，PiCIE（使用不变性和等方差的像素级特征聚类），是第一种能够在没有任何超参数调优或任务特定的预处理的情况下同时分割事物和东西类别的方法。

在本文中，我们朝着一个实际有用的无监督语义分割系统迈出了一步：我们提出了一种方法，能够以比现有技术更细的粒度分割所有像素，无论是thing还是stuff。我们的方法是基于一个直接的目标，它只编码了两个常识性的约束。首先，具有相似外观的像素（即，它们在一个学习到的特征空间中聚集在一起）应该被类似地标记，反之亦然。其次，像素标签应该与颜色空间变换不变，与几何变换等变。我们的结果表明，单独使用这两个目标，我们可以训练一个基于ConvNet的语义分割系统的端到端，没有任何标签。
我们发现，尽管它很简单，但我们的方法在这个任务上远远优于之前的工作，比现有技术的准确性提高了一倍多（图1，底部）。与之前的工作相比，我们的基于聚类的损失函数（上面的第一个目标）导致了一个更简单、更容易的学习问题，而之前的工作则试图学习参数像素分类器。但不变性和等方差目标是关键。它们允许卷积网络将跨尺度、姿态和颜色变化的像素连接在一起，而这是之前的系统无法做到的。
我们将无监督图像分割的任务表示为像素级聚类，其中每个像素都被分配给一个聚类。聚类通常需要一个良好的特征空间，但事先不存在这样的特征表示。因此，我们提出了一种与聚类联合学习特征表示的方法。PiCIE的整体管道，代表使用不变性和等方差的像素级特征聚类，如图2所示。我们将在下面描述我们的方法。

`2 理论介绍`

`2.1 A baseline clustering approach`

本文所解决的关键问题是，将图像聚类为类需要强特征表示，而对于训练强特征表示，则需要类标签。为了解决这个鸡和蛋的问题，最简单的解决方案是由深度聚类确定一个。交替使用当前的特征表示进行聚类，并使用聚类标签作为伪标签来训练特征表示。对于无监督的语义分割任务，人们也可以遵循类似的策略。唯一的区别是，我们需要使用一个嵌入函数 f ( θ ) f(θ) f(θ)来生成一个特征映射，为每个像素生成一个特征向量。分类器还必须对单个像素进行操作。然后，我们可以交替聚类像素特征向量来得到像素伪标签，并使用这些伪标签来训练像素特征表示。
具体地说，假设我们有一组未标记的图像 x i ， i = 1 ， … … ， n x_i，i = 1，……，n xi，i=1，……，n。假设我们的嵌入，用 f θ f_θ fθ表示，产生了一个特征张量 f θ ( x ) f_θ(x) fθ(x)。这将为图像 x x x中的每个像素p产生一个特征表示。用 f θ ( x ) [ p ] f_θ(x)[p] fθ(x)[p]表示这个像素级的特征表示。用 g w ( ⋅ ) g_w(·) gw(⋅)表示在这些像素特征向量上操作的分类器。然后，我们的基线方法在两个步骤之间交替进行：
1.使用当前的嵌入和k-means来聚类数据集中的像素。

其中， y i p y_{ip} yip为第 i i i幅图像中第 p p p个像素的聚类标签， µ k µ_k µk为第 k k k个聚类质心。（我们使用小批量k-means [39](Web-scale k-means clustering)）。

2.使用聚类标签标准交叉熵损失训练像素分类器。

其中 s k s_k sk是分类器 g w ( f θ ( x i ， p ) g_w(f_θ(x_i，p) gw(fθ(xi，p)输出的第k类分数。

`2.2 Non-parametric prototype-based classifiers`

上面的深度聚类启发框架使用了一个单独的、已学习的分类器。然而，在不断变化的伪标签的无监督设置中，联合训练分类器是具有挑战性的。一个训练不足的分类器可以将噪声梯度输入到特征提取器中，导致下一轮训练的噪声聚类。
因此，我们建议完全抛弃参数化像素分类器 g w g_w gw。相反，我们根据像素到kmeans估计的质心（“原型”[41]( Prototypical networks for few-shot learning)）的距离来标记像素。这就导致了以下目标的改变。

`2.3 Invariance and Equivariance`

将上述特征表示与聚类联合学习肯定会产生在特征空间中紧凑的聚类，但没有理由说这些集群必须是语义的。为了得到像素的语义分组，我们需要引入一个额外的归纳偏差。如果我们没有标签，这种归纳偏差又会是什么呢？
我们引入的归纳偏差是光度变换的不变性和几何变换的等方差：如果像素颜色轻微抖动，标记不应该改变，当图像被几何扭曲时，标记应该是类似的扭曲。具体地说，如果Y是上述图像的输出的语义标记为 x x x，如果 P P P和 G G G分别是光度变换和几何变换，则变换后的图像 G ( P ( x ) ) G(P(x)) G(P(x))的输出语义标记应为 G ( Y ) G(Y) G(Y)。
在联合聚类和学习框架中实现这个约束是很棘手的，因为每个图像都没有真实标签。伪真实标记本身来自于聚类，而聚类本身是由特征映射产生的，因此它本身对输入转换很敏感。因此，在这种情况下，不变性/等方差意味着两件事：第一，我们应该产生相同的簇，而不考虑转换，第二，预测的像素标签应该显示所期望的等方差。

`2.4 Invariance to photometric transformations`

我们首先讨论不变性的问题。对于数据集中的每一个图像 x i x_i xi，我们随机抽取两个光度变换， P i ( 1 ) P_i^{(1)} Pi(1)和 P i ( 2 ) P_i^{(2)} Pi(2)。这将为每个图像 x i xi xi中的每个像素 p p p产生两个特征向量：

然后，我们分别在两个“视图”中进行聚类，得到两组伪标签和质心：

给定这两组质心和这两组伪标签，我们使用了两组损失函数：
与前面一样，我们希望特征向量坚持聚类标签。现在我们有了两个视图，我们希望在每个视图中都是这样的：

因为我们假设聚类应该对光度转换是不变的，所以我们也希望来自一个视图的特征向量来匹配另一个视图的聚类标签和质心：

这个多视图框架和交叉视图损失实现了两件事。首先，通过强制来自一个转换的特征向量与另一个转换产生的标签保持一致，它鼓励网络学习特征表示，这些特征表示将被相同地标记，而不考虑任何光度变换。其次，通过强制相同的特征表示与两种不同的聚类解决方案保持一致，它鼓励了对这两种解决方案本身进行匹配，从而确保通过聚类发现的概念集对光度变换是不变的。

`2.4 Equivariance to geometric transformations`

房子和放大的房子应该贴上类似的标签，但可能产生截然不同的特征。更准确地说，放大内部的分割应该是原始分割的放大版本。这是几何变换（如随机作物）的等方差的概念，我们在接下来添加它。
为了学习关于几何变换的等方差，我们对每幅图像采样一个几何变换（具体来说，随机裁剪和水平翻转） G i G_i Gi。然后，在上述框架中，一个视图使用变换后的图像的特征向量，而另一个视图使用原始图像的变换后的特征向量：

其他的步骤也是完全相同的。这两个视图分别聚类，最终的训练目标是视图内目标和交叉视图目标的组合：

`3 结论`

本文介绍了一种基于聚类的无监督语义分割框架。我们的主要贡献是将几何一致性作为归纳偏差，以学习光度和几何变化的不变性和等方差。我们的新交叉视图丢失是简单但非常有效的学习高级视觉概念必要的分割事物的类别。我们的方法是第一个无监督的语义分割，它既适用于东西和事物类别，又没有严格的超参数调整或特定于任务的预处理。

`3 源码测试`

源码
问题1

File "/home/jinglai.whr/code/PiCIE/data/coco_eval_dataset.py", line 104, in
fine_to_coarse_map = np.vectorize(lambda x: fine_to_coarse_dict[x]) # not in-place.
KeyError: 185

解决方法：
这里报错的原因是因为json文件中的类别跟fine_to_coarse_dict不一致导致的，因此这里手动给它对齐。

运行结果：
数据按照官方的格式放置即可(注意这里val数据集需要对应标签数据，训练数据不需要)

test.py

import os
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision
import numpy as np
import matplotlib.pyplot as plt from torchvision import transforms
from modules import fpn
from PIL import Image
import cv2class Args:def __init__(self):passdef collate_eval(batch):indice = [b[0] for b in batch]image = torch.stack([b[1] for b in batch])if batch[0][-1] is not None:label = torch.stack([b[2] for b in batch])return indice, image, labelelse:return indice, imagedef get_score_histogram(label_trues, label_preds, n_class):hist = np.zeros((n_class, n_class))for lt, lp in zip(label_trues, label_preds):hist += _fast_hist(lt.flatten(), lp.flatten(), n_class)return histdef compute_dist(featmap, metric_function, euclidean_train=True):centroids = metric_function.module.weight.dataif euclidean_train:return - (1 - 2*metric_function(featmap)\+ (centroids*centroids).sum(dim=1).unsqueeze(0)) # negative l2 squared else:return metric_function(featmap)preprocess = transforms.Compose([transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])def show_neighbors(img_raw, q_loc, imglist_c, loclist_c):fig = plt.figure(figsize=(12, 12))fig.add_subplot(1, 3, 2)plt.imshow(img_raw)ax = plt.gca()rect = Rectangle(q_loc, 5, 5, linewidth=8, edgecolor='r', facecolor='none')ax.add_patch(rect)plt.show()fig = plt.figure(figsize=(20, 100))for i, (img, loc) in enumerate(zip(imglist_c, loclist_c)):fig.add_subplot(1, 5, i+1)plt.imshow(img)ax = plt.gca()_, h, w = locrect = Rectangle((4*w, 4*h), 5, 5, linewidth=8, edgecolor='r', facecolor='none')ax.add_patch(rect)plt.show()if __name__ == '__main__':args = Args()# Modelargs.pretrain = True args.ssl = False args.ssl_dir = False args.arch = 'resnet18'# Classifierargs.in_dim = 128args.K = 27# Dataloaderargs.data_root = ''picie_load = torch.load('picie.pkl', 'cpu')model = fpn.PanopticFPN(args)# model = nn.DataParallel(model)model = torch.nn.DataParallel(model, device_ids=[1])model.load_state_dict(picie_load['state_dict'])classifier = nn.Conv2d(args.in_dim, args.K, kernel_size=1, stride=1, padding=0, bias=True)classifier = nn.DataParallel(classifier, device_ids=[1])classifier.load_state_dict(picie_load['classifier1_state_dict'])# Important ! model.eval()classifier.eval()histogram_train = torch.load('picie_histogram_coco.pkl')histogram_train = histogram_train.astype(int)histogram_train = torch.tensor(histogram_train).topk(1, dim=0)[1].flatten()map_fn_train = lambda x: histogram_train[x]mapper_train = np.vectorize(map_fn_train)fig = plt.figure(figsize=(10, 10))# img_raw1 = Image.open('assets/BGCT-0824_20200617_1234_20220708093113_000016.png').convert('RGB')img = cv2.imread("assets/115.png", 0)img_raw1 = cv2.merge((img, img, img))img_raw2 = Image.open('assets/115.png').convert('RGB')fig.add_subplot(1, 3, 1)plt.imshow(img_raw1)fig.add_subplot(1, 3, 3)plt.imshow(img_raw2)plt.show()print('\t Animal\t\t\t\t\t\tBuilding')img1 = preprocess(img_raw1)[None, :]out1 = model(img1)out1 = F.normalize(out1, dim=1, p=2)prb1 = compute_dist(out1, classifier)lbl1 = prb1.topk(1, dim=1)[1]lbl1 = lbl1.squeeze(0).squeeze(0)img2 = preprocess(img_raw2)[None, :]out2 = model(img2)out2 = F.normalize(out2, dim=1, p=2)prb2 = compute_dist(out2, classifier)lbl2 = prb2.topk(1, dim=1)[1]lbl2 = lbl2.squeeze(0).squeeze(0)res = mapper_train(lbl1.cpu().numpy())# print(res)cv2.imwrite("result.png", res)from matplotlib import colorsfig = plt.figure(figsize=(10, 10))fig.add_subplot(1, 3, 1)plt.imshow(mapper_train(lbl1.cpu().numpy()))fig.add_subplot(1, 3, 3)plt.imshow(mapper_train(lbl2.cpu().numpy()))plt.show()print('\t Animal\t\t\t\t\t\tBuilding')f = plt.gcf() f.savefig(r'{}.png'.format(1))f.clear()cv2.imwrite("result1.png", mapper_train(lbl2.cpu().numpy()))from matplotlib.patches import Rectangle fig = plt.figure(figsize=(10, 10))fig.add_subplot(1, 3, 1)plt.imshow(img_raw1)ax = plt.gca()# w, hpos1  = (150, 50)rect1 = Rectangle(pos1, 5, 5, linewidth=5, edgecolor='r', facecolor='none')ax.add_patch(rect1)fig.add_subplot(1, 3, 3)plt.imshow(img_raw2)ax = plt.gca()# w, hpos2  = (80, 120)rect2 = Rectangle(pos2, 5, 5, linewidth=5, edgecolor='r', facecolor='none')ax.add_patch(rect2)plt.show()print('\t Animal\t\t\t\t\t\tBuilding')print('Image 1 shape       : ', img_raw1.size[::-1])print('Feature map 1 shape : ', lbl1.cpu().numpy().shape)print('Image 2 shape       : ', img_raw2.size[::-1])print('Feature map 2 shape : ', lbl2.cpu().numpy().shape)# h, wfpos1 = (pos1[1]*55//217, pos1[0]*73//289)fpos2 = (pos2[1]*55//217, pos2[0]*36//144)query1 = out1[:, :, fpos1[0], fpos1[1]]query2 = out2[:, :, fpos2[0], fpos2[1]]querys = torch.cat([query1, query2])print('Image 1 query position : ', fpos1)print('Image 2 query position : ', fpos2)imglist, loclist = torch.load('picie_retrieval_result_coco.pkl')print('\n\n\n')print('\t\t\t\t\t\t\tAnimal')show_neighbors(img_raw1, pos1, imglist[0], loclist[0])print('\n\n\n')print('\t\t\t\t\t\t\tBuilding')show_neighbors(img_raw2, pos2, imglist[1], loclist[1])