Kornia RandomRain 雨滴增强修复解析：均匀起始位置与“最后一行/列可被绘制”-育师

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

【免费下载链接】kornia

🐍 Geometric Computer Vision Library for Spatial AI

项目地址：https://gitcode.com/gh_mirrors/ko/kornia

RandomRain是 Kornia 随机增强模块（kornia.augmentation）中用于为图像叠加雨滴纹理的强度类变换。本篇文章以 changelog.d/4612.fixed.md 这一修复记录为主体，深入解析本次变更的核心内容——雨滴起始位置改为在“所有能让整条雨滴完整落在图像内”的候选位置间均匀分布，从而使得图像的最后一行与最后一列首次可以被绘制。读完本文，你将掌握RandomRain的参数语义、雨滴在图像上的绘制几何、起始位置合法域的推导方式，以及对应的源码实现与测试验证路径，可直接用于在实际增强流水线中规避边界无法命中、尺寸越界等陷阱。

变更记录说了什么

该 changelog 片段全文只有一句话：

RandomRainnow places drops uniformly over every start position that keeps the drop inside the image, so the last row and column can be painted.

它描述了一次针对bug 的行为修正：修复前，某些合法的起始位置无法被命中，导致雨滴永远画不到图像的最后一行与最后一列；修复后，每个“能把整条雨滴保持在图像内”的起始位置被等概率选取。注意其中两个关键词：

uniformly（均匀）：不是简单地“允许”画到最后一行，而是所有合法起始位置共享同等概率；
keeps the drop inside the image（保持雨滴完整在图像内）：合法域的判定标准是雨滴的“远端”不越界，而不是起点不越界。

这段记录属于 Kornia 的 changelog.d 碎片化变更机制：每个带用户可见变更的 PR 在changelog.d下以<PR>.<type>.md命名（4612.fixed.md即 PR 4612 的fixed类条目），发布时由 Towncrier 汇总进CHANGELOG.md。因此本文所述行为以当前仓库源码为准。

RandomRain 是什么

RandomRain定义在 kornia/augmentation/_2d/intensity/random_rain.py，继承自IntensityAugmentationBase2D，构造签名如下：

RandomRain( number_of_drops: tuple[int, int] = (1000, 2000), drop_height: tuple[int, int] = (5, 20), drop_width: tuple[int, int] = (-5, 5), same_on_batch: bool = False, p: float = 0.5, keepdim: bool = False, )

各参数的含义与约束（均来自类文档字符串）：

参数	默认值	语义与约束
`number_of_drops`	`(1000, 2000)`	每张图像的雨滴数量区间（闭区间均匀整数采样）
`drop_height`	`(5, 20)`	雨滴高度，对应图像行方向；采样值须严格小于图像高度
`drop_width`	`(-5, 5)`	雨滴宽度，对应图像列方向，可正可负；绝对值须严格小于图像宽度
`same_on_batch`	`False`	为`True`时整批样本共享相同的雨滴数量、尺寸与坐标
`p`	`0.5`	应用该变换的概率
`keepdim`	`False`	是否保持输入维度（与`IntensityAugmentationBase2D`约定一致）

输入形状为(C, H, W)或(B, C, H, W)，输出统一为(B, C, H, W)。前向时输入通道数必须是 1 或 3，否则抛出“Number of color channels should be 1 or 3.”。

一个可直接运行的例子（取自类文档）：

import torch import kornia as K torch.manual_seed(0) input = torch.rand(1, 1, 5, 5) rain = K.augmentation.RandomRain( p=1.0, drop_height=(1, 2), drop_width=(1, 2), number_of_drops=(1, 1) ) rain(input) # 输出形状仍为 (1, 1, 5, 5)，部分像素被替换为 200/255

雨滴的绘制几何：span 是“跨度”而非“数量”

要理解“最后一行能否被画到”，必须先理解雨滴的绘制模型。在 apply_transform 中，每条雨滴的线条形状通过torch.linspace生成：

