一点就分享系列(理解篇6）BBA出品 Painter—>SegGPT，主打推理的图生图视觉模型
文章目录
- 一点就分享系列(理解篇6）BBA出品 Painter—>SegGPT，主打推理的图生图视觉模型
- 前言背景
- 下面结合代码使用和论文阅读进行一次浅析，写的比较匆忙，欢迎批评和指正。
- 一、Painter 和SegGPT
- 1.1 Painter 简介
- 1.2 SegGPT
- 1.3 Painter—->SegGPT
- 1.4 Painter详解
- 1.4.1 数据集介绍
- 1.4.2 项目快速推理使用——展示迁移能力
- 1.4.1 模型结构以及训练设计（本节提到的详细代码可以看源码，这里DEIT和VIT比较成熟的原理不作赘述）
- 1.4.2 训练损失函数
- 2. 额外补充
- 2.1 细节实现
- 2.1.1语义分割：
- 2.1.2 关键点检测
- 2.1.2 全景分割：
- 2.1.3 图像重建
- 2.2 实现的代码优化细节（必看）
- 2.2.1.patch merging
- 2.2.2 encoder
- 2.2.3 Head
- 总结
前言背景
今天继续AIGC领域的学习，是由BAAI发布的两个2023 CVPR论文，论文地址分别是Images Speak in Images:
A Generalist Painter for In-Context Visual Learning，SegGPT
论文项目地址：github， hugging_Face_Painter, hugging_Face_Seggpt
虽然目前作者开源这个项目并加入这两部分，整个代码结构是建立在一些开源的CV框架上去开发的如detectron2/mmcv等，算法上的想法很大胆，利用MAE的思路去统一多种视觉任务，但是这俩部分的代码只有Painter算是完全开源的，SegGPT只是调用的webui，并没有开源；并且该项目代码对于宣传的多任务的调用成熟度并不够，需要自己再完善逻辑或等作者更新
下面结合代码使用和论文阅读进行一次浅析，写的比较匆忙，欢迎批评和指正。
一、Painter 和SegGPT
1.1 Painter 简介
Painter是一个图生图的纯视觉模型，其特点如下：
- 纯图像通用模型： 以输入图和其GT作为Propmpt，重新定义多任务为输出为图像，以图像为中心；
- 上下文推理能力：zero-shot ，泛化性极强；
- 任务多样性：支持7种视觉经典任务：语义分割、全景分割、深度估计、降噪、去雨、图像增强、关键点检测都在数据集上取得了”具备竞争力”性能；
- 任务的拓展性，就像上篇SAM的介绍，这个模型依旧依赖于PROMPT，这就带来了更多模型组合的可能，同时对工业界的实时需求部署也提升了难度。
1.2 SegGPT
SegGPT是致力于分割任务的通用视觉模型，其特点如下：
- 上下文推理能力:zero-shot ，泛化性极强;
- 同样具备任务拓展性和图生图的通用模型定义，看作是Painter的分支即可。
1.3 Painter—->SegGPT
正如官方展示的 SegGPT作为单一的分割任务是集成到Painter这个多任务模型的分支中，Painter如同“GPT3+”一样作为图像的propmpt去完成任务的自动化识别和输出，如官方给出的图例，SegGpt作为其中的一个分割任务分支。
1.4 Painter详解
1.4.1 数据集介绍
文章中提到针对7类视觉任务：uses COCO, ADE20K, NYUDepthV2, Synthetic Rain Datasets, SIDD, and LoL datasets.开源代码中提供了详细的使用脚本。
1.4.2 项目快速推理使用——展示迁移能力
提示： 首先项目请仔细阅读安装环境，提供了数据制作、训练、推理的代码和预训练权重 官方推荐Linux，而我使用的windows，可能因为其中含有detectron2的安装，这里我推荐一个简单方法，就是直接把detectron2/detectron2 这个目录PIP或者手动下载后，放在这个项目的Painter/ 下面，然后直接运行，如果出现了_six.py的torch版本报错，可以尝试如下解决方式： 再_six,py中加入： import collections.abc as container_abcs int_classes = int
另外从代码看出，这个开源的比较急，你可以下载作者提供的toy_datasets，包含了他们预训练任务的基本样例。
在安装基本环境后和模型权重下载后，你可以使用utils/或者eval下的painter_inference_demo.py脚本
修改其模型和路径参数，如下代码：
我们只需要修改模型的权重路径，并且提供prompt的图像，一对输入：标注图像和原始图，img_path为测试图
model_painter = prepare_model(ckpt_path, model)
print('Model loaded.')
device = torch.device("cuda")
model_painter.to(device)
img2_path = "data/toy_datasets/coco_pose/data_pair/train_256x192_aug0/000000006293_box7_image.png"
tgt2_path = "data/toy_datasets/coco_pose/data_pair/train_256x192_aug0/000000006293_box7_label.png"
img_path = "data/toy_datasets/coco_pose/data_pair/val_256x192/000000188465_box0_image.png"
比如我想进行关键点检测任务，首先提供Prompt的图像
对要推理的图像img推理后得到如下结果
1.4.1 模型结构以及训练设计（本节提到的详细代码可以看源码，这里DEIT和VIT比较成熟的原理不作赘述）
整个模型还是基于VIT系列的骨干，通过代码可以简要看出VIT的BACKBONE由两部分组成encoder、decoder, 编码器是由VIT的骨干BLOCK堆积而成，解码器其实是卷积层构成的。
这里的相对复杂的其实是encoder，如下图展示
这是基于掩码机制的图像训练流程，这个框架统一了以上7类视觉任务的输入和输出（这里跳过了对论文的7类任务设计的介绍）：
1.通过来自于同一个视觉任务的图像和其GT标注图像编队，经过常规的数据增强，模型处理的图像是默认（896，448）的比例。
2.随机的masking，论文中提到是75%的ratio，目的希望模型能够重建像素
3.在VIT骨干结构中，encoder是一个比较大的骨干VIT模型（24blocks），在encoder的过程中，经过patch_embeding后，可学习的token 会替代被mask的图像patch，将输入和其GT加入Posembeds，segment_token后拼接送进VIT中，具体核心的Mask token和分割像素分类的token处理代码大致如下：
self.mask_token = nn.Parameter(torch.zeros(1, 1, 1, embed_dim))
mask_token = self.mask_token.expand(batch_size, Hp, Wp, -1)
# replace the masked visual tokens by mask_token，w是一个掩码矩阵，这里masktoken是可学习的
w = bool_masked_pos.unsqueeze(-1).type_as(mask_token).reshape(-1, Hp, Wp, 1)
y = y * (1 - w) + mask_token * w
# add pos embed w/o cls token
x = x + self.segment_token_x
y = y + self.segment_token_y
if self.pos_embed is not None:
x = x + get_abs_pos(
self.pos_embed, self.pretrain_use_cls_token, (x.shape[1], x.shape[2])
)
y = y + get_abs_pos(
self.pos_embed, self.pretrain_use_cls_token, (y.shape[1], y.shape[2])
)
- decoder是一组卷积层做的解码部分，其中源码有穿插patch正反化的操作，个人认为解码器输出的肯定是为图像，大致操作如下，
(在平时训练过程中LOSS的计算迭代，使得模型可以“脑补”mask的像素，这是经典的MAE套路)
loss, y, mask = model(x.float().to(device), tgt.float().to(device), bool_masked_pos.to(device), valid.float().to(device))
// y代表输出图像,我注释掉了pathch
//y = model.unpatchify(y)
print("unpatchify :",y.size())
y = torch.einsum('nchw->nhwc', y).detach().cpu()
# 插值回原始尺度，MEAN,STD处理，作Image数据类型转换
output = y[0, y.shape[1]//2:, :, :]
output = torch.clip((output * imagenet_std + imagenet_mean) * 255, 0, 255)
output = F.interpolate(output[None, ...].permute(0, 3, 1, 2), size=[size[1], size[0]], mode='nearest').permute(0, 2, 3, 1)[0]
output = output.int()
output = Image.fromarray(output.numpy().astype(np.uint8))
output.save(out_path)
1.4.2 训练损失函数
这里一致强调输入和输出都是图像，那么LOSS自然也是像素相关的，文章中使用了smooth-L1 去计算掩码像素损失，通过实验比较L1,L2以及组合，发现smooth-L1效果最好。
2. 额外补充
论文中提到，
不同的任务提示会导致不同的结果。因此，如何选择或生成更合适的任务提示可以是一个探索的新方向。在这里，我们介绍
两个简单的基线，我们将更多的探索作为未来的工作。第一个基线是通过选择获得更好的提示，我们以启发式方式遍历整个训练集，并为每个任务选择性能最佳的示例对。第二个基线是生成任务提示。我们将任务提示定义为可学习的张量，冻结整个模型，然后使用训练损失来优化任务提示。
表格证明，学习prompt的方式或许能够带来更好的性能。
2.1 细节实现
作者团队为不同的任务做出相应处理：
2.1.1语义分割：
我们在 RGB 空间中使用不同的颜色来制定不同的语义类别。为此，我们定义背景并忽略黑色区域，即彩色像素（0, 0, 0），并使用伪代码阐述的伪代码为前景类别生成颜色。
def generate colors dict(b, K):
# b: the number of used values in a single channel
# K: the number of classes
m = 255 // b # get margin
colors = []
for class id in range(K):
# compute margin multiplier
r mult = class id // b∗∗2
g mult = (class id %
b mult = class id %
# compute r, g, b values
r = 255 − r mult ∗ m
g = 255 − g mult ∗ m
b = 255 − b mult ∗ m
colors.append((r, g, b))
return colors
在推理过程中，要将输出图像解码为单通道 ID 映射，其中每个像素代表一个类 ID，通过L1距离计算每个输出像素与每个语义类别的近似度，并将最接近颜色的 ID 作为预测类别。T
def forward(image, colors):
# image: (H, W, 3)
# colors: (K, 3), where K is the number of classes
# get distance between pixels and pre−defined colors
dist = (image.view(H, W, 1, 3) − colors.view(1, 1, K, 3)
).abs().sum(−1) # (H, W, K)
segm = dist.argmin(dim=−1) # (H, W)
return segm
2.1.2 关键点检测
对于关键点检测，输出图像由表示与类无关的热图的 R 通道和表示关键点类别的 G/B 通道组成。经过后处理得到关键点位置：
def forward(image, colors):
# image: (H, W, 3)
# colors: (K, 3), where K is the number of keypoints
# r for heatmaps and gb for keypoint classes
r = images[..., 0] # (H, W)
gb = images[..., 1:] # (H, W, 2)
# get keypoint class of each pixel
dist = (gb.view(H, W, 1, 2) − colors.view(1, 1, K, 2)).
abs().sum(−1) # (H, W, K)
segm = dist.argmin(dim=−1) # (H, W)
for idx in range(K):
mask = segm == idx
heatmap = mask ∗ r # (H, W)
heatmaps.append(heatmap)
heatmaps = stack(heatmaps) # (K, H, W)
return heatmaps
2.1.2 全景分割：
将全景分割任务分解为语义分割和与类无关的实例分割,所以你从代码来看，其实是跑了两个任务，然后将这两类分割任务结果合并：
在训练期间，语义分割子任务使用，与 ADE-20K 上的语义分割设置相同。除了在分配颜色时设置了基数 b =
7。类似语义分割的颜色生成过程，为与类无关的实例分割中使用的每个位置类别去生成颜色。每个实例mask的颜色由其中心的位置决定；
在推理期间，采用DMatrix NMS
以删除重复的实例预测。我们应用语义预测中的像素多数投票来获取每个实例掩码的语义类，最后，将语义分割和实例分割预测合并，得到全景分割结果。
全景分割的设计思路，论文中也提到了是参考的SOLO（算法将图像分为网格，如果某一对象的中心落在一个网格单元中，该网格单元负责预测语义类别和分割该实例）这样一个”yolo”化的设计思想，同时使用了Matrix-NMS去做后处理，直接用和预测的某个mask的IOU最大的其余mask预测来近似表达一个概率，然后求出一个影响因子Factor，去和置信度做乘积，也就是说取决于iou最高的mask。
对于这一任务的后处理代码中提供了Soft-nms的更多选择。推理部分作者使用的通用处理，只有在评估的代码中，作者加入了以上很对不同任务的一些处理，具体可以看Painter/eval/oco_panoptic/COCOInstSegEvaluatorCustom.py下的脚本说明。
这样一种根据输入的prompt提示去完成多个视觉任务的设计思路很经验，确实是趋向于CV任务一统化的设计，看了不少CV大模型的论文，这里VIT和DEIT的研究贡献真的是基石。
2.1.3 图像重建
这类low-level的任务，推理伪代码：
def forward(segm dist, inst masks):
# segm dist: (H, W, K), where K is the number of thing
classes
# inst masks: (N, H, W), where N is the number of
instances
# turn distances to scores
segm scores = 1. − semseg dist / max(semseg dist)
# majority vote
class probs = einsum("nhw,hwk−>nk", inst masks,
segm scores) # (N, K)
pred classes = class probs.argmax(dim=−1) # (N,)
return pred classes
这类任务代码给出的是常规的psnr和ssim指标来评价效果，直接给论文的数据体现：
2.2 实现的代码优化细节（必看）
在3K迭代下进行消融实验，证明如下工作是否Work.
2.2.1.patch merging
在训练中每个输入样本都是由输入图像和输出图像成对组成,这样内存的开销会增大，在每3个block之后逐个path添加它们的特征；
通过合并输入图像和输出图像的早期特征减少了将近一半的计算成本。如上表所示，Patch-merging设计通过将输入和输出堆叠在一起，进一步提供了输入和输出之间的像素到像素对应关系。但是在原始设置中，这些关系需要由模型学习，尤其是在短时间内这将使优化更加困难,代码如下：
目的：减少计算量&&增强输入和输出关联性
patch merging 逻辑：
merge_idx = 2
x = torch.cat((x, y), dim=0)
# apply Transformer blocks
out = []
for idx, blk in enumerate(self.blocks):
x = blk(x)
if idx == merge_idx: ##每三个vit的block后进行输入和输出的拆分
x = (x[:x.shape[0]//2] + x[x.shape[0]//2:]) * 0.5
if idx in [5, 11, 17, 23]: ##这里实际压入了4个tenor=相当于bLOCK模块均匀采样的4个特征图的串联和
out.append(self.norm(x))
return out
2.2.2 encoder
如上图表所示，使用 ViT-L 的模型比使用 ViT-B 的模型性能非常大。这种观察是直观的，也就是模型越大会产生更好的性能的可能性也越大。对于通用模型，可以使用更多的数据，但在方法设计之前使用较少的任务特定数据，因此可能需要更多的模型容量而不是特定于任务的模型。
2.2.3 Head
使用由线性（1×1卷积）、3×3卷积层和另一个线性层组成的轻量三层Head，代码如下：
self.decoder_embed = nn.Linear(embed_dim*4, patch_size ** 2 * self.decoder_embed_dim, bias=True)
self.decoder_pred = nn.Sequential(
nn.Conv2d(self.decoder_embed_dim, self.decoder_embed_dim, kernel_size=3, padding=1, ),
LayerNorm2D(self.decoder_embed_dim),
nn.GELU(),
nn.Conv2d(self.decoder_embed_dim, 3, kernel_size=1, bias=True),
)
将每个补丁的特征映射到其原始分辨率，例如，16×16×3。每个patch特征是从vit模块均匀采样的 4 个特征图的串联.
这里代码可以解释：decode中接受输入是一个list为4的featurepae，VIT模型就是encoder中的VIT的每4次Block后加起来再存一次，一共存了4个结果，2.2.1节提到过。
def forward_decoder(self, x):
# predictor projection
x = torch.cat(x, dim=-1)
x = self.decoder_embed(x)
p = self.patch_size
h, w = x.shape[1], x.shape[2]
x = x.reshape(shape=(x.shape[0], h, w, p, p, self.decoder_embed_dim))
x = torch.einsum('nhwpqc->nchpwq', x)
x = x.reshape(shape=(x.shape[0], -1, h * p, w * p))
x = self.decoder_pred(x) # Bx3xHxW
return x
总结
大概这篇只消化到这种程度了，这类项目因为其代码成熟度还不够，并且很多细节本人也没有完全看完，如果读者想了解更多可以留言，怕很多人觉得写的比较乱和晦涩就不继续补一些东西了，特别是主要原因在工业界这类任务的性能成熟度远远不够，不过值得学习的是作者团队如何运用DEIT且集合CV多个经典模型框架去设计这套“通用框架”的想法，其次是实现网络设计时候针对不同任务的解析兼容和模型的计算优化都使我开阔了不少思路，并且最后文章的work的编码细节也是值得学习的， 下篇由于SEGGPT没有开源核心代码，故暂时不作。
希望我5小时的奋笔疾书能换你的点赞关注~
文章出处登录后可见！
