博客

将 TorchVision 的转换扩展到目标检测、分割和视频任务

注意：本文的前一个版本发布于 2022 年 11 月。鉴于 2023 年 3 月 torchvision 随 PyTorch 2.0 一同发布的 0.15 版本，我们已对本文内容进行了最新信息的更新。

TorchVision 正在扩展其 Transforms API！新功能如下：

  • 您可以将它们不仅用于图像分类，还用于目标检测、实例和语义分割以及视频分类。
  • 您可以使用新的函数式转换来转换视频、边界框和分割掩码。

该 API 与旧版本完全向后兼容，保持不变以辅助迁移和采用。我们现已在 torchvision.transforms.v2 命名空间中以 Beta 版形式发布此新 API，非常期待收到您的早期反馈以改进其功能。如果您有任何问题或建议，请联系我们。

当前变换的局限性

TorchVision 现有的 Transforms API（即 V1）仅支持单张图像。因此，它只能用于分类任务。

from torchvision import transforms
trans = transforms.Compose([
   transforms.ColorJitter(contrast=0.5),
   transforms.RandomRotation(30),
   transforms.CenterCrop(480),
])
imgs = trans(imgs)

上述方法不支持目标检测或分割。这一局限使得任何非分类计算机视觉任务都处于“二等公民”地位，因为无法利用 Transforms API 执行必要的增强操作。从历史上看，这使得利用 TorchVision 的原语训练高精度模型变得困难，从而导致我们的模型库（Model Zoo）较 SOTA（当前最优）水平落后数个百分点。

为了规避这一局限，TorchVision 在其参考脚本中提供了自定义实现，展示了如何在每个任务中执行增强。尽管这种做法使我们能够训练出高精度的分类、目标检测和分割模型，但它是一种“黑客”式的变通方法，导致这些变换无法直接从 TorchVision 二进制文件中导入。

全新的 Transforms API

Transforms V2 API 支持视频、边界框（bounding boxes）和分割掩码（segmentation masks），这意味着它为许多计算机视觉任务提供了原生支持。新方案可作为直接替代品。

import torchvision.transforms.v2 as transforms

# Exactly the same interface as V1:
trans = transforms.Compose([
    transforms.ColorJitter(contrast=0.5),
    transforms.RandomRotation(30),
    transforms.CenterCrop(480),
])
imgs, bboxes, labels = trans(imgs, bboxes, labels)

新的变换类（Transform Classes）可以接收任意数量的输入，而无需强制特定的顺序或结构。

# Already supported:
trans(imgs)  # Image Classification
trans(videos)  # Video Tasks
trans(imgs, bboxes, labels)  # Object Detection
trans(imgs, bboxes, masks, labels)  # Instance Segmentation
trans(imgs, masks)  # Semantic Segmentation
trans({"image": imgs, "box": bboxes, "tag": labels})  # Arbitrary Structure

# Future support:
trans(imgs, bboxes, labels, keypoints)  # Keypoint Detection
trans(stereo_images, disparities, masks)  # Depth Perception
trans(image1, image2, optical_flows, masks)  # Optical Flow
trans(imgs_or_videos, labels)  # MixUp/CutMix-style Transforms

变换类可确保将相同的随机变换应用于所有输入，从而保证结果的一致性。

功能性 API 已更新，以支持所有输入类型所需的所有信号处理内核（调整大小、裁剪、仿射变换、填充等）。

from torchvision.transforms.v2 import functional as F


# High-level dispatcher, accepts any supported input type, fully BC
F.resize(inpt, size=[224, 224])
# Image tensor kernel
F.resize_image_tensor(img_tensor, size=[224, 224], antialias=True) 
# PIL image kernel
F.resize_image_pil(img_pil, size=[224, 224], interpolation=BILINEAR)
# Video kernel
F.resize_video(video, size=[224, 224], antialias=True) 
# Mask kernel
F.resize_mask(mask, size=[224, 224])
# Bounding box kernel
F.resize_bounding_box(bbox, size=[224, 224], spatial_size=[256, 256])

