物件偵測

物件偵測模型用於辨識圖像中的目標，而物件偵測資料集則應用於自動駕駛及偵測野火等自然災害。本指南將向您展示如何依照 Albumentations 的教學課程，對物件偵測資料集套用轉換。

若要執行這些範例，請確保您已安裝最新版本的 albumentations 和 cv2。

pip install -U albumentations opencv-python

在此範例中，您將使用 cppe-5 資料集，用於在 COVID-19 疫情背景下辨識醫療個人防護裝備 (PPE)。

載入資料集並查看其中一個範例。

>>> from datasets import load_dataset

>>> ds = load_dataset("rishitdagli/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：包含圖像中物件之邊界框（bounding box）詮釋資料的字典。
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() 函數即時 (on-the-fly) 套用轉換，這可以節省磁碟空間。資料增強的隨機性可能導致您在存取同一個範例兩次時得到不同的圖像。這在訓練模型多個 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']]
...     )
... )

現在您已了解如何處理物件偵測的資料集，請進一步學習如何訓練物件偵測模型並將其用於推論。

在 GitHub 上更新

資料集 (Datasets)

物件偵測