Lớp SingleRoIExtractor, được định nghĩa trong tệp single_level_roi_extractor.py, có nhiệm vụ trích xuất đặc trưng từ các vùng quan tâm (RoI) từ một hoặc nhiều lớp đặc trưng. Nó kế thừa từ lớp cơ sở BaseRoIExtractor.
import torch
from mmcv.runner import force_fp32
from mmdet.models.builder import ROI_EXTRACTORS
from .base_roi_extractor import BaseRoIExtractor
@ROI_EXTRACTORS.register_module()
class RoIFeatureExtractor(BaseRoIExtractor):
"""Trích xuất đặc trưng RoI từ một lớp đặc trưng duy nhất.
Nếu có nhiều lớp đầu vào, mỗi RoI sẽ được ánh xạ đến một lớp phù hợp dựa trên
tỷ lệ của nó. Luật ánh xạ này được đề xuất trong
`FPN <https://arxiv.org/abs/1612.03144>`_.
Args:
roi_cfg (dict): Cấu hình lớp RoI.
output_channels (int): Số kênh đầu ra của lớp RoI.
feature_strides (list[int]): Các bước nhảy của các bản đồ đặc trưng đầu vào.
scale_threshold (int): Ngưỡng tỷ lệ để ánh xạ đến lớp 0. Mặc định: 56.
"""
def __init__(self,
roi_cfg,
output_channels,
feature_strides,
scale_threshold=56):
super(RoIFeatureExtractor, self).__init__(roi_cfg, output_channels,
feature_strides)
self.scale_threshold = scale_threshold
def assign_roi_levels(self, rois, num_levels):
"""Ánh xạ các RoI đến các lớp đặc trưng tương ứng dựa trên tỷ lệ.
- scale < scale_threshold * 2: level 0
- scale_threshold * 2 <= scale < scale_threshold * 4: level 1
- scale_threshold * 4 <= scale < scale_threshold * 8: level 2
- scale >= scale_threshold * 8: level 3
Args:
rois (Tensor): Đầu vào RoI, có hình dạng (k, 5).
num_levels (int): Tổng số lớp.
Returns:
Tensor: Chỉ số lớp (bắt đầu từ 0) của mỗi RoI, có hình dạng (k, )
"""
scale = torch.pow(
(rois[:, 3] - rois[:, 1]) * (rois[:, 4] - rois[:, 2]), 0.5)
target_lvls = torch.floor(torch.log2(scale / self.scale_threshold + 1e-6))
target_lvls = target_lvls.clamp(min=0, max=num_levels - 1).long()
return target_lvls
@force_fp32(apply_to=('features', ), out_fp16=True)
def extract_features(self, features, rois, roi_scale_factor=None):
"""Hàm lan truyền tiến."""
out_size = self.roi_operators[0].output_size
num_levels = len(features)
extracted_features = features[0].new_zeros(
rois.size(0), self.output_channels, *out_size)
# TODO: loại bỏ khi parrots hỗ trợ
if torch.__version__ == 'parrots':
extracted_features.requires_grad = True
if num_levels == 1:
if len(rois) == 0:
return extracted_features
return self.roi_operators[0](features[0], rois)
target_lvls = self.assign_roi_levels(rois, num_levels)
if roi_scale_factor is not None:
rois = self.scale_roi_coordinates(rois, roi_scale_factor)
for i in range(num_levels):
indices = target_lvls == i
if indices.any():
rois_ = rois[indices, :]
level_features = self.roi_operators[i](features[i], rois_)
extracted_features[indices] = level_features
else:
extracted_features += sum(
x.view(-1)[0]
for x in self.parameters()) * 0. + features[i].sum() * 0.
return extracted_features
Lớp RoIFeatureExtractor có các hàm chính sau:
(1) __init__(): Hàm khởi tạo, thiết lập giá trị scale_threshold làm cơ sở để phân bổ mẫu vào lớp FPN nào;
(2)assign_roi_levels(): Phân bổ mẫu vào một lớp FPN cụ thể, tính toán dựa trên công thức trong bài báo;
(3)extract_features(): Lan truyền tiến, có các tham số features, rois, trong đó features là các lớp đặc trưng của FPN, rois là tọa độ của ROI.
Lớp RoIFeatureExtractor có ít hàm, nhiều chức năng được thực hiện bằng cách gọi các phương thức từ lớp BaseRoIExtractor. Lớp BaseRoIExtractor được định nghĩa trong tệp base_roi_extractor.py:
from abc import ABCMeta, abstractmethod
import torch
import torch.nn as nn
from mmcv import ops
class AbstractRoIExtractor(nn.Module, metaclass=ABCMeta):
"""Lớp cơ sở cho bộ trích xuất RoI.
Args:
roi_config (dict): Cấu hình lớp RoI.
output_dim (int): Số kênh đầu ra của lớp RoI.
strides (list[int]): Các bước nhảy của bản đồ đặc trưng đầu vào.
"""
def __init__(self, roi_config, output_dim, strides):
super(AbstractRoIExtractor, self).__init__()
self.roi_operators = self.construct_roi_operators(roi_config, strides)
self.output_dim = output_dim
self.strides = strides
self.fp16_enabled = False
@property
def num_inputs(self):
"""int: Số lượng bản đồ đặc trưng đầu vào."""
return len(self.strides)
def init_weights(self):
pass
def construct_roi_operators(self, layer_cfg, strides):
"""Xây dựng toán tử RoI để trích xuất đặc trưng từ mỗi lớp bản đồ đặc trưng.
Args:
layer_cfg (dict): Từ điển để xây dựng và cấu hình toán tử RoI.
Các tùy chọn là các mô-đun dưới ``mmcv/ops`` như ``RoIAlign``.
strides (list[int]): Bước nhảy của bản đồ đặc trưng đầu vào so với
kích thước ảnh gốc, sẽ được sử dụng để điều chỉnh tọa độ RoI
(hệ tọa độ ảnh gốc) sang hệ tọa độ đặc trưng.
Returns:
nn.ModuleList: Các mô-đun trích xuất RoI cho mỗi lớp bản đồ đặc trưng.
"""
cfg = layer_cfg.copy()
layer_type = cfg.pop('type')
assert hasattr(ops, layer_type)
layer_cls = getattr(ops, layer_type)
roi_operators = nn.ModuleList(
[layer_cls(spatial_scale=1 / s, **cfg) for s in strides])
return roi_operators
def scale_roi_coordinates(self, rois, scale_factor):
"""Thay đổi tỷ lệ tọa độ RoI.
Args:
rois (torch.Tensor): RoI (Vùng quan tâm), hình dạng (n, 5)
scale_factor (float): Tỷ lệ nhân với RoI.
Returns:
torch.Tensor: RoI đã được thay đổi tỷ lệ.
"""
center_x = (rois[:, 1] + rois[:, 3]) * 0.5
center_y = (rois[:, 2] + rois[:, 4]) * 0.5
width = rois[:, 3] - rois[:, 1]
height = rois[:, 4] - rois[:, 2]
new_width = width * scale_factor
new_height = height * scale_factor
new_x1 = center_x - new_width * 0.5
new_y1 = center_y - new_height * 0.5
new_x2 = center_x + new_width * 0.5
new_y2 = center_y + new_height * 0.5
new_rois = torch.stack((rois[:, 0], new_x1, new_y1, new_x2, new_y2), dim=-1)
return new_rois
@abstractmethod
def forward(self, feats, rois, roi_scale_factor=None):
pass
Lớp AbstractRoIExtractor có các hàm chính sau:
(1) __init__(): Hàm khởi tạo, có các tham số roi_config, output_dim, strides, xây dựng các lớp RoI thông qua hàm construct_roi_operators();
(2) construct_roi_operators(): Xây dựng các lớp RoI, có các tham số layer_cfg, strides, sử dụng nn.ModuleList() để tạo ra các lớp RoI;
(3) scale_roi_coordinates(): Thay đổi tỷ lệ RoI;
(4) forward(): Phương thức trừu tượng, được triển khai bởi lớp con.