ROI Align原理介绍

介绍

torch.ops.torchvision.roi_align是PyTorch的一个操作函数，用于在给定的RoIs（感兴趣区域）上执行空间变换网络。该操作函数可以被用于目标检测和物体识别等计算机视觉任务中，通常在RoI池化之前被执行。

RoI Align（Region of Interest Align）是RoI池化的一种改进，它可以更加精确地对齐RoI中的特征，避免了RoI池化过程中的信息损失。具体来说，RoI Align操作首先将RoI区域划分为若干个小的网格，然后在每个网格内执行双线性插值，最终得到固定大小的特征图。

torch.ops.torchvision.roi_align函数的输入包括特征图、RoIs、输出特征图的大小和插值的空间大小等参数。它的输出是对应于每个RoI的特征图，每个RoI的特征图的大小相同。

示例1

以下是一个使用torch.ops.torchvision.roi_align函数的简单示例：

import torch
import torchvision.ops as ops

# 假设特征图大小为[1, 256, 56, 56]
features = torch.rand(1, 256, 56, 56).float()

# 假设RoIs为[batch_index, x1, y1, x2, y2]
# 其中batch_index表示RoI所在的图像在batch中的索引
rois = torch.tensor([[0, 10, 10, 50, 50], [0, 20, 20, 40, 40]]).float()

# 定义RoI Align的参数
output_size = (7, 7)
spatial_scale = 1.0 / 16

# 执行RoI Align操作
output = ops.roi_align(features, rois, output_size, spatial_scale)
print(output.shape) # torch.Size([2, 256, 7, 7])

在上述示例中，features是输入的特征图，rois是输入的RoIs，output_size是输出特征图的大小，spatial_scale是特征图和原始图像之间的空间比例因子。最终，output是RoI Align操作的输出特征图。

示例2

自定义实现方法：

import torch
import torch.nn.functional as F


def roi_align(x, rois, output_size, spatial_scale=1.0, sampling_ratio=-1):
    # 将感兴趣区域（ROI）坐标转换为特征图上的坐标
    rois = rois / spatial_scale
    # 向上取整，保证感兴趣区域完全覆盖
    rois[:, 2:] = torch.ceil(rois[:, 2:])
    # 向下取整，保证感兴趣区域完全覆盖
    rois[:, :2] = torch.floor(rois[:, :2])
    # 将坐标转换为整数类型
    rois = rois.int()
    roi_features = []
    for roi in rois:
        # 获取感兴趣区域内的特征图
        x_ = x[..., roi[0]:roi[2], roi[1]:roi[3]]
        h, w = x_.shape[-2:]
        if sampling_ratio > 0:
            # 如果采样率大于0，则进行双线性插值
            x_ = F.interpolate(x_, size=(sampling_ratio * h, sampling_ratio * w), mode='bilinear', align_corners=False)
        # 进行最大池化操作，得到感兴趣区域内的特征向量
        roi_feature = F.adaptive_max_pool2d(x_, output_size)
        roi_features.append(roi_feature)
    # 将所有感兴趣区域内的特征向量拼接起来，得到最终的输出结果
    return torch.cat(roi_features, dim=0)


# 测试代码
feature_map = torch.randn(1, 3, 64, 64)
rois = torch.tensor([[0, 10, 10, 50, 50], [0, 20, 20, 40, 40]])
output_size = (7, 7)
output = roi_align(feature_map, rois, output_size, sampling_ratio=5)
print(output.shape) # torch.Size([2, 3, 7, 7])

补充：ROI Pooling

ROI Pooling是一种池化方法，用于从输入特征映射中提取感兴趣的区域信息。在目标检测任务中，ROI Pooling通常用于将不同大小的候选框（proposals）映射到相同大小的特征图上，以便进行后续的特征提取和分类。

具体来说，ROI Pooling首先将输入特征图与每个提议的边界框进行匹配，然后将匹配的区域划分为大小相同的网格，并对每个网格进行最大池化操作，以得到固定大小的特征向量。这样，不同大小的提议都可以映射到相同大小的特征图上，以便进行后续的特征提取和分类操作。

ROI Pooling与传统的池化方法（如最大池化、平均池化等）不同，因为它只对感兴趣的区域进行池化，而非在整个特征图上进行池化。这种选择性池化方法有助于更好地捕捉目标对象的局部特征，提高目标检测的准确性和鲁棒性。

以下是使用Python和PyTorch实现ROI Pooling的示例代码：

import torch  
import torch.nn.functional as F  
  
def roi_pooling(inputs, rois, output_size):  
    # inputs: [batch_size, channels, height, width]  
    # rois: [num_proposals, 4] (x1, y1, x2, y2)  
    # output_size: [height, width]  
  
    # 将提议的边界框转换为相对于输入特征图的坐标  
    rois = rois.t()  
    batch_ids = rois[:, 0].long()  
    roi_start_w = rois[:, 1] * inputs.size(3)  
    roi_start_h = rois[:, 2] * inputs.size(2)  
    roi_end_w = rois[:, 3] * inputs.size(3)  
    roi_end_h = rois[:, 4] * inputs.size(2)  
  
    # 对每个输入特征图进行ROI Pooling  
    pooled = []  
    for i in range(inputs.size(0)):  
        # 提取与当前输入特征图相关的提议  
        roi_batch_ids = batch_ids == i  
        if roi_batch_ids.sum() == 0:  
            # 如果没有与当前输入特征图相关的提议，则使用零填充  
            pooled.append(torch.zeros(0, inputs.size(1), output_size[0], output_size[1]).float().to(inputs.device))  
            continue  
        rois_i = rois[roi_batch_ids]  
  
        # 将提议的边界框转换为相对于输入特征图的坐标  
        roi_start_w_i = roi_start_w[roi_batch_ids]  
        roi_start_h_i = roi_start_h[roi_batch_ids]  
        roi_end_w_i = roi_end_w[roi_batch_ids]  
        roi_end_h_i = roi_end_h[roi_batch_ids]  
  
        # 将输入特征图分割成等大小的网格，并选择落在提议区域内的网格进行最大池化  
        grid_size = inputs.size(2) // output_size[0]  
        input_rois = torch.stack([roi_start_w_i / grid_size, roi_start_h_i / grid_size, roi_end_w_i / grid_size, roi_end_h_i / grid_size], dim=1).round().type(torch.int)  
        input_rois = torch.clamp(input_rois, 0, torch.max(input_rois[:, 2:]) - 1)  
        input_rois[:, 0] = torch.clamp(input_rois[:, 0], 0, inputs.size(3) - 1)  
        input_rois[:, 1] = torch.clamp(input_rois[:, 1], 0, inputs.size(2) - 1)  
        input = inputs[i][:, input_rois[:, 0], input_rois[:, 1]]  
        input = F.max_pool2d(input, kernel_size=(grid_size, grid_size))  
        pooled.append(input)  
    pooled = torch.stack(pooled, dim=0)  
    return pooled

这段代码接受两个输入参数：inputs和rois，其中inputs是输入特征图，rois是候选框的坐标信息。该函数通过将候选框转换为相对于输入特征图的坐标，然后将输入特征图分割成等大小的网格，并选择落在候选框内的网格进行最大池化，从而实现了ROI Pooling操作。