size_of_line = max(height_of_drop, abs(width_of_drop)) lines[i, 0, :size_of_line] = torch.linspace(0, height_of_drop, steps=size_of_line, dtype=torch.long) lines[i, 1, :size_of_line] = torch.linspace(0, width_of_drop, steps=size_of_line, dtype=torch.long)

由此可推出几个关键几何事实：

跨度而非数量：一条drop_height=h、drop_width=w的雨滴，其行/列偏移量取自linspace(0, h, steps=max(h, |w|))，端点包含在内。当max(h, |w|) >= 2时，雨滴实际占h + 1行（或|w| + 1列）；当两者都<= 1时只有一个采样点，雨滴退化为单像素。
尺寸差会导致内部空隙：由于步数是max(h, |w|)，当两尺寸不等时，被涂色的格子在该跨度内并不连续。例如drop_height=5, drop_width=0在6 x 10图像上涂色的行是[0, 1, 2, 3, 5]（行 4 有空隙）。
尺寸以图像为参照：drop_height沿行方向、drop_width沿列方向，负的drop_width使雨滴向另一侧倾斜——参数命名的是图像坐标轴而非雨滴局部坐标。
雨滴值是常量：被涂像素固定写为200 / 255，与图像内容无关（见 random_rain.py 的赋值语句），因此对任意取值范围的输入，“雨”都表现为深色痕迹。这一点由测试 test_convention_random_rain_drop_value_is_two_hundred_over_255 固定：全2.0的图像经处理后min=200/255、max=2.0。

修复核心：起始位置的合法域与均匀性

合法域的推导

关键代码在 random_rain.py 的 meta 构造段：

last_dy, last_dx = (height_of_drop, width_of_drop) if size_of_line > 1 else (0, 0) meta[i, 0] = image_height - last_dy # 行方向允许的起始范围 meta[i, 1] = image_width - abs(last_dx) # 列方向允许的起始范围 meta[i, 2] = max(-last_dx, 0) # 负宽度时的列偏移

last_dy/last_dx即雨滴线条的“最远端偏移量”（端点包含在采样中）。合法域要求：起始行 + 最远行偏移 ≤ image_height - 1，即起始行 ∈[0, image_height - last_dy]；列同理，且负宽度时起点还需加上max(-last_dx, 0)使远端（而非起点）约束生效。注释对此的概括是：“a drop may start anywhere its far end stays inside the image”——雨滴可在远端仍留在图内的任意位置起始。

均匀采样与 MPS 钳制

起始坐标来自生成器产出的coordinates_factor（形状(B, max_drops, 2)，取值[0, 1]），由 random_rain.py 的 start 计算段换算：

start_rows = torch.minimum((coordinates[..., 0] * rows_c).long(), rows_c - 1) start_cols = torch.minimum((coordinates[..., 1] * cols_c).long(), cols_c - 1) + shift_c

坐标在[0, 1]上均匀采样，乘以合法范围后再取整，因此每个合法起始位置被等概率选中——这正是 changelog 中 “uniformly over every start position” 的落地方式。torch.minimum(..., rows_c - 1)的钳制针对 MPS 半精度rand可能恰好返回1.0的问题（源码注释引用 #4553）：它把这种采样保留在最后一个合法起始位置，而不是越界到范围之外。

于是：对高度H的图像，最大起始行为H - h - 1，雨滴占据h + 1行，末端恰好落在第H - 1行——最后一行可达；同理最后一列也可达。修复前部分合法起点被排除，导致最后一行/列永远画不上，这正是本片段修复的 bug。

边界行为一览（来自类文档与测试）

尺寸等于图像尺寸（drop_height == H或abs(drop_width) == W）会在前向抛出BaseError（“Height of drop should be greater than zero and less than image height.” / “Width of drop should be less than image width.”）；比图像小 1 个像素则正常运行。校验发生在forward（此时才知道图像形状），构造阶段不报错——见测试 test_convention_random_rain_rejects_a_drop_as_large_as_the_image。
由于默认drop_height=(5, 20)、drop_width=(-5, 5)的上界对尺寸规则“实时生效”，20 像素高的图像或 5 像素宽的图像在部分随机种子下会抛错；高度 ≤ 5 像素的图像则因不存在任何合法高度而必然抛错。
均匀性还带来一个推论：跨度内部的空隙可能使某列仍未被涂到——例如drop_height=5, drop_width=9作用于6 x 10图像时只有一个合法起点，且列 8 永远不会被画到。“可达”指可由某个起点命中，而非必然被涂色。