在底层，该 API 使用 Tensor 子类化来包装输入、附加有用的元数据并分发到正确的内核。为了使您的数据与这些新变换兼容，您可以使用提供的数据集包装器（该包装器应适用于大多数 torchvision 内置数据集），或者手动将数据包装为 Datapoints。

from torchvision.datasets import wrap_dataset_for_transforms_v2
ds = CocoDetection(..., transforms=v2_transforms)
ds = wrap_dataset_for_transforms_v2(ds) # data is now compatible with transforms v2!

# Or wrap your data manually using the lower-level Datapoint classes:
from torchvision import datapoints

imgs = datapoints.Image(images)
vids = datapoints.Video(videos)
masks = datapoints.Mask(target["masks“])
bboxes = datapoints.BoundingBox(target["boxes“], format=”XYXY”, spatial_size=imgs.shape)

除了新 API 之外，我们现在还为 SOTA 研究中使用的多种数据增强方法提供了可导入的实现，例如大规模抖动（Large Scale Jitter）、自动增强（AutoAugmentation）方法以及若干新的几何、颜色和类型转换变换。

该 API 继续支持 PIL 和 Tensor 后端的图像输入，支持单输入或批处理输入，并保持了函数式和类 API 的 JIT 脚本化能力。经验证，新 API 可实现与先前实现相同的精度。

端到端示例

以下是使用此图像演示新 API 的示例。它既适用于 PIL 图像也适用于 Tensor。如需更多示例和教程，请查看我们的图库！

from torchvision import io, utils
from torchvision import datapoints
from torchvision.transforms import v2 as T
from torchvision.transforms.v2 import functional as F

# Defining and wrapping input to appropriate Tensor Subclasses
path = "COCO_val2014_000000418825.jpg"
img = datapoints.Image(io.read_image(path))
# img = PIL.Image.open(path)
bboxes = datapoints.BoundingBox(
    [[2, 0, 206, 253], [396, 92, 479, 241], [328, 253, 417, 332],
     [148, 68, 256, 182], [93, 158, 170, 260], [432, 0, 438, 26],
     [422, 0, 480, 25], [419, 39, 424, 52], [448, 37, 456, 62],
     [435, 43, 437, 50], [461, 36, 469, 63], [461, 75, 469, 94],
     [469, 36, 480, 64], [440, 37, 446, 56], [398, 233, 480, 304],
     [452, 39, 463, 63], [424, 38, 429, 50]],
    format=datapoints.BoundingBoxFormat.XYXY,
    spatial_size=F.get_spatial_size(img),
)
labels = [59, 58, 50, 64, 76, 74, 74, 74, 74, 74, 74, 74, 74, 74, 50, 74, 74]
# Defining and applying Transforms V2
trans = T.Compose(
    [
        T.ColorJitter(contrast=0.5),
        T.RandomRotation(30),
        T.CenterCrop(480),
    ]
)
img, bboxes, labels = trans(img, bboxes, labels)
# Visualizing results
viz = utils.draw_bounding_boxes(F.to_image_tensor(img), boxes=bboxes)
F.to_pil_image(viz).show()

开发里程碑与未来工作

以下是我们目前的开发进度：

  • 设计 API
  • 编写用于变换视频、边界框、掩码和标签的内核
  • 基于新 API 重写所有现有变换类（稳定版 + 参考版）
    • 图像分类
    • 视频分类
    • 目标检测
    • 实例分割
    • 语义分割
  • 验证新 API 在所有支持的任务和后端上的精度
  • 速度基准测试和性能优化（进行中 – 计划于 12 月完成）
  • 脱离原型阶段（计划于第一季度完成）
  • 增加对深度感知、关键点检测、光流等功能的支持（未来计划）
  • 增加对 MixUp 和 CutMix 等批处理级变换的平滑支持

我们非常期待收到您的反馈以改进其功能。如果您有任何问题或建议，请联系我们。