目录
- 一、BiCornerPool 类
- 二、CornerHead 类
- 1. 构造函数
- 2. _make_layers 方法
- 3. _init_corner_kpt_layers 方法
- 4. _init_corner_emb_layers 方法
- 5. _init_layers 方法
- 6. init_weights 方法
- 7. forward 方法
- 8. forward_single 方法
- 9. get_targets 方法
- 10. loss_by_feat 方法
- 11. loss_by_feat_single 方法
- 12. predict_by_feat 方法
- 13. _predict_by_feat_single 方法
- 14. _bboxes_nms 方法
- 15. _decode_heatmap 方法
《MMDetection 中的 CornerHead 解析》
在目标检测领域,MMDetection 是一个强大的开源工具包。其中,CornerHead
是一个具有独特设计的目标检测头,它源自 CornerNet 并进行了一些扩展和优化。本文将深入解析CornerHead
中的各个关键函数,帮助读者更好地理解其工作原理和实现细节。
一、BiCornerPool 类
1. 构造函数
- 功能:创建一个双向角落池化模块,用于提取图像中的特定方向的特征。
- 参数解释:
in_channels
:输入模块的通道数。directions
:一个包含两个字符串的列表,分别指定两个角落池化的方向,例如['top', 'left']
。feat_channels
:中间特征的通道数。out_channels
:输出模块的通道数。norm_cfg
:用于构建和配置归一化层的字典。init_cfg
:用于控制初始化的配置,默认为None
。
2. forward 方法
- 功能:对输入的特征图进行双向角落池化操作,并返回处理后的特征图。
- 执行过程:
- 首先,对输入特征图分别进行两个方向的卷积操作,得到
direction1_conv
和direction2_conv
。 - 然后,使用指定的方向进行角落池化操作,得到
direction1_feat
和direction2_feat
。 - 将两个方向的池化结果相加后进行卷积操作,得到
aftpool_conv
。 - 对输入特征图进行一次卷积操作得到
conv1
。 - 将
aftpool_conv
和conv1
相加后经过 ReLU 激活函数得到relu
。 - 最后,对
relu
进行一次卷积操作得到最终的输出特征图conv2
。
- 首先,对输入特征图分别进行两个方向的卷积操作,得到
二、CornerHead 类
1. 构造函数
- 功能:初始化 CornerHead,用于检测目标的角点。
- 参数解释:
num_classes
:除背景类别外的类别数量。in_channels
:输入特征图的通道数。num_feat_levels
:特征层级的数量,默认为 2(对于 HourglassNet-104)或 1(对于 HourglassNet-52)。corner_emb_channels
:角点嵌入向量的通道数,默认为 1。train_cfg
:训练配置,在 CornerHead 中暂未使用。test_cfg
:测试配置。loss_heatmap
:角点热图损失的配置。loss_embedding
:角点嵌入损失的配置。loss_offset
:角点偏移损失的配置。init_cfg
:初始化配置,为防止异常初始化行为,不允许设置。
2. _make_layers 方法
- 功能:初始化一个卷积序列,用于 CornerHead 的不同分支。
- 参数解释:
out_channels
:输出通道数。in_channels
:输入通道数,默认为 256。feat_channels
:中间特征的通道数,默认为 256。
3. _init_corner_kpt_layers 方法
- 功能:初始化角点关键点(corner keypoint)相关的层,包括角点热图分支和角点偏移分支。
- 执行过程:
- 创建两个模块列表
tl_pool
和br_pool
,分别用于存储 top-left 和 bottom-right 方向的双向角落池化模块。 - 创建两个模块列表
tl_heat
和br_heat
,分别用于存储 top-left 和 bottom-right 方向的角点热图生成模块。 - 创建两个模块列表
tl_off
和br_off
,分别用于存储 top-left 和 bottom-right 方向的角点偏移生成模块。
- 创建两个模块列表
4. _init_corner_emb_layers 方法
- 功能:初始化角点嵌入(corner embedding)相关的层,仅包括 top-left 和 bottom-right 方向的模块。
- 执行过程:
- 创建两个模块列表
tl_emb
和br_emb
,分别用于存储 top-left 和 bottom-right 方向的角点嵌入生成模块。
- 创建两个模块列表
5. _init_layers 方法
- 功能:初始化 CornerHead 的所有层,包括角点关键点层和角点嵌入层(如果需要)。
6. init_weights 方法
- 功能:初始化 CornerHead 的权重。
- 执行过程:
- 对不同层级的角点热图和偏移模块的最后一层卷积进行初始化,并设置偏置项。
- 如果有角点嵌入层,也对其最后一层卷积进行初始化。
7. forward 方法
- 功能:对输入的特征图进行前向传播,得到角点热图、偏移热图和嵌入热图。
- 参数解释:
feats
:来自上游网络的特征图元组,每个特征图都是一个 4D 张量。
- 执行过程:
- 使用
multi_apply
函数对每个特征层级进行前向传播,得到多个输出结果。 - 返回的结果包括 top-left 和 bottom-right 方向的角点热图、嵌入热图(如果有)和偏移热图。
- 使用
8. forward_single 方法
- 功能:对单个特征层级进行前向传播,得到 CornerHead 在该层级的输出。
- 参数解释:
x
:单个特征层级的特征图。lvl_ind
:特征层级的索引。return_pool
:是否返回角落池化特征,默认为False
。
- 执行过程:
- 依次进行 top-left 和 bottom-right 方向的角落池化、热图生成、嵌入生成(如果有)和偏移生成操作。
- 根据
return_pool
参数决定是否返回角落池化特征。
9. get_targets 方法
以下是对这段代码的详细解析:
9.1 函数功能
这个函数的目的是生成角点目标,包括角点热图、角点偏移以及可选的角点嵌入、引导偏移和向心偏移。它根据输入的真实边界框、标签、特征图形状和图像形状等信息来生成这些目标。
9.2 参数解释
gt_bboxes
:一个包含每个图像的真实边界框的列表,每个边界框是一个形状为(num_gt, 4)
的张量,表示左上角和右下角的坐标。gt_labels
:一个包含每个边界框真实标签的列表,每个标签是一个形状为(num_gt,)
的张量。feat_shape
:输出特征的形状,是一个包含[batch, channel, height, width]
的序列。img_shape
:输入图像的形状,是一个包含[height, width, channel]
的序列。with_corner_emb
:一个布尔值,表示是否生成角点嵌入目标,默认为False
。with_guiding_shift
:一个布尔值,表示是否生成引导偏移目标,默认为False
。with_centripetal_shift
:一个布尔值,表示是否生成向心偏移目标,默认为False
。
9.3 函数执行过程
-
提取特征图和图像的形状信息:
batch_size, _, height, width = feat_shape
:从特征图形状中提取批量大小、通道数、高度和宽度。img_h, img_w = img_shape[:2]
:从图像形状中提取高度和宽度。- 计算宽度和高度的比例:
width_ratio = float(width / img_w)
和height_ratio = float(height / img_h)
。
-
初始化目标张量:
- 创建用于存储左上角和右下角角点热图、偏移的张量:
gt_tl_heatmap = gt_bboxes[-1].new_zeros([batch_size, self.num_classes, height, width])
:左上角角点热图。gt_br_heatmap = gt_bboxes[-1].new_zeros([batch_size, self.num_classes, height, width])
:右下角角点热图。gt_tl_offset = gt_bboxes[-1].new_zeros([batch_size, 2, height, width])
:左上角角点偏移。gt_br_offset = gt_bboxes[-1].new_zeros([batch_size, 2, height, width])
:右下角角点偏移。
- 如果
with_corner_emb
为True
,初始化一个空列表用于存储角点嵌入。 - 如果
with_guiding_shift
为True
,创建用于存储左上角和右下角引导偏移的张量。 - 如果
with_centripetal_shift
为True
,创建用于存储左上角和右下角向心偏移的张量。
举例解析:
1). 首先理解变量的含义:
batch_size
表示这批图像的数量。例如,假设batch_size = 4
,意味着有 4 张图像在同一批次中处理。self.num_classes
表示要检测的目标类别数量。比如有 3 个类别:人、车、自行车,那么self.num_classes = 3
。height
和width
是特征图的高度和宽度。假设特征图大小为32x32
,即height = 32
,width = 32
。
2). 解析代码:
-
gt_tl_heatmap = gt_bboxes[-1].new_zeros([batch_size, self.num_classes, height, width])
:- 这里创建了一个四维张量来表示左上角角点热图。每个元素初始值为 0。这个张量的形状为
[batch_size, num_classes, height, width]
。 - 例如,对于上述假设的参数值,这个张量的形状就是
[4, 3, 32, 32]
。这意味着对于每一张图像(共 4 张),对于每一个类别(共 3 个类别),都有一个32x32
的热图来表示左上角角点的可能性分布。
- 这里创建了一个四维张量来表示左上角角点热图。每个元素初始值为 0。这个张量的形状为
-
gt_br_heatmap = gt_bboxes[-1].new_zeros([batch_size, self.num_classes, height, width])
:- 类似地,这创建了一个四维张量表示右下角角点热图,形状也为
[batch_size, num_classes, height, width]
。初始值全为 0。
- 类似地,这创建了一个四维张量表示右下角角点热图,形状也为
-
gt_tl_offset = gt_bboxes[-1].new_zeros([batch_size, 2, height, width])
:- 创建一个四维张量表示左上角角点的偏移量。形状为
[batch_size, 2, height, width]
。这里的“2”表示有两个维度的偏移量,通常是在水平和垂直方向上的偏移。例如,对于一张图像上的某个位置,这个张量存储了该位置左上角角点相对于真实左上角角点的水平和垂直偏移量。如果该位置不是左上角角点,那么偏移量可能不是零。对于每一张图像,都有一个这样的2x32x32
的张量来存储所有位置的左上角角点偏移量信息。
- 创建一个四维张量表示左上角角点的偏移量。形状为
-
gt_br_offset = gt_bboxes[-1].new_zeros([batch_size, 2, height, width])
:- 与
gt_tl_offset
类似,这个张量表示右下角角点的偏移量,形状也是[batch_size, 2, height, width]
。
- 与
-
if with_corner_emb:
- 如果
with_corner_emb
为真,这里会初始化一个名为match
的列表,但目前代码中没有进一步对这个列表进行操作,可能在后续的代码中会被填充或使用。
- 如果
-
if with_guiding_shift:
- 当
with_guiding_shift
为真时:- 创建一个四维张量
gt_tl_guiding_shift
,形状为[batch_size, 2, height, width]
。这个张量用于存储左上角角点的引导偏移(guiding shift)。初始值全为 0,数据类型与gt_bboxes[-1]
相同,并且在相同的设备上。 - 同样,创建另一个四维张量
gt_br_guiding_shift
用于存储右下角角点的引导偏移,形状也为[batch_size, 2, height, width]
。
- 创建一个四维张量
- 当
-
if with_centripetal_shift:
- 当
with_centripetal_shift
为真时:- 创建一个四维张量
gt_tl_centripetal_shift
,形状为[batch_size, 2, height, width]
。这个张量可能用于存储左上角角点的向心偏移(centripetal shift)。初始值全为 0,数据类型与gt_bboxes[-1]
相同,并且在相同的设备上。 - 创建另一个四维张量
gt_br_centripetal_shift
用于存储右下角角点的向心偏移,形状也为[batch_size, 2, height, width]
。
- 创建一个四维张量
- 当
- 创建用于存储左上角和右下角角点热图、偏移的张量:
-
遍历批次和边界框:
- 对于每个批次中的每个边界框:
- 提取边界框的坐标和中心坐标,以及对应的标签。
- 根据宽度和高度比例计算在特征图上的坐标。
- 将坐标转换为整数,并确保不超出特征图的范围。
- 生成高斯热图:
- 计算边界框在特征图上的宽度和高度。
- 根据宽度和高度计算高斯半径。
- 使用
gen_gaussian_target
函数生成左上角和右下角的高斯热图。见:〖MMDetection〗解析自定义文件:mmdet/models/utils/gaussian_target.py
- 生成角点偏移:
-
计算边界框四个角点在特征图上的偏移量。
-
将偏移量存储在对应的张量中。
(1)代码功能
这段代码的作用是计算并设置左上角和右下角角点的偏移量,并将这些偏移量存储在对应的张量中。它通常在目标检测任务中用于生成角点的偏移目标,以帮助模型更准确地定位目标的边界框。
参数解释
left_offset
、top_offset
、right_offset
、bottom_offset
:分别表示左上角的水平偏移量、左上角的垂直偏移量、右下角的水平偏移量和右下角的垂直偏移量。既取整与没有取整坐标之间的差值。scale_left
、scale_top
、scale_right
、scale_bottom
:是经过某种比例缩放后的边界框的坐标值。left_idx
、top_idx
、right_idx
、bottom_idx
:是整数坐标,表示在特征图或真实值张量上的位置索引。gt_tl_offset
和gt_br_offset
:分别是左上角和右下角角点偏移张量,形状为[batch_id, 2, height, width]
,其中batch_id
表示批次索引,2
分别对应水平和垂直两个方向的偏移,height
和width
是特征图的高度和宽度。
执行过程
-
计算偏移量:
left_offset = scale_left - left_idx
:计算左上角角点的水平偏移量,即缩放后的左侧边界坐标减去对应的整数索引。top_offset = scale_top - top_idx
:计算左上角角点的垂直偏移量,即缩放后的顶部边界坐标减去对应的整数索引。right_offset = scale_right - right_idx
:计算右下角角点的水平偏移量,即缩放后的右侧边界坐标减去对应的整数索引。bottom_offset = scale_bottom - bottom_idx
:计算右下角角点的垂直偏移量,即缩放后的底部边界坐标减去对应的整数索引。
-
设置偏移张量的值:
gt_tl_offset[batch_id, 0, top_idx, left_idx] = left_offset
:将左上角角点的水平偏移量存储在对应的张量位置中。其中batch_id
指定批次,0
表示水平方向的偏移,top_idx
和left_idx
指定在特征图上的位置索引。gt_tl_offset[batch_id, 1, top_idx, left_idx] = top_offset
:将左上角角点的垂直偏移量存储在对应的张量位置中。其中1
表示垂直方向的偏移。gt_br_offset[batch_id, 0, bottom_idx, right_idx] = right_offset
:将右下角角点的水平偏移量存储在对应的张量位置中。gt_br_offset[batch_id, 1, bottom_idx, right_idx] = bottom_offset
:将右下角角点的垂直偏移量存储在对应的张量位置中。
通过以上步骤,这段代码实现了计算并存储角点偏移量的功能,为目标检测模型提供了重要的训练目标之一。
-
- 如果
with_corner_emb
为True
,生成角点嵌入:- 将左上角和右下角的坐标添加到角点嵌入列表中。
- 如果
with_guiding_shift
为True
,生成引导偏移:- 计算从中心到角点的引导偏移,并存储在对应的张量中。
- 如果
with_centripetal_shift
为True
,生成向心偏移:- 计算从中心到角点的向心偏移,并存储在对应的张量中。
笔记: 这里首先将bbox转换到特征图上,然后计算左上角和右下角坐标对应的高斯热力图。 这里是可以可视化查看的。
- 对于每个批次中的每个边界框:
-
构建目标结果字典:
假设我们有一个目标检测的场景,有一批图像正在处理,其中batch_id = 0
,图像的尺寸为height = 8
和width = 8
。并且我们有一些已知的参数设置,比如with_corner_emb = True
、with_guiding_shift = True
、with_centripetal_shift = True
,以及一些特定的坐标和缩放值。
-
生成角点嵌入(Generate corner embedding)
如果
with_corner_emb
为真,就会执行以下操作:corner_match.append([[top_idx, left_idx], [bottom_idx, right_idx]])
假设
top_idx = 2
,left_idx = 3
,bottom_idx = 5
,right_idx = 4
,那么这行代码就会将[[2, 3], [5, 4]]
添加到corner_match
列表中。这个操作可能是为了记录图像中不同角点的匹配关系,以便后续处理。 -
生成引导偏移(Generate guiding shift)
如果
with_guiding_shift
为真,会执行以下操作来计算并设置引导偏移:1). 对于左上角角点:
gt_tl_guiding_shift[batch_id, 0, top_idx, left_idx] = scale_center_x - left_idx
gt_tl_guiding_shift[batch_id, 1, top_idx, left_idx] = scale_center_y - top_idx
假设
scale_center_x = 4.5
,scale_center_y = 3.5
,那么对于左上角角点,假设其索引为top_idx = 2
,left_idx = 3
,引导偏移在水平方向上的值为4.5 - 3 = 1.5
,在垂直方向上的值为3.5 - 2 = 1.5
。这个引导偏移表示从中心坐标到左上角角点在水平和垂直方向上的偏移量。没有取整的中心点坐标与取整的左上角坐标的差
2). 对于右下角角点:
gt_br_guiding_shift[batch_id, 0, bottom_idx, right_idx] = right_idx - scale_center_x
gt_br_guiding_shift[batch_id, 1, bottom_idx, right_idx] = bottom_idx - scale_center_y
假设
bottom_idx = 5
,right_idx = 4
,那么对于右下角角点,引导偏移在水平方向上的值为4 - 4.5 = -0.5
,在垂直方向上的值为5 - 3.5 = 1.5
。同样,这个引导偏移表示从中心坐标到右下角角点在水平和垂直方向上的偏移量。没有取整的中心点坐标与取整的右下角坐标的差
-
生成向心偏移(Generate centripetal shift)
如果
with_centripetal_shift
为真,会执行以下操作来计算并设置向心偏移:1). 对于左上角角点:
gt_tl_centripetal_shift[batch_id, 0, top_idx, left_idx] = log(scale_center_x - scale_left)
gt_tl_centripetal_shift[batch_id, 1, top_idx, left_idx] = log(scale_center_y - scale_top)
假设
scale_center_x = 4.5
,scale_left = 2
,scale_center_y = 3.5
,scale_top = 1
,那么对于左上角角点,向心偏移在水平方向上的值为log(4.5 - 2) = log(2.5)
,在垂直方向上的值为log(3.5 - 1) = log(2.5)
。向心偏移通过对数进行归一化,并且表示从中心坐标到左上角角点在水平和垂直方向上的一种特殊偏移量,可能用于强调目标的中心向角点的吸引力。没有取整的中心点坐标与没有取整的左上角坐标的差
2). 对于右下角角点:
gt_br_centripetal_shift[batch_id, 0, bottom_idx, right_idx] = log(scale_right - scale_center_x)
gt_br_centripetal_shift[batch_id, 1, bottom_idx, right_idx] = log(scale_bottom - scale_center_y)
假设
scale_right = 6
,scale_bottom = 7
,那么对于右下角角点,向心偏移在水平方向上的值为log(6 - 4.5) = log(1.5)
,在垂直方向上的值为log(7 - 3.5) = log(3.5)
。同样,这个向心偏移表示从中心坐标到右下角角点在水平和垂直方向上的特殊偏移量。没有取整的中心点坐标与没有取整的右下角坐标的差
- 返回目标结果字典:
整体功能:
这段代码的作用是将生成的各种目标检测相关的张量和数据整理到一个字典中,以便后续使用和返回。它根据不同的配置选项,选择性地将角点嵌入、引导偏移和向心偏移等信息添加到结果字典中。
具体解析:
-
target_result = dict(...)
:- 创建一个字典
target_result
,其中包含了左上角热图gt_tl_heatmap
、左上角偏移gt_tl_offset
、右下角热图gt_br_heatmap
和右下角偏移gt_br_offset
。这些是基本的目标检测结果,无论其他配置选项如何,都会包含在结果字典中。
- 创建一个字典
-
if with_corner_emb:
:- 如果
with_corner_emb
为真,表示需要生成角点嵌入。 target_result.update(corner_embedding=match)
:将角点嵌入列表match
添加到结果字典中,键为corner_embedding
。
- 如果
-
if with_guiding_shift:
:- 如果
with_guiding_shift
为真,表示需要生成引导偏移。 target_result.update(...)
:将左上角引导偏移张量gt_tl_guiding_shift
和右下角引导偏移张量gt_br_guiding_shift
添加到结果字典中,键分别为topleft_guiding_shift
和bottomright_guiding_shift
。
- 如果
-
if with_centripetal_shift:
:- 如果
with_centripetal_shift
为真,表示需要生成向心偏移。 target_result.update(...)
:将左上角向心偏移张量gt_tl_centripetal_shift
和右下角向心偏移张量gt_br_centripetal_shift
添加到结果字典中,键分别为topleft_centripetal_shift
和bottomright_centripetal_shift
。
- 如果
10. loss_by_feat 方法
10.1 函数功能
这个函数的目的是根据检测头提取的特征计算损失。它接受一系列输入,包括不同层级的角点热图、嵌入和偏移张量,以及真实实例信息和图像元信息,最后返回一个包含不同损失组件的字典。
10.2 参数解释
self
:指向当前对象的引用。tl_heats
:一个包含张量的列表,每个张量表示不同层级的左上角角点热图,形状为(N, num_classes, H, W)
,其中N
是批量大小,num_classes
是类别数,H
和W
是特征图的高度和宽度。br_heats
:与tl_heats
类似,但是表示右下角角点热图。tl_embs
:左上角角点嵌入张量列表,形状为(N, corner_emb_channels, H, W)
。br_embs
:右下角角点嵌入张量列表。tl_offs
:左上角角点偏移张量列表,形状为(N, corner_offset_channels, H, W)
。br_offs
:右下角角点偏移张量列表。batch_gt_instances
:一个包含真实实例数据的列表,通常包括边界框和标签属性。batch_img_metas
:图像的元信息列表,例如图像大小、缩放因子等。batch_gt_instances_ignore
(可选):指定在计算损失时可以忽略的边界框列表。
10.3 函数执行过程
-
提取真实边界框和标签:
gt_bboxes = [gt_instances.bboxes for gt_instances in batch_gt_instances]
:从batch_gt_instances
中提取每个图像的真实边界框列表。gt_labels = [gt_instances.labels for gt_instances in batch_gt_instances]
:提取每个图像的真实标签列表。
-
生成目标:
targets = self.get_targets(gt_bboxes, gt_labels, tl_heats[-1].shape, batch_img_metas[0]['batch_input_shape'], with_corner_emb=self.with_corner_emb)
:调用get_targets
函数生成角点目标,包括角点热图、偏移和可选的角点嵌入等。这个函数根据真实边界框、标签、特征图形状和图像元信息来生成目标。mlvl_targets = [targets for _ in range(self.num_feat_levels)]
:为每个特征层级复制目标,因为不同层级可能需要相同的目标信息。
-
计算损失:
det_losses, pull_losses, push_losses, off_losses = multi_apply(self.loss_by_feat_single, tl_heats, br_heats, tl_embs, br_embs, tl_offs, br_offs, mlvl_targets)
:使用multi_apply
函数在每个特征层级上调用loss_by_feat_single
函数来计算损失。这个函数返回不同类型的损失,包括角点关键点损失(det_losses
)、关联嵌入损失的第一部分(pull_losses
)、关联嵌入损失的第二部分(push_losses
)和角点偏移损失(off_losses
)。
-
构建损失字典:
loss_dict = dict(det_loss=det_losses, off_loss=off_losses)
:首先创建一个字典,包含角点关键点损失和角点偏移损失。- 如果有角点嵌入,则更新字典:
if self.with_corner_emb: loss_dict.update(pull_loss=pull_losses, push_loss=push_losses)
。
-
返回损失字典:
- 返回包含不同损失组件的字典,可用于后续的优化和训练过程。
11. loss_by_feat_single 方法
11.1 函数功能
这个函数用于计算在单个尺度级别上目标检测的损失。它接收由检测头提取的不同特征张量以及目标字典,然后分别计算角点检测损失、关联嵌入损失(如果启用了角点嵌入)和角点偏移损失,并将这些损失组合后返回。
11.2 参数解释
tl_hmp
:当前尺度级别的左上角角点热图张量,形状为(N, num_classes, H, W)
,其中N
是批量大小,num_classes
是类别数,H
和W
是热图的高度和宽度。br_hmp
:当前尺度级别的右下角角点热图张量,形状与tl_hmp
相同。tl_emb
:当前尺度级别的左上角角点嵌入张量(可选),形状为(N, corner_emb_channels, H, W)
。br_emb
:当前尺度级别的右下角角点嵌入张量(可选),形状与tl_emb
相同。tl_off
:当前尺度级别的左上角角点偏移张量,形状为(N, corner_offset_channels, H, W)
。br_off
:当前尺度级别的右下角角点偏移张量,形状与tl_off
相同。targets
:由get_targets
函数生成的角点目标字典,包含了真实的左上角热图'topleft_heatmap'
、右下角热图'bottomright_heatmap'
、左上角偏移'topleft_offset'
、右下角偏移'bottomright_offset'
以及角点嵌入'corner_embedding'
(如果启用了角点嵌入)。
11.3 函数执行过程
-
提取目标张量:
gt_tl_hmp = targets['topleft_heatmap']
:从目标字典中提取真实的左上角角点热图。gt_br_hmp = targets['bottomright_heatmap']
:提取真实的右下角角点热图。gt_tl_off = targets['topleft_offset']
:提取真实的左上角角点偏移。gt_br_off = targets['bottomright_offset']
:提取真实的右下角角点偏移。gt_embedding = targets['corner_embedding']
:如果启用了角点嵌入,提取角点嵌入目标。
-
计算检测损失:
tl_det_loss = self.loss_heatmap(tl_hmp.sigmoid(), gt_tl_hmp, avg_factor=max(1, gt_tl_hmp.eq(1).sum()))
:计算左上角角点的检测损失。首先对预测的热图进行 sigmoid 激活,然后使用自定义的loss_heatmap
函数计算损失。avg_factor
用于归一化损失,这里根据真实热图中值为 1 的元素数量进行计算,确保在样本数量较少时也能合理计算损失。br_det_loss = self.loss_heatmap(br_hmp.sigmoid(), gt_br_hmp, avg_factor=max(1, gt_br_hmp.eq(1).sum()))
:类似地,计算右下角角点的检测损失。det_loss = (tl_det_loss + br_det_loss) / 2.0
:将两个角点的检测损失平均,得到最终的检测损失。
-
计算关联嵌入损失(如果启用):
if self.with_corner_emb and self.loss_embedding is not None:
:如果启用了角点嵌入并且存在关联嵌入损失函数。pull_loss, push_loss = self.loss_embedding(tl_emb, br_emb, gt_embedding)
:调用自定义的loss_embedding
函数计算关联嵌入损失,该损失通常包括拉近正样本对(pull loss)和推开负样本对(push loss)两部分。
else:
:如果未启用角点嵌入或者没有关联嵌入损失函数。pull_loss, push_loss = None, None
:将关联嵌入损失设置为 None。
-
计算偏移损失:
- 首先计算偏移损失的掩码:
tl_off_mask = gt_tl_hmp.eq(1).sum(1).gt(0).unsqueeze(1).type_as(gt_tl_hmp)
:对于左上角角点,在真实热图中找到值为 1 的元素(表示真实角点位置),然后对每个样本在类别维度上求和,判断是否大于 0,得到一个布尔掩码,再将其扩展为与热图相同的数据类型。br_off_mask = gt_br_hmp.eq(1).sum(1).gt(0).unsqueeze(1).type_as(gt_br_hmp)
:类似地,计算右下角角点的偏移损失掩码。
- 然后计算偏移损失:
tl_off_loss = self.loss_offset(tl_off, gt_tl_off, tl_off_mask, avg_factor=max(1, tl_off_mask.sum()))
:使用自定义的loss_offset
函数计算左上角角点的偏移损失,考虑了掩码以仅在真实角点位置计算损失,并根据掩码中值为 1 的元素数量进行归一化。br_off_loss = self.loss_offset(br_off, gt_br_off, br_off_mask, avg_factor=max(1, br_off_mask.sum()))
:计算右下角角点的偏移损失。
off_loss = (tl_off_loss + br_off_loss) / 2.0
:将两个角点的偏移损失平均,得到最终的偏移损失。
模拟数据查看
br_off_mask
:# 1. gaussian2D() radius = 1 x = torch.arange( -radius, radius + 1).view(1, -1) print(x) y = torch.arange( -radius, radius + 1).view(-1, 1) print(y) diameter = 2 * radius + 1 sigma = diameter / 6 h = (-(x * x + y * y) / (2 * sigma * sigma)).exp() print(h) h[h < torch.finfo(h.dtype).eps * h.max()] = 0 print(h) # 2. gen_gaussian_target() # 初始化数据 gaussian_kernel = h center = [2, 2] heatmap = torch.tensor(torch.rand((5, 5))) # 这里手动获取后两个唯独:宽度和高度。 因为之前做了[left_idx, top_idx] print("heatmap:",heatmap) # 开始 x, y = center height, width = heatmap.shape[:2] print("height, width:",height, width) print("x:",x) print("y",y) print("radius:",radius) left, right = min(x, radius), min(width - x, radius + 1) top, bottom = min(y, radius), min(height - y, radius + 1) print("left, right, top, bottom:",left, right, top, bottom) masked_heatmap = heatmap[y - top:y + bottom, x - left:x + right] print("masked_heatmap:",masked_heatmap) masked_gaussian = gaussian_kernel[radius - top:radius + bottom, radius - left:radius + right] print("masked_gaussian:",masked_gaussian) k=1 # 默认为1 out_heatmap = heatmap print("out_heatmap:",out_heatmap) print("masked_gaussian * k:",masked_gaussian * k) torch.max( masked_heatmap, masked_gaussian * k, out=out_heatmap[y - top:y + bottom, x - left:x + right]) print(out_heatmap) # 假设out_heatmap为左上角热图 gt_tl_hmp = out_heatmap tl_off_mask = gt_tl_hmp.eq(1) print("tl_off_mask:",tl_off_mask) tl_off_mask = gt_tl_hmp.eq(1).sum(1) print("tl_off_mask:",tl_off_mask) tl_off_mask = gt_tl_hmp.eq(1).sum(1).gt(0) print("tl_off_mask:",tl_off_mask) tl_off_mask = gt_tl_hmp.eq(1).sum(1).gt(0).unsqueeze(1) print("tl_off_mask:",tl_off_mask) tl_off_mask = gt_tl_hmp.eq(1).sum(1).gt(0).unsqueeze(1).type_as( gt_tl_hmp) print("tl_off_mask:",tl_off_mask) tensor([[-1, 0, 1]]) tensor([[-1], [ 0], [ 1]]) tensor([[0.0183, 0.1353, 0.0183], [0.1353, 1.0000, 0.1353], [0.0183, 0.1353, 0.0183]]) tensor([[0.0183, 0.1353, 0.0183], [0.1353, 1.0000, 0.1353], [0.0183, 0.1353, 0.0183]]) heatmap: tensor([[0.6907, 0.7162, 0.5568, 0.7927, 0.4490], [0.1968, 0.5420, 0.2634, 0.8898, 0.6967], [0.7910, 0.6914, 0.3164, 0.8162, 0.4236], [0.4870, 0.5177, 0.9171, 0.7794, 0.6405], [0.2565, 0.9226, 0.3157, 0.8301, 0.4038]]) height, width: 5 5 x: 2 y 2 radius: 1 left, right, top, bottom: 1 2 1 2 masked_heatmap: tensor([[0.5420, 0.2634, 0.8898], [0.6914, 0.3164, 0.8162], [0.5177, 0.9171, 0.7794]]) masked_gaussian: tensor([[0.0183, 0.1353, 0.0183], [0.1353, 1.0000, 0.1353], [0.0183, 0.1353, 0.0183]]) out_heatmap: tensor([[0.6907, 0.7162, 0.5568, 0.7927, 0.4490], [0.1968, 0.5420, 0.2634, 0.8898, 0.6967], [0.7910, 0.6914, 0.3164, 0.8162, 0.4236], [0.4870, 0.5177, 0.9171, 0.7794, 0.6405], [0.2565, 0.9226, 0.3157, 0.8301, 0.4038]]) masked_gaussian * k: tensor([[0.0183, 0.1353, 0.0183], [0.1353, 1.0000, 0.1353], [0.0183, 0.1353, 0.0183]]) tensor([[0.6907, 0.7162, 0.5568, 0.7927, 0.4490], [0.1968, 0.5420, 0.2634, 0.8898, 0.6967], [0.7910, 0.6914, 1.0000, 0.8162, 0.4236], [0.4870, 0.5177, 0.9171, 0.7794, 0.6405], [0.2565, 0.9226, 0.3157, 0.8301, 0.4038]]) tl_off_mask: tensor([[False, False, False, False, False], [False, False, False, False, False], [False, False, True, False, False], [False, False, False, False, False], [False, False, False, False, False]]) tl_off_mask: tensor([0, 0, 1, 0, 0]) tl_off_mask: tensor([False, False, True, False, False]) tl_off_mask: tensor([[False], [False], [ True], [False], [False]]) tl_off_mask: tensor([[0.], [0.], [1.], [0.], [0.]])
- 首先计算偏移损失的掩码:
-
返回损失:
return det_loss, pull_loss, push_loss, off_loss
:返回包含检测损失、关联嵌入损失(如果有)和偏移损失的元组。
12. predict_by_feat 方法
- 功能:将提取的特征转换为检测结果。
- 参数解释:
tl_heats
:每个层级的 top-left 角点热图列表。br_heats
:每个层级的 bottom-right 角点热图列表。tl_embs
:每个层级的 top-left 角点嵌入列表。br_embs
:每个层级的 bottom-right 角点嵌入列表。tl_offs
:每个层级的 top-left 角点偏移列表。br_offs
:每个层级的 bottom-right 角点偏移列表。batch_img_metas
:每个图像的元信息列表。rescale
:如果为True
,则将检测框转换回原始图像空间,默认为False
。with_nms
:如果为True
,则在返回检测框之前进行非极大值抑制,默认为True
。
- 执行过程:
- 对于每个图像,调用
_predict_by_feat_single
方法进行预测。 - 返回包含每个图像检测结果的列表。
- 对于每个图像,调用
13. _predict_by_feat_single 方法
- 功能:对单个图像进行预测,将其特征转换为检测结果。
- 参数解释:
tl_heat
:当前图像的 top-left 角点热图。br_heat
:当前图像的 bottom-right 角点热图。tl_off
:当前图像的 top-left 角点偏移。br_off
:当前图像的 bottom-right 角点偏移。img_meta
:当前图像的元信息。tl_emb
:当前图像的 top-left 角点嵌入(可选)。br_emb
:当前图像的 bottom-right 角点嵌入(可选)。tl_centripetal_shift
:当前图像的 top-left 角点向心偏移(可选)。br_centripetal_shift
:当前图像的 bottom-right 角点向心偏移(可选)。rescale
:如果为True
,则将检测框转换回原始图像空间,默认为False
。with_nms
:如果为True
,则在返回检测框之前进行非极大值抑制,默认为True
。
- 执行过程:
- 使用输入的特征图、偏移和嵌入(如果有)以及图像元信息,解码得到原始边界框预测。
- 如果需要,进行缩放和非极大值抑制操作。
- 返回包含检测结果的
InstanceData
对象。
14. _bboxes_nms 方法
- 功能:对检测框进行非极大值抑制操作。
- 参数解释:
bboxes
:检测框张量,形状为[num_boxes, 5]
,最后一维包含检测分数。labels
:检测框的标签张量。cfg
:测试配置字典。
- 执行过程:
- 根据配置进行非极大值抑制操作,返回抑制后的检测框和标签。
15. _decode_heatmap 方法
- 功能:将输出转换为检测框的原始预测,包括边界框坐标、分数和类别。
- 参数解释:
tl_heat
:top-left 角点热图。br_heat
:bottom-right 角点热图。tl_off
:top-left 角点偏移。br_off
:bottom-right 角点偏移。tl_emb
:top-left 角点嵌入(可选)。br_emb
:bottom-right 角点嵌入(可选)。tl_centripetal_shift
:top-left 角点向心偏移(可选)。br_centripetal_shift
:bottom-right 角点向心偏移(可选)。img_meta
:图像元信息。k
:从热图中获取的 top k 角点数量。kernel
:用于提取局部最大像素的最大池化核大小。distance_threshold
:角点之间的距离阈值,用于判断是否来自同一对象。num_dets
:进行非极大值抑制之前的原始检测框数量。
- 执行过程:
- 首先,根据是否有嵌入或向心偏移进行相应的处理。
- 对热图进行局部最大池化操作,并获取 top k 角点的位置、分数和类别。
- 根据角点位置和偏移计算边界框的四个顶点坐标。
- 如果有向心偏移,进行额外的处理以计算更准确的边界框。
- 根据各种条件(如类别不一致、距离过大、宽度或高度不合理等)拒绝不合适的边界框。
- 对分数进行排序,获取前
num_dets
个检测框。 - 返回包含边界框、分数和类别的张量。
通过对CornerHead
中的这些关键函数的解析,我们可以更好地理解其在目标检测任务中的作用和实现方式,为进一步研究和使用 MMDetection 中的目标检测头提供了深入的理解。