Datasets documentation
Обнаружение объектов
Обнаружение объектов
Модели обнаружения объектов определяют объекты на изображении, а наборы данных для обнаружения объектов используются в таких приложениях, как автономное вождение и выявление природных опасностей, например лесных пожаров. В этом руководстве показано, как применять преобразования к набору данных для обнаружения объектов, следуя туториалу от Albumentations.
Чтобы запустить эти примеры, убедитесь, что у вас установлены актуальные версии albumentations и cv2:
pip install -U albumentations opencv-python
В этом примере вы будете использовать набор данных cppe-5 для распознавания медицинских средств индивидуальной защиты (СИЗ) в контексте пандемии COVID-19.
Загрузите набор данных и рассмотрите пример:
>>> 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: идентификатор изображения.height: высота изображения.width: ширина изображения.objects: Словарь, содержащий метаданные ограничивающих прямоугольников для объектов на изображении:id: идентификатор аннотации.area: площадь ограничивающего прямоугольника.bbox: ограничивающая рамка объекта (в формате coco).category: Категория объекта, возможные значения включаютCoverall (0),Face_Shield (1),Gloves (2),Goggles (3)иMask (4).
Вы можете визуализировать bboxes на изображении, используя некоторые внутренние утилиты torch. Для этого вам нужно обратиться к функции 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(), чтобы применять преобразование на лету, что позволяет экономить место на диске. Случайность аугментации данных может привести к тому, что при повторном обращении к одному и тому же примеру будет возвращено другое изображение. Это особенно полезно при обучении модели в течение нескольких эпох.
>>> 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']]
... )
... )
Update on GitHubТеперь, когда вы знаете, как обработать набор данных для обнаружения объектов, узнайте как обучить модель обнаружения объектов и использовать её для инференса.