Detr部署

使用代码及环境配置

代码

使用D:\学习文件\研究生文件\部署相关文件\DETR\代码\test1.ipynb文件

环境配置

1.创建环境

conda create -n detr python==3.x(3.8版本以上)

#安装 PyTorch 1.5+ 和 torchvision 0.6+

conda install -c pytorch pytorch torchvision

2.git文件

git clone https://github.com/facebookresearch/detr.git

3.requirementstxt文件

pip install requirements.txt -i Simple Index

代码流程

1.导包

from PIL import Image

import requests

import matplotlib.pyplot as plt

# %config InlineBackend.figure_format = 'retina'

import torch

from torch import nn

from torchvision.models import resnet50

import torchvision.transforms as T

torch.set_grad_enabled(False);

2.设置类DETRdemo

class DETRdemo(nn.Module):
    def __init__(self, num_classes, hidden_dim=256, nheads=8,
                 num_encoder_layers=6, num_decoder_layers=6):
        super().__init__()

        # 创建ResNet-50主干网络
        self.backbone = resnet50()
        del self.backbone.fc  # 删除ResNet-50中的全连接层

        # 创建转换层
        self.conv = nn.Conv2d(2048, hidden_dim, 1)  # 将2048个通道转换为hidden_dim（256）个通道

        # 创建一个默认的PyTorch transformer
        self.transformer = nn.Transformer(
            hidden_dim, nheads, num_encoder_layers, num_decoder_layers)

        # 预测头，一个额外的类用于预测非空槽
        # 注意，在基线DETR中，linear_bbox层是一个3层的MLP
        self.linear_class = nn.Linear(hidden_dim, num_classes + 1)  # 线性层用于类别预测
        self.linear_bbox = nn.Linear(hidden_dim, 4)  # 线性层用于边界框预测

        # 输出位置编码（对象查询）
        self.query_pos = nn.Parameter(torch.rand(100, hidden_dim))  # 对象查询位置编码

        # 空间位置编码
        # 注意，在基线DETR中我们使用正弦位置编码
        self.row_embed = nn.Parameter(torch.rand(50, hidden_dim // 2))  # 行位置编码
        self.col_embed = nn.Parameter(torch.rand(50, hidden_dim // 2))  # 列位置编码

    def forward(self, inputs):
        # 通过ResNet-50的前几层传播输入直到avg-pool层
        x = self.backbone.conv1(inputs)  # 通过ResNet的conv1层
        x = self.backbone.bn1(x)  # 通过ResNet的bn1层
        x = self.backbone.relu(x)  # 通过ResNet的relu层
        x = self.backbone.maxpool(x)  # 通过ResNet的maxpool层

        x = self.backbone.layer1(x)  # 通过ResNet的layer1层
        x = self.backbone.layer2(x)  # 通过ResNet的layer2层
        x = self.backbone.layer3(x)  # 通过ResNet的layer3层
        x = self.backbone.layer4(x)  # 通过ResNet的layer4层

        # 从2048个特征平面转换为transformer所需的256个特征平面
        h = self.conv(x)  # 通过转换层

        # 构建位置编码
        H, W = h.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)  # 合并并展平位置编码

        # 通过transformer传播
        h = self.transformer(pos + 0.1 * h.flatten(2).permute(2, 0, 1),
                             self.query_pos.unsqueeze(1)).transpose(0, 1)
        # 将特征图展平成二维，加入位置编码，通过transformer传播并转置回原始形状
        
        # 最后，将transformer输出投影到类别标签和边界框
        return {'pred_logits': self.linear_class(h), 
                'pred_boxes': self.linear_bbox(h).sigmoid()}  # 通过线性层得到类别和边界框，并应用sigmoid到边界框

3.加载与训练好的DETR模型权重

#这行代码创建了一个DETRdemo模型实例，num_classes参数设为91。
#这意味着该模型可以识别91种不同的类别（比如COCO数据集中的类别数量）。
detr = DETRdemo(num_classes=91)



#这行代码从指定的URL下载预训练好的模型权重。
#url参数指定权重文件的下载地址。
#map_location='cpu'参数表示将下载的权重映射到CPU设备上，这样即使没有GPU也可以加载模型。
#check_hash=True参数用于验证下载文件的哈希值，以确保文件完整性。
state_dict = torch.hub.load_state_dict_from_url(
    url='https://dl.fbaipublicfiles.com/detr/detr_demo-da2a99e9.pth',
    map_location='cpu', check_hash=True)


#将下载的权重加载到DETR模型实例中。
#state_dict是包含模型参数的字典，通过load_state_dict方法将这些参数加载到模型中。
detr.load_state_dict(state_dict)



#模型设置为评估模式。
#eval()方法将模型切换到评估模式，这对某些模块（如Dropout和BatchNorm）的行为有所影响，使其在推理时的表现与训练时不同。
detr.eval();

4.定义coco数据集类别名称及可视化颜色

# 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'
]
# 定义了一个包含COCO数据集中所有类别名称的列表，共91个类别，其中一些索引位置上是'N/A'，表示没有对应类别

# 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]]
# 定义了一组用于可视化的颜色，每个颜色是一个RGB值的列表

5.将pytorch图像输入归一化

# standard PyTorch mean-std input image normalization
transform = T.Compose([
    T.Resize(800),
    # 将图像的较小边缩放到800像素，同时保持长宽比

    T.ToTensor(),
    # 将图像从PIL格式或numpy数组转换为PyTorch张量，并且将图像的像素值从0-255缩放到0-1之间

    T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    # 使用ImageNet数据集上的均值和标准差对图像进行归一化
    # 均值为[0.485, 0.456, 0.406]，标准差为[0.229, 0.224, 0.225]
])

6.对DETR模型输出进行后处理

定义了一些函数，用于对DETR模型的输出进行后处理，包括将边界框格式从中心坐标格式转换为左上角-右下角格式，以及根据图像大小对边界框进行缩放。最后定义了一个detect函数，用于对输入图像进行检测并返回预测的概率和边界框。

# 用于输出边界框的后处理
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 detect(im, model, transform):
    # 对输入图像进行均值-标准差归一化（批量大小：1）
    img = transform(im).unsqueeze(0)
    # 使用定义的transform对输入图像进行归一化处理，并增加一个batch维度

    # 演示模型默认仅支持宽高比在0.5到2之间的图像
    # 如果你想使用宽高比在此范围之外的图像
    # 需要重新调整图像的大小，以确保最大尺寸不超过1333像素，效果最佳
    assert img.shape[-2] <= 1600 and img.shape[-1] <= 1600, '演示模型仅支持每边最多1600像素的图像'
    # 确保输入图像的长宽不超过1600像素

    # 通过模型进行传播
    outputs = model(img)
    # 将图像输入模型，得到输出

    # 仅保留置信度大于0.7的预测
    probas = outputs['pred_logits'].softmax(-1)[0, :, :-1]
    keep = probas.max(-1).values > 0.7
    # 仅保留置信度大于0.7的预测结果

    # 将边界框从[0; 1]的比例转换为图像的实际尺寸
    bboxes_scaled = rescale_bboxes(outputs['pred_boxes'][0, keep], im.size)
    # 将边界框从[0; 1]的比例转换为图像的实际尺寸

    return probas[keep], bboxes_scaled
    # 返回保留的预测概率和缩放后的边界框

7.导入图片及获取检测框和得分

url1 = '1.png'
im = Image.open(url1)

scores, boxes = detect(im, detr, transform)

8.画结果图

def plot_results(pil_img, prob, boxes):
    # 设置图形大小
    plt.figure(figsize=(16,10))
    # 显示图像
    plt.imshow(pil_img) 
    # 获取当前的图形坐标轴
    ax = plt.gca()
    # 对于每个预测概率、边界框和颜色的组合，绘制边界框
    for p, (xmin, ymin, xmax, ymax), c in zip(prob, boxes.tolist(), COLORS * 100):
        # 添加一个矩形框表示检测到的物体
        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()

8.结果示意图