基于注意力机制的目标检测代码

基于注意力机制的目标检测算法有很多，这里提供一种基于注意力机制的 Faster R-CNN 的代码实现。Faster R-CNN 是一种两阶段目标检测算法，其基本思路是通过区域提取网络 (Region Proposal Network, RPN) 提取出可能包含目标的区域，再通过一个分类和回归网络对这些区域进行分类和定位。

以下是一个基于 Pytorch 的 Faster R-CNN 代码实现，其中使用了注意力机制来改善分类和回归的准确性：

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

class Attention(nn.Module):
    def __init__(self, in_channels):
        super(Attention, self).__init__()
        self.conv1 = nn.Conv2d(in_channels, in_channels // 8, kernel_size=1)
        self.conv2 = nn.Conv2d(in_channels // 8, in_channels, kernel_size=1)
        self.pool = nn.AdaptiveAvgPool2d(1)
        self.fc1 = nn.Linear(in_channels, in_channels // 8)
        self.fc2 = nn.Linear(in_channels // 8, in_channels)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        y = self.conv1(x)
        y = F.relu(y)
        y = self.conv2(y)
        y = self.sigmoid(y)
        y = x * y
        y = self.pool(y)
        y = y.view(y.size(0), -1)
        y = self.fc1(y)
        y = F.relu(y)
        y = self.fc2(y)
        y = self.sigmoid(y)
        y = y.view(y.size(0), y.size(1), 1, 1)
        return x * y

class FasterRCNN(nn.Module):
    def __init__(self, num_classes):
        super(FasterRCNN, self).__init__()
        self.backbone = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Conv2d(256, 512, kernel_size=3, stride=1, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1),
            nn.ReLU(inplace=True),
        )
        self.rpn = nn.Sequential(
            nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(512, 512, kernel_size=1, stride=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(512, 18, kernel_size=1, stride=1),
        )
        self.roi_pool = nn.AdaptiveMaxPool2d((7, 7))
        self.fc1 = nn.Linear(7 * 7 * 512, 4096)
        self.fc2 = nn.Linear(4096, 4096)
        self.cls_score = nn.Linear(4096, num_classes)
        self.bbox_pred = nn.Linear(4096, num_classes * 4)
        self.attention = Attention(512)

    def forward(self, x):
        x = self.backbone(x)
        attention = self.attention(x)
        rpn = self.rpn(attention)
        rpn_probs = F.softmax(rpn[:, :9, :, :], dim=1)
        rpn_deltas = rpn[:, 9:, :, :]
        rois = self.generate_rois(rpn_probs, rpn_deltas)
        roi_features = self.roi_pool(attention, rois)
        roi_features = roi_features.view(roi_features.size(0), -1)
        fc1 = F.relu(self.fc1(roi_features))
        fc2 = F.relu(self.fc2(fc1))
        cls_score = self.cls_score(fc2)
        bbox_pred = self.bbox_pred(fc2)
        return cls_score, bbox_pred

    def generate_rois(self, rpn_probs, rpn_deltas):
        _, _, H, W = rpn_probs.shape
        A = self.anchor_box(H, W)
        A = A.to(rpn_probs.device)
        rpn_probs = rpn_probs.detach().cpu().numpy()
        rpn_deltas = rpn_deltas.detach().cpu().numpy()
        rois = []
        for i in range(H):
            for j in range(W):
                for k in range(9):
                    if rpn_probs[0, k, i, j] > 0.5:
                        dy, dx, dh, dw = rpn_deltas[0, :, i, j]
                        center_y = i * 16 + A[k, 0] * 16
                        center_x = j * 16 + A[k, 1] * 16
                        h = A[k, 2] * np.exp(dh)
                        w = A[k, 3] * np.exp(dw)
                        y1 = center_y - h / 2
                        x1 = center_x - w / 2
                        y2 = center_y + h / 2
                        x2 = center_x + w / 2
                        rois.append([0, y1, x1, y2, x2])
        rois = torch.tensor(rois).float()
        return rois

    def anchor_box(self, H, W):
        scales = [8, 16, 32]
        ratios = [0.5, 1, 2]
        base_anchor = np.array([1, 1, 16, 16]) - 1
        anchors = self.generate_anchors(base_anchor, scales, ratios)
        shift_y = np.arange(0, H) * 16
        shift_x = np.arange(0, W) * 16
        shift_x, shift_y = np.meshgrid(shift_x, shift_y)
        shift = np.stack([shift_y.ravel(), shift_x.ravel(),
                          shift_y.ravel(), shift_x.ravel()], axis=1)
        A = anchors.shape[0]
        K = shift.shape[0]
        all_anchors = (anchors.reshape((1, A, 4)) +
                       shift.reshape((1, K, 4)).transpose((1, 0, 2)))
        all_anchors = all_anchors.reshape((A * K, 4))
        return all_anchors

    def generate_anchors(self, base_anchor, scales, ratios):
        h, w, cy, cx = self.anchor_geometry(base_anchor)
        scales_grid, ratios_grid = np.meshgrid(scales, ratios)
        scales_grid = scales_grid.reshape(-1)
        ratios_grid = ratios_grid.reshape(-1)
        h_grid = h.reshape((1, -1)) * np.sqrt(ratios_grid).reshape((-1, 1)) * scales_grid.reshape((-1, 1))
        w_grid = w.reshape((1, -1)) / np.sqrt(ratios_grid).reshape((-1, 1)) * scales_grid.reshape((-1, 1))
        cy_grid = np.repeat(cy.reshape((1, -1)), len(scales) * len(ratios), axis=0).reshape(-1)
        cx_grid = np.repeat(cx.reshape((1, -1)), len(scales) * len(ratios), axis=0).reshape(-1)
        anchors = np.stack([cy_grid, cx_grid, h_grid, w_grid], axis=1)
        return anchors

    def anchor_geometry(self, base_anchor):
        h, w, cy, cx = base_anchor
        return (np.array([h, w, h, w]), np.array([w, h, w, h]), np.array([cy, cy, cy, cy]), np.array([cx, cx, cx, cx]))

上面的代码中，`Attention` 类定义了注意力模块，`FasterRCNN` 类定义了整个 Faster R-CNN 网络。在 `forward` 方法中，先将输入图像经过卷积网络得到特征图，然后使用注意力模块对特征图进行加权，然后将加权后的特征图输入到区域提取网络（`rpn`）中，得到可能包含目标的候选框（`rois`），再将这些候选框输入到分类和回归网络中，得到最终的分类和定位结果。在 `generate_rois` 方法中，使用 RPN 网络输出的概率和边框偏移量生成候选框。