Fast RCNN中的ROI pooling代码实现

简单版的roi pooling的代码实现，方便理解。
实际上就是原图经过cnn以后得到feature map，同时对原图进行selective search得到若干个候选区域（坐标为归一化的），在feature map中获取候选区域的对应feature，再将该feature resize为roi pooling的输出尺寸即可。

import cv2
import numpy as np


def roi_pooling(feature_map, rois, output_size):
    '''
    根据指定的 RoI 坐标，在特征图上进行池化
    
    参数：
    feature_map: 特征图
    rois: RoI 坐标
    output_size: 池化输出大小
    
    返回：
    pooled_features: 池化后的 RoI 特征
    '''
    pooled_features = []
    height, width, channels = feature_map.shape
    
    # 遍历每一个 RoI
    for roi in rois:
        # 将归一化的坐标映射回特征图上的位置
        start_x = int(np.floor(roi[0] * width))
        start_y = int(np.floor(roi[1] * height))
        end_x = int(np.ceil(roi[2] * width))
        end_y = int(np.ceil(roi[3] * height))
        
        # 获取在特征图上对应的 RoI 区域
        roi_feature = feature_map[start_y:end_y, start_x:end_x, :]
        
        # 对 RoI 区域进行池化，并将其调整为指定的输出大小
        pooled_roi = np.zeros(output_size + (channels,))
        for c in range(channels):
            channel_roi = roi_feature[:, :, c]
            pooled_roi[:, :, c] = cv2.resize(channel_roi, output_size, interpolation=cv2.INTER_AREA)
        
        pooled_features.append(pooled_roi)
    
    return np.array(pooled_features)

def main():
    # 示例输入数据
    # 假设 feature_map 是一个 32x32x8 的特征图，rois 是一组有selective search得到的候选框归一化后的坐标信息
    feature_map = np.random.rand(32, 32, 8)
    rois = [[0.2, 0.3, 0.6, 0.8], [0.4, 0.5, 0.8, 0.9]]  # 两个示例的候选框区域坐标信息（在特征图上的比例）

    # 指定输出的固定大小
    output_size = (4, 4)

    # 调用 RoI Pooling 函数
    pooled_features = roi_pooling(feature_map, rois, output_size)
    print(pooled_features.shape)  # 打印池化后的特征的形状 shape:(2, 4, 4, 8)

if __name__ == '__main__':
    main()