参数生成器：闭区间均匀整数采样

起始坐标之外，number_of_drops、drop_height、drop_width三个整数区间也要求“闭区间均匀”。实现位于 kornia/augmentation/random_generator/_2d/random_rain.py 的RainGenerator：

_closed_integer_range把闭区间(lo, hi)转成半开采样区间[lo, hi + 1)，并拒绝小数、颠倒区间（lower > upper）与非有限值；
_draw_closed_integer使用floor()（而非向零截断）并对上端点hi + 1做clamp。注释给出了一个精妙的数值分析：对float32下最大的u = 1 - 2**-24，5 + u * 16恰好等于21.0，若不加钳制，约每2**24次采样会有一次越界到21；而向零截断则会让(-1, 0)与(0, 1)都折叠到0，使0被抽到两倍概率。

因此默认drop_height=(5, 20)可以抽到20，drop_width=(-5, 5)可以抽到-5与5，0不比任何其他宽度权重更大。该行为由 test_convention_random_rain_integer_ranges_are_closed_and_uniform 用 22000 次采样验证：每个整数出现次数与均匀期望偏差在 15% 以内，且极值（hi）确实会被抽到。值得注意的是，这段闭区间修复本身是另一个 changelog 条目（changelog.d/4568.fixed.md）的主题，与本文的起始位置修复互为表里：前者保证“参数值可被均匀抽到”，后者保证“抽到的值能均匀落位”。

源码级实现：批量化光栅化与单次索引写入

本次修复所依托的apply_transform采用了一次批量化重写（源码注释引用 #4530）。核心手段包括：

主机端一次性构建所有雨滴形状：lines（(B, 2, longest_line)）与meta（(B, 5)：合法起始行、合法起始列、列偏移、线条长度、雨滴数量）在 CPU 上按样本填充，再以一次拼接拷贝上设备（host = torch.cat([lines.reshape(-1), meta.reshape(-1)]).to(device)），避免每样本一次的主机-设备拷贝。
显式固定 CPU 设备：中间缓冲区用device="cpu"显式指定，避免调用方设置torch.set_default_device后主机循环退化为逐样本设备启动，或产生无法从meta张量拷贝的问题。
单一nonzero完成掩码解析：drop_valid（雨滴序号 < 该样本雨滴数）与step_valid（步序号 < 线条长度）按位与后展平，用一次nonzero()取得有效索引，再从扁平位置算术反推批次号，避免对三个张量分别做布尔索引而触发三次nonzero。
一次索引赋值写入：所有有效 (batch, channel, row, col) 一次性执行output[...] = 200 / 255。由于所有像素写的是同一个常量，写入顺序不影响结果。

这版实现与“逐条雨滴逐步写入”的朴素参考实现做了逐像素相等对比，见测试 test_batched_rasterisation_matches_per_drop_reference_4530——只要重写版移动了哪怕一个像素，断言都会失败。

测试如何锁定新行为

围绕本次修复的行为，仓库提供了多层次的回归保护：

tests/augmentation/test_augmentation.py 的 TestRandomRain：覆盖形状保持（test_cardinality）、尺寸边界（等于图像尺寸时报错、小 1 像素正常）、零概率（p=0.0透传）、same_on_batch=True时全批参数一致，以及构造时不受环境默认设备影响的特性（test_construction_ignores_the_ambient_default_device，在torch.device("meta")上下文内构造仍得到 CPU 上的采样器）。
tests/augmentation/test_conventions_intensity_ops.py：以“约定测试”形式固定行列命名（drop_height沿行、drop_width沿列，换轴后外延随之互换）、常量雨滴值200/255、跨度语义（h尺寸的雨滴占h+1行）、尺寸等于图像时报错、闭区间均匀性，以及1 - 2**-24端点钳制行为。