Kornia 2D 边界框形状校验修复：`infer_bbox_shape` 与 `bbox_to_mask` 拒绝 rank-4 批处理输入并抛出 `ShapeError`

• 计算机视觉
• 人工智能
• 深度学习
• 图像处理

本文解读 Kornia 几何模块中的一项行为修复（对应 changelog 条目 changelog.d/+migration-090.fixed.md，issue #4218）：两个 2D 边界框辅助函数infer_bbox_shape与bbox_to_mask现在会明确拒绝(B, N, 4, 2)这种 rank-4 批处理输入并抛出ShapeError，取代了过去N < 3时抛出偶发IndexError、N >= 3时静默返回由错误顶点算出的(B, 2)尺寸的隐蔽行为。读完本文，你将理解该修复的动机、底层校验实现、测试覆盖情况，以及迁移现有调用代码的正确姿势。

一、变更背景：旧行为为何危险

在修复之前，infer_bbox_shape和bbox_to_mask对四维输入的处理是"碰运气"式的：

• 当批内框数N < 3时，函数在索引顶点坐标时（例如boxes[:, 1, 0]）会意外地触发IndexError。这个异常信息与"形状不合法"毫无关联，用户无法据此判断问题出在输入形状上。
• 当批内框数N >= 3时，索引不会越界，函数会"成功"执行并返回一个形状为(B, 2)的结果——但这个结果是从错误的顶点计算出来的（因为 rank-4 张量在boxes[:, 1, 0]这类索引下，1被当成了第 1 个框的编号，而不是第 1 个顶点），数值完全不可信，却又没有任何报错提示。

也就是说，旧的实现把"输入形状错误"这一缺陷藏在了两个极端里：要么报出误导性的异常，要么悄悄返回错误数值。这类静默失败在几何计算管线中尤其危险——下游的裁剪、掩码、数据增强可能基于错误的宽高继续计算，最终污染整个训练或推理结果。

二、新行为：统一以ShapeError拒绝 rank-4 输入

修复后，两个函数的行为一致且明确：

• infer_bbox_shape与bbox_to_mask只接受非批处理的(N, 4, 2)/(B, 4, 2)输入；
• 传入(B, N, 4, 2)形式的 rank-4 张量时，立即抛出kornia.core.exceptions.ShapeError；
• 不再存在"N 较小时报IndexError、N 较大时静默算错"的分叉行为。

ShapeError是 Kornia 在 kornia/core/exceptions.py 中定义的校验异常基类BaseError的子类，专门用于"张量形状校验失败"场景，并带有actual_shape与expected_shape两个结构化属性，便于上层捕获后给出清晰的诊断信息。

三、源码实现剖析：校验是如何落地的

两个函数的形状校验都落在 kornia/geometry/bbox.py 的入口处：

# infer_bbox_shape 入口 KORNIA_CHECK_SHAPE(boxes, ["N", "4", "2"]) # bbox_to_mask 入口 KORNIA_CHECK_SHAPE(boxes, ["B", "4", "2"])

KORNIA_CHECK_SHAPE定义于 kornia/core/check.py，它按给定维度列表校验张量形状。当维度数不匹配时，会抛出ShapeError，错误消息形如：

Shape dimension mismatch: expected 3 dimensions, got 4. Expected shape: ['N', '4', '2'] Actual shape: [1, 2, 4, 2]

这正是测试中match="expected 3 dimensions, got 4"断言所匹配的消息。值得注意的是，该校验支持通过KORNIA_CHECKS环境变量或disable_checks()全局关闭（TorchScript 编译路径下则始终执行），但形状错误属于结构性错误，建议保持默认开启。

从源码注释可以看出，这次的 2D 修复还"传染"给了 3D 孪生函数：infer_bbox_shape3d与bbox_to_mask3d在 kornia/geometry/bbox.py 中同样加入了KORNIA_CHECK_SHAPE校验（分别要求["N", "8", "3"]与["B", "8", "3"]），源码注释明确写道这是 "the way #4218 did for the 2D twin"。因此本文的迁移建议对 3D 版本同样适用。

底层实现细节：为什么旧代码会算错

看 kornia/geometry/bbox.py 中infer_bbox_shape的核心计算：

width: torch.Tensor = boxes[:, 1, 0] - boxes[:, 0, 0] + 1 height: torch.Tensor = boxes[:, 2, 1] - boxes[:, 0, 1] + 1 return height, width

