如何实现ImageOps.Solarize的PyTorch Tensor批量处理版本

PIL的ImageOps.solarize逻辑非常简单：对所有像素值高于设定阈值的通道值做反相处理，不存在复杂的自定义算子依赖，PyTorch生态下有成熟的原生张量实现，完全可以直接对批量张量做处理，不需要循环转PIL单张操作。

官方内置实现 #

TorchVision 0.9.0及以上版本已经自带支持批量张量输入的Solarize接口，处理结果和PIL实现完全一致：

• 函数式接口：torchvision.transforms.functional.solarize
  支持输入形状为(B, C, H, W)的批量张量，同时兼容两种常用图像张量格式：0-255值域的uint8张量、归一化到0-1值域的浮点张量，整个处理流程不经过PIL格式转换，批量处理效率远高于循环调用PIL接口。
• 增强类接口：torchvision.transforms.RandomSolarize
  可以直接接入现有PyTorch数据增强流水线，支持按指定概率对输入批量执行Solarize增强，和其他TorchVision的张量变换完全兼容。

基础调用示例：

import torch
import torchvision.transforms.functional as F

# 模拟输入：8张3通道224*224的归一化图像张量
batch_tensor = torch.rand(8, 3, 224, 224)
# 阈值设为0.5，对应uint8格式下的常用阈值128
solarized_batch = F.solarize(batch_tensor, threshold=0.5)
# 输出张量和输入形状、设备、dtype完全一致，可直接送入后续模型
print(solarized_batch.shape)  # 输出: torch.Size([8, 3, 224, 224])

低版本兼容手动实现 #

如果使用的TorchVision版本低于0.9.0，没有内置接口，仅需2行核心代码即可实现逻辑完全对齐的批量处理函数，支持任意batch size：

def batch_solarize(img_tensor: torch.Tensor, threshold: float) -> torch.Tensor:
    # 自动识别张量值域，适配uint8/浮点两种输入格式
    max_pixel = 1.0 if torch.is_floating_point(img_tensor) else 255
    return torch.where(img_tensor >= threshold, max_pixel - img_tensor, img_tensor)

使用注意事项 #

• 输入张量必须使用PyTorch默认的(批量维, 通道, 高度, 宽度)维度顺序，不要传入PIL默认的(高度, 宽度, 通道)格式张量，否则会出现处理逻辑错误。
• 阈值参数需要和输入张量的值域匹配：处理0-255值域的uint8张量时，阈值通常取128；处理0-1值域的浮点张量时，阈值对应取0.5即可，输出结果和PIL单张处理的结果完全一致。

内容的提问来源于stack exchange，提问作者Gooby