从DETR到Mask2Former（1）：DETR-segmentation结构全解析

网上关于DETR做的detection的解析很多，但是DETR做Segmentation的几乎没有，本文结合DETR的论文与代码，对DETR做一个详细的拆解。理解DETR是理解Mask2Former的基础。

首先得把DETR-segmentation给run起来。Github上DETR的repository，下载了也只能run起来detection，run不起来segmentation功能，但还是下载下来，后面留着有用。我们用torch的hub里集成的DETR segmentation模型，运行下面的代码

import torch
models_list = torch.hub.list('facebookresearch/detr', force_reload=True)
print(models_list)

你可以看到torch.hub中所有关于detr的模型。我们选择  detr_resnet50_panoptic

model = torch.hub.load('facebookresearch/detr', 'detr_resnet50_panoptic', pretrained=True)

再新建一个py文件，把以下代码放进去：

import math
from PIL import Image
import requests
import matplotlib.pyplot as plt

import ipywidgets as widgets
from IPython.display import display, clear_output

import torch
from torch import nn
import torch.nn.functional as F
from torchvision.models import resnet50
import torchvision.transforms as T
torch.set_grad_enabled(False)

import matplotlib.pyplot as plt


# COCO classes
CLASSES = [
	'N/A', 'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus',
	'train', 'truck', 'boat', 'traffic light', 'fire hydrant', 'N/A',
	'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse',
	'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe', 'N/A', 'backpack',
	'umbrella', 'N/A', 'N/A', 'handbag', 'tie', 'suitcase', 'frisbee', 'skis',
	'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove',
	'skateboard', 'surfboard', 'tennis racket', 'bottle', 'N/A', 'wine glass',
	'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple', 'sandwich',
	'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake',
	'chair', 'couch', 'potted plant', 'bed', 'N/A', 'dining table', 'N/A',
	'N/A', 'toilet', 'N/A', 'tv', 'laptop', 'mouse', 'remote', 'keyboard',
	'cell phone', 'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'N/A',
	'book', 'clock', 'vase', 'scissors', 'teddy bear', 'hair drier',
	'toothbrush'
]

# colors for visualization
COLORS = [[0.000, 0.447, 0.741], [0.850, 0.325, 0.098], [0.929, 0.694, 0.125],
		  [0.494, 0.184, 0.556], [0.466, 0.674, 0.188], [0.301, 0.745, 0.933]]