这里宽度和高度是从固定顶点索引读取的（第 1 个顶点与第 0 个顶点之差），而不是通过max - min归约得到。对于(B, N, 4, 2)输入，boxes[:, 1, 0]中的1会被解释为"第 1 个框"而非"第 1 个顶点"，于是计算出的 extents 来自错误的坐标组合——这就是旧行为中N >= 3时"静默返回错误结果"的根源。新增的形状校验正是为了在进入这段索引逻辑之前拦截这类输入。

bbox_to_mask同样只读取第 0 个（左上）与第 2 个（右下）顶点：

x_min = boxes[:, 0, 0].view(-1, 1, 1) y_min = boxes[:, 0, 1].view(-1, 1, 1) x_max = boxes[:, 2, 0].view(-1, 1, 1) y_max = boxes[:, 2, 1].view(-1, 1, 1)

对于 rank-4 输入，这些索引同样会错位，最终生成完全错误的掩码。形状校验保证了这一路径永远不会被错误形状的输入触达。

四、测试验证：修复被如何锁定

test_bbox.py 中新增了两个参数化测试，将修复行为固化下来：

@pytest.mark.parametrize("num_boxes", [2, 3]) def test_infer_bbox_shape_rejects_rank4_4180(self, device, dtype, num_boxes): boxes = torch.zeros(1, num_boxes, 4, 2, device=device, dtype=dtype) with pytest.raises(ShapeError, match="expected 3 dimensions, got 4"): infer_bbox_shape(boxes) @pytest.mark.parametrize("num_boxes", [2, 3]) def test_bbox_to_mask_rejects_rank4_4180(self, device, dtype, num_boxes): boxes = torch.zeros(1, num_boxes, 4, 2, device=device, dtype=dtype) with pytest.raises(ShapeError, match="expected 3 dimensions, got 4"): bbox_to_mask(boxes, width=5, height=5)

关键点在于参数化取值num_boxes覆盖了2 和 3——这正是旧行为的分界点：N = 2曾触发偶发的IndexError，N = 3曾静默返回错误结果。两个取值现在都必须统一抛出ShapeError，从而完整覆盖了旧缺陷的两个表现形态。

五、迁移指引：收到ShapeError后怎么办

如果你的代码之前直接向这两个函数传入过(B, N, 4, 2)张量（例如来自kornia.geometry.boxes.Boxes的data属性或to_tensor('vertices')导出结果），修复后会开始抛出ShapeError。迁移方式很简单：先展平再调用。

import torch from kornia.geometry.bbox import infer_bbox_shape, bbox_to_mask # 之前的错误用法：(B, N, 4, 2) rank-4 输入会被 ShapeError 拒绝 # boxes_batched = torch.zeros(2, 3, 4, 2) # 正确用法一：整体展平，批量计算全部 N*B 个框 boxes_flat = torch.zeros(2 * 3, 4, 2) heights, widths = infer_bbox_shape(boxes_flat) # 形状均为 (N*B,) masks = bbox_to_mask(boxes_flat, width=32, height=32) # 形状为 (N*B, 32, 32) # 正确用法二：按批次索引出单个 batch 再调用 heights_b, widths_b = infer_bbox_shape(boxes_batched[0]) # 形状为 (N,)

与Boxes面向对象 API 的对照关系（见 kornia/geometry/boxes.py 的约定说明）：

• infer_bbox_shape等价于Boxes.get_boxes_shape()，但前者要求非批处理(N, 4, 2)，而后者直接支持(B, N, 4, 2)的存储结构；且get_boxes_shape基于amin/amax归约计算，两者对非规范顶点顺序的框结果可能不同。
• bbox_to_mask(boxes, width, height)与Boxes.to_mask(height, width)都用于生成掩码，但参数顺序相反（函数版是(width, height)，方法版是(height, width)），迁移时务必注意。
• 两个函数都按**闭区间（inclusive）**语义读取顶点，即width = xmax - xmin + 1。因此如果数据来自Boxes，应传入'vertices_plus'导出（而不是'vertices'，后者会被读成每个轴大 1 像素）。

六、总结

本次修复（issue #4218）把infer_bbox_shape与bbox_to_mask对 rank-4 输入的两种故障模式——误导性的IndexError与静默的错误(B, 2)结果——统一收敛为清晰的ShapeError，并通过参数化测试在N = 2与N = 3两个关键边界上加以锁定。对于使用这批低层几何辅助函数的代码，只需在调用前将(B, N, 4, 2)展平为(B*N, 4, 2)或按批次索引，即可无缝适配新行为；3D 版本的infer_bbox_shape3d与bbox_to_mask3d也采用了相同的校验策略，迁移思路一致。

• 计算机视觉
• 人工智能
• 深度学习
• 图像处理