目标检测
目标检测模型用于识别图像中的物体,而目标检测数据集则用于自动驾驶、火灾等自然灾害检测等应用。本指南将向您展示如何根据来自 Albumentations 的 教程 对目标检测数据集应用转换。
要运行这些示例,请确保您已安装最新版本的 albumentations
和 cv2
。
pip install -U albumentations opencv-python
在本示例中,您将使用 cppe-5
数据集来识别 COVID-19 大流行背景下医疗个人防护装备 (PPE)。
加载数据集并查看示例
>>> from datasets import load_dataset
>>> ds = load_dataset("cppe-5")
>>> example = ds['train'][0]
>>> example
{'height': 663,
'image': <PIL.JpegImagePlugin.JpegImageFile image mode=RGB size=943x663 at 0x7FC3DC756250>,
'image_id': 15,
'objects': {'area': [3796, 1596, 152768, 81002],
'bbox': [[302.0, 109.0, 73.0, 52.0],
[810.0, 100.0, 57.0, 28.0],
[160.0, 31.0, 248.0, 616.0],
[741.0, 68.0, 202.0, 401.0]],
'category': [4, 4, 0, 0],
'id': [114, 115, 116, 117]},
'width': 943}
数据集包含以下字段
image
: 包含图像的 PIL.Image.Image 对象。image_id
: 图像 ID。height
: 图像高度。width
: 图像宽度。objects
: 一个字典,包含图像中对象的边界框元数据id
: 标注 ID。area
: 边界框的面积。bbox
: 对象的边界框(采用 coco 格式)。category
: 对象的类别,可能的值包括Coverall (0)
、Face_Shield (1)
、Gloves (2)
、Goggles (3)
和Mask (4)
。
您可以使用一些内部 torch 工具在图像上可视化 bboxes
。为此,您需要参考与类别 ID 关联的 ClassLabel 特性,以便查找字符串标签
>>> import torch
>>> from torchvision.ops import box_convert
>>> from torchvision.utils import draw_bounding_boxes
>>> from torchvision.transforms.functional import pil_to_tensor, to_pil_image
>>> categories = ds['train'].features['objects'].feature['category']
>>> boxes_xywh = torch.tensor(example['objects']['bbox'])
>>> boxes_xyxy = box_convert(boxes_xywh, 'xywh', 'xyxy')
>>> labels = [categories.int2str(x) for x in example['objects']['category']]
>>> to_pil_image(
... draw_bounding_boxes(
... pil_to_tensor(example['image']),
... boxes_xyxy,
... colors="red",
... labels=labels,
... )
... )
使用 albumentations
,您可以应用会影响图像并相应更新 bboxes
的转换。在本例中,图像调整大小为 (480, 480),水平翻转并变亮。
>>> import albumentations
>>> import numpy as np
>>> transform = albumentations.Compose([
... albumentations.Resize(480, 480),
... albumentations.HorizontalFlip(p=1.0),
... albumentations.RandomBrightnessContrast(p=1.0),
... ], bbox_params=albumentations.BboxParams(format='coco', label_fields=['category']))
>>> image = np.array(example['image'])
>>> out = transform(
... image=image,
... bboxes=example['objects']['bbox'],
... category=example['objects']['category'],
... )
现在,当您可视化结果时,图像应该被翻转,但 bboxes
应该仍然在正确的位置。
>>> image = torch.tensor(out['image']).permute(2, 0, 1)
>>> boxes_xywh = torch.stack([torch.tensor(x) for x in out['bboxes']])
>>> boxes_xyxy = box_convert(boxes_xywh, 'xywh', 'xyxy')
>>> labels = [categories.int2str(x) for x in out['category']]
>>> to_pil_image(
... draw_bounding_boxes(
... image,
... boxes_xyxy,
... colors='red',
... labels=labels
... )
... )
创建一个函数来将转换应用于一批示例
>>> def transforms(examples):
... images, bboxes, categories = [], [], []
... for image, objects in zip(examples['image'], examples['objects']):
... image = np.array(image.convert("RGB"))
... out = transform(
... image=image,
... bboxes=objects['bbox'],
... category=objects['category']
... )
... images.append(torch.tensor(out['image']).permute(2, 0, 1))
... bboxes.append(torch.tensor(out['bboxes']))
... categories.append(out['category'])
... return {'image': images, 'bbox': bboxes, 'category': categories}
使用 set_transform() 函数即时应用转换,从而减少磁盘空间的使用。数据增强的随机性可能会在您两次访问同一示例时返回不同的图像。在训练模型多个 epoch 时,这尤其有用。
>>> ds['train'].set_transform(transforms)
您可以通过可视化第 10 个示例来验证转换是否有效
>>> example = ds['train'][10]
>>> to_pil_image(
... draw_bounding_boxes(
... example['image'],
... box_convert(example['bbox'], 'xywh', 'xyxy'),
... colors='red',
... labels=[categories.int2str(x) for x in example['category']]
... )
... )
现在您已经了解了如何处理用于对象检测的数据集,请学习 如何训练对象检测模型 并将其用于推理。