transform = T.Compose([
	T.Resize(800),
	T.ToTensor(),
	T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

# for output bounding box post-processing
def box_cxcywh_to_xyxy(x):
	x_c, y_c, w, h = x.unbind(1)
	b = [(x_c - 0.5 * w), (y_c - 0.5 * h),
		 (x_c + 0.5 * w), (y_c + 0.5 * h)]
	return torch.stack(b, dim=1)

def rescale_bboxes(out_bbox, size):
	img_w, img_h = size
	b = box_cxcywh_to_xyxy(out_bbox)
	b = b * torch.tensor([img_w, img_h, img_w, img_h], dtype=torch.float32)
	return b

def plot_results(pil_img, prob, boxes):
	plt.figure(figsize=(16,10))
	plt.imshow(pil_img)
	ax = plt.gca()
	colors = COLORS * 100
	for p, (xmin, ymin, xmax, ymax), c in zip(prob, boxes.tolist(), colors):
		ax.add_patch(plt.Rectangle((xmin, ymin), xmax - xmin, ymax - ymin,
								   fill=False, color=c, linewidth=3))
		# cl = p.argmax()
		# text = f'{CLASSES[cl]}: {p[cl]:0.2f}'
		# ax.text(xmin, ymin, text, fontsize=15,
		#         bbox=dict(facecolor='yellow', alpha=0.5))
	plt.axis('off')
	plt.show()


model = torch.hub.load('facebookresearch/detr', 'detr_resnet50_panoptic', pretrained=True)
model.eval()

url = 'http://images.cocodataset.org/val2017/000000039769.jpg'
im = Image.open(requests.get(url, stream=True).raw)

# mean-std normalize the input image (batch-size: 1)
img = transform(im).unsqueeze(0)

# propagate through the model
outputs = model(img)

# show result of detection
# keep only predictions with 0.7+ confidence
probas = outputs['pred_logits'].softmax(-1)[0, :, :-1]
keep = probas.max(-1).values > 0.9
# convert boxes from [0; 1] to image scales
bboxes_scaled = rescale_bboxes(outputs['pred_boxes'][0, keep], im.size)
plot_results(im, probas[keep], bboxes_scaled)


# show segmentation
# compute the scores, excluding the "no-object" class (the last one)
scores = outputs['pred_logits']
scores = scores.softmax(dim=-1)
scores = scores[..., :-1]
scores_onehot = scores.max(-1)
scores_onehot = scores_onehot[0]
# threshold the confidence
keep = scores_onehot > 0.85

scores_selected = scores[keep]
labels = torch.argmax(scores_selected, dim=-1)

masks = outputs['pred_masks'][keep].detach().cpu().numpy()
mask_i = masks[0, :, :]

# plt.imshow(mask_i, cmap="viridis")
# plt.show()

# ## Plot all the remaining masks
ncols = 5
fig, axs = plt.subplots(ncols=ncols, nrows=math.ceil(keep.sum().item() / ncols), figsize=(18, 10))
for line in axs:
	for a in line:
		a.axis('off')
for i, mask in enumerate(outputs['pred_masks'][keep].detach().cpu().numpy()):
	ax = axs[i // ncols, i % ncols]
	ax.imshow(mask, cmap="cividis")
	ax.text(0, 0, labels[i].cpu().numpy().item(0))
	ax.axis('off')
fig.tight_layout()

plt.show()

mask_pred = outputs['pred_masks'].sigmoid()
mask_pred = F.interpolate(mask_pred, size=(480, 640), mode='bilinear', align_corners=False)

semseg = torch.einsum("bqc,bqhw->bchw", scores, mask_pred)
result = torch.argmax(semseg, dim=1)


plt.figure(figsize=(12, 8))
# 第一个子图
plt.subplot(1, 1, 1)
# image1_rgb = cv2.cvtColor(img_color, cv2.COLOR_BGR2RGB)
plt.imshow(result[0, :, :].to("cpu").numpy())
plt.title('Image 1')
plt.axis('off')

plt.show()

解释一下上面的代码，

plot_results(im, probas[keep], bboxes_scaled)

是把detection的结果可视化出来，本文略过的detection部分的内容。  plot_results()  下面的代码就是segmentation内容结果的可视化。模型的输出包括三个内容

pred_logits(1, 100, 251)中三个数字的含义是：1是batchsize，100是query的个数，251是分类数，去掉最后一个no-object不要，实际是250个类。可以这样理解，每一个query都会进入网络，但不是每一个query都能从图片找到东西，而找到东西的query，找到的instance所属于的类就是这个query所在的那一行中，最大的数所在列的index。如果一个图片里有两只猫，那结果就是有两个query，分别各自找到一只猫（也就是一个instance），这样也就实现了全景分割的功能。所以说DETR的结构是Mask2Former实现全景分割功能的基础。

pred_boxes是query所找到的instance对应的box。

pred_masks是query所找到的instance所在的像素。

这段代码：

# ## Plot all the remaining masks
ncols = 5
fig, axs = plt.subplots(ncols=ncols, nrows=math.ceil(keep.sum().item() / ncols), figsize=(18, 10))
for line in axs:
	for a in line:
		a.axis('off')
for i, mask in enumerate(outputs['pred_masks'][keep].detach().cpu().numpy()):
	ax = axs[i // ncols, i % ncols]
	ax.imshow(mask, cmap="cividis")
	ax.text(0, 0, labels[i].cpu().numpy().item(0))
	ax.axis('off')
fig.tight_layout()

plt.show()

会画出这样的图像：

有94个query没有找到instance，有6个query找到了instance，2个query找到了猫，也就是250列的17列，两个query找到了遥控器（250列的74列）。

好了，现在DETR-seg的输出我们弄清楚了，接下来进到DETR内部去看看，这个模型封装在了torch.hub中，进入内部的正确方法，说实话笔者也不知道，这里笔者耍了一个不正经的小trick：

import torch

model = torch.hub.load('facebookresearch/detr', 'detr_resnet50_panoptic', pretrained=True)

from models.segmentation import DETRsegm

model = DETRsegm(model)

上面的代码中，models.segmentation文件是DETR的github repository下载下来，里面有的内容，运行上面的代码，会报错

但是通过这个报错，我们可以找到torch.hub中DETR的源代码。这样就可以在源代码里打断点，看DETR 内部了。进入源代码，我们能看到这样的内容，注释的代码是笔者加进去的，不是torch自带的。

self.detr.backbone是resnet18。 这里，我们打开DETR的github代码——detr/models/backbone.py 找到  class BackboneBase 这个类。

这里的 def forward 中的out，就是 DETRsegm 中的 features

features, pos = self.detr.backbone(samples)

pos是正弦函数的 position embeding，就是输入transformer的encoder的位置编码。

在获取features时，DETR detection和Segmentation的区别在与，detection只拿resnet第四层的输出，而segmentation将每层的输出都拿出来。通过这段代码实现（缩进乱了，忽略）。

if return_interm_layers:
            return_layers = {"layer1": "0", "layer2": "1", "layer3": "2", "layer4": "3"}
        else:
            return_layers = {'layer4': "0"}
        self.body = IntermediateLayerGetter(backbone, return_layers=return_layers)

这里扯远一点，复习一下resnet18  resnet18在torch中也有集成，通过下面的代码得到，

import torch
import torchvision.models as models
from torchvision.models._utils import IntermediateLayerGetter
import cv2
import numpy as np


# 加载一个预训练的 ResNet 模型
model = models.resnet18(pretrained=True)
model.to("cuda")

# 定义要获取的中间层
layers = {'layer1': 'layer1', 'layer2': 'layer2', 'layer3': 'layer3', 'layer4': 'layer4'}

# 创建 IntermediateLayerGetter
intermediate_layers = IntermediateLayerGetter(model, layers)

image0 = cv2.imread("/home/robotics/dino/img/kitaku/002.jpg")
height0, width0, channels = image0.shape
image = cv2.resize(image0, (640, 480), interpolation=cv2.INTER_AREA)
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

mean = np.array([0.485, 0.456, 0.406]) * 255
std = np.array([0.229, 0.224, 0.225]) * 255
image = image.astype(float)
for i in range(3):
	image[:, :, i] = (image[:, :, i] - mean[i]) / std[i]

input_data = torch.from_numpy(image).permute(2, 0, 1).unsqueeze(dim=0).to(torch.float32).to("cuda")

# 获取中间层的输出
outputs = intermediate_layers(input_data)

# 打印输出
for key, value in outputs.items():
	print(f"{key}: {value.shape}")


from torchview import draw_graph
model_graph = draw_graph(model, input_size=(1, 3, 480, 640))
model_graph.resize_graph(scale=5.0)
model_graph.visual_graph.render(format='svg')

通过draw_graph得到一个网络结构可视化图，格式是svg，如下图所示

也可以自己写代码，构建一个resnet18

import torch
import torch.nn as nn
from torch.nn import functional as F


class RestNetBasicBlock(nn.Module):
	def __init__(self, in_channels, out_channels, stride):
		super(RestNetBasicBlock, self).__init__()
		self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1)
		self.bn1 = nn.BatchNorm2d(out_channels)
		self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=stride, padding=1)
		self.bn2 = nn.BatchNorm2d(out_channels)

	def forward(self, x):
		output = self.conv1(x)
		output = F.relu(self.bn1(output))
		output = self.conv2(output)
		output = self.bn2(output)
		return F.relu(x + output)


class RestNetDownBlock(nn.Module):
	def __init__(self, in_channels, out_channels, stride):
		super(RestNetDownBlock, self).__init__()
		self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride[0], padding=1)
		self.bn1 = nn.BatchNorm2d(out_channels)
		self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=stride[1], padding=1)
		self.bn2 = nn.BatchNorm2d(out_channels)

	def forward(self, x):
		output = self.conv1(x)
		out = F.relu(self.bn1(output))

		out = self.conv2(out)
		out = self.bn2(out)
		return F.relu(x + out)


class RestNet18(nn.Module):
	def __init__(self):
		super(RestNet18, self).__init__()
		self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3)
		self.bn1 = nn.BatchNorm2d(64)
		self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)

		self.layer1 = nn.Sequential(RestNetBasicBlock(64, 64, 1),
									RestNetBasicBlock(64, 64, 1))

		self.layer2 = nn.Sequential(RestNetDownBlock(64, 128, [2, 1]),
									RestNetBasicBlock(128, 128, 1))

		self.layer3 = nn.Sequential(RestNetDownBlock(128, 256, [2, 1]),
									RestNetBasicBlock(256, 256, 1))

		self.layer4 = nn.Sequential(RestNetDownBlock(256, 512, [2, 1]),
									RestNetBasicBlock(512, 512, 1))

		self.avgpool = nn.AdaptiveAvgPool2d(output_size=(1, 1))

		self.fc = nn.Linear(512, 10)

	def forward(self, x):
		out = self.conv1(x)
		out = self.layer1(out)
		out = self.layer2(out)
		out = self.layer3(out)
		out = self.layer4(out)
		out = self.avgpool(out)
		out = out.reshape(x.shape[0], -1)
		out = self.fc(out)
		return out


if __name__ == "__main__":
	resnet18 = RestNet18()
	
	from torchview import draw_graph
	model_graph = draw_graph(resnet18, input_size=(1, 3, 480, 640))
	model_graph.resize_graph(scale=5.0)
	model_graph.visual_graph.render(format='svg')

好，回到DETR，还是看torch.hub中的这段detr代码：

 上面的代码中：

hs, memory = self.detr.transformer(src_proj, mask, self.detr.query_embed.weight, pos[-1])

self.detr.transformer就是下图经典的transformer结构，memory是encoder的输出  hs是decoder的输出。mask对于理解detr不重要，是masked attention中用到了。

pos和query_embed.weight分别对应上图的Positional Encoding和decoder中，output上面的OutputEmbeding。

DETR中transformer的理解，可以参看下面的代码，这段代码是笔者自己写的，对于DETR本身是没有用处的，只是为了方便理解，有错误的地方（feature那里只拿出resnet18最后一层的输出，是错的，不影响大局，懒得改了），自己理解时候做一个参考就好。

import torch
from torch import nn
from torchvision.models import resnet50
from models.transformer import TransformerEncoderLayer, TransformerEncoder, TransformerDecoderLayer, TransformerDecoder
from models.segmentation import MHAttentionMap, MaskHeadSmallConv
import time


class DETR(nn.Module):
    def __init__(self, num_classes, d_model, nheads, dim_feedforward=2048,
                 num_encoder_layers=6, num_decoder_layers=6, dropout=0.9, activation='relu', normalize_before=True):
        super().__init__()
        # We take only convolutional layers from ResNet-50 model
        self.backbone = nn.Sequential(*list(resnet50(pretrained=True).children())[:-2])
        self.conv = nn.Conv2d(dim_feedforward, d_model, 1)
        
        encoder_layer = TransformerEncoderLayer(d_model, nheads, dim_feedforward,
                                                dropout, activation, normalize_before)
        encoder_norm = nn.LayerNorm(d_model) if normalize_before else None
        self.encoder = TransformerEncoder(encoder_layer, num_encoder_layers, encoder_norm)

        decoder_layer = TransformerDecoderLayer(d_model, nheads, dim_feedforward,
                                                dropout, activation, normalize_before)
        decoder_norm = nn.LayerNorm(d_model)
        self.decoder = TransformerDecoder(decoder_layer, num_decoder_layers, decoder_norm,
                                          return_intermediate=True)

        self.linear_class = nn.Linear(d_model, num_classes + 1)
        self.linear_bbox = nn.Linear(d_model, 4)
        
        num_queries = 100
        self.query_embed = nn.Embedding(num_queries, d_model)
        
        #position embedding
        self.row_embed = nn.Parameter(torch.rand(50, d_model // 2))
        self.col_embed = nn.Parameter(torch.rand(50, d_model // 2))
        
        self.bbox_attention = MHAttentionMap(d_model, d_model, nheads, dropout=0.0)
        self.mask_head = MaskHeadSmallConv(d_model + nheads, [1024, 512, 256], d_model)
 
    def forward(self, inputs):
        #inputs是[1,3,800,1200]
        features = self.backbone(inputs)
        #x是[1,2048,25,38]
        hh = self.conv(features)
        #hh是[1,256,25,38]
        H, W = hh.shape[-2:]
        pos = torch.cat([
            self.col_embed[:W].unsqueeze(0).repeat(H, 1, 1),
            self.row_embed[:H].unsqueeze(1).repeat(1, W, 1),
        ], dim=-1).flatten(0, 1).unsqueeze(1)
        
        bs, c, h, w = hh.shape
        
        query_embed = self.query_embed.weight.unsqueeze(1).repeat(1, bs, 1)
        tgt = torch.zeros_like(query_embed)
        src = hh.flatten(2).permute(2, 0, 1)
        memory = self.encoder(src=src, pos=pos)
        hs = self.decoder(tgt, memory, pos=pos, query_pos=query_embed)
        
        hs = hs.transpose(1, 2)
        memory = memory.permute(1, 2, 0).view(bs, c, h, w)
        bbox_mask = self.bbox_attention(hs[-1], memory)
        
        seg_masks = self.mask_head(hh, bbox_mask, [features[0], features[0], features[0]])
        outputs_seg_masks = seg_masks.view(bs, self.detr.num_queries, seg_masks.shape[-2], seg_masks.shape[-1])
        
        #self.query_pos是[100,256]
        #src是encoder输入，tgt是decoder输入
        

        #h是[100,1,256]
        return outputs_seg_masks
 
 
#coco是91个类, hidden dimension是256, 多头注意力是8, encoder,decoder layer都是6
device = torch.device("cuda")
detr = DETR(num_classes=91, d_model=256, nheads=8, num_encoder_layers=6, num_decoder_layers=6)
detr.eval().cuda()


inputs = torch.randn(1, 3, 800, 1200).cuda()

outputs_seg_masks = detr(inputs)
# print(logits, bboxes)
#logits是[100,1,92]
#bboxes是[100,1,4]

还是回到torch.hub那个detr代码，下面结合DETR论文中的那张图来理解一下这段代码

图中multi head attention部分，对应代码中

bbox_mask = self.bbox_attention(hs[-1], memory, mask=mask)

hs[-1]是左侧四个不同颜色的小框，memory是Encoded image，mask可以忽略，bbox_mask的尺寸是：（batch_size，query的个数，注意力头的个数，attention map的高和宽）

        temp = bbox_mask[0][20].squeeze().cpu().numpy()
        import matplotlib.pyplot as plt
        for i in range(8):
            plt.imshow(temp[i, :, :], cmap="viridis")
            plt.show()

上面那段代码就是拿出第1个batchsize的第21个query的所有注意力头的attention map，也就是可以画出这几张图：

seg_masks = self.mask_head(src_proj, bbox_mask, [features[2].tensors, features[1].tensors, features[0].tensors])

mask_head对应图中这个结构：

输入是多头注意力map，以及resnet18这个backbone的4层的特征。

本文先到这，下一篇写DETR的损失函数。

未完待续