Datasets

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v4.8.4).

Join the Hugging Face community

and get access to the augmented documentation experience

Collaborate on models, datasets and Spaces

Faster examples with accelerated inference

Switch between documentation themes

to get started

Detecção de objetos

Modelos de detecção de objetos identificam algo em uma imagem, e conjuntos de dados de detecção de objetos são usados em aplicações como direção autônoma e detecção de perigos naturais, como incêndios florestais. Este guia mostrará como aplicar transformações a um conjunto de dados de detecção de objetos seguindo o tutorial da Albumentations.

Para executar esses exemplos, verifique se você tem as versões mais recentes de albumentations e cv2 instaladas:

pip install -U albumentations opencv-python

Neste exemplo, você usará o conjunto de dados cppe-5 para identificar equipamentos de proteção individual (EPI) médicos no contexto da pandemia de COVID-19.

Carregue o conjunto de dados e veja um exemplo:

>>> 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}

O conjunto de dados possui os seguintes campos:

image: objeto PIL.Image.Image contendo a imagem.
image_id: O ID da imagem.
height: A altura da imagem.
width: A largura da imagem.
objects: Um dicionário contendo metadados de bounding box para os objetos na imagem:
- id: O id da anotação.
- area: A área da caixa delimitadora.
- bbox: A caixa delimitadora do objeto (no formato coco).
- category: A categoria do objeto, com valores possíveis incluindo Coverall (0), Face_Shield (1), Gloves (2), Goggles (3) e Mask (4).

Você pode visualizar os bboxes na imagem usando alguns utilitários internos do torch. Para isso, você precisará referenciar o recurso ClassLabel associado aos IDs de categoria para que possa consultar os rótulos em string:

>>> 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,
...     )
... )

Com o albumentations, você pode aplicar transformações que afetam a imagem e, ao mesmo tempo, atualizam os bboxes de forma correspondente. Neste caso, a imagem é redimensionada para (480, 480), invertida horizontalmente e clareada.

>>> 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'],
... )

Agora, ao visualizar o resultado, a imagem deve estar invertida, mas as bboxes devem permanecer nas posições corretas.

>>> 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
...     )
... )

Crie uma função para aplicar a transformação a um lote de exemplos:

>>> 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}

Use a função set_transform() para aplicar a transformação em tempo de execução, o que consome menos espaço em disco. A aleatoriedade da ampliação de dados pode retornar uma imagem diferente se você acessar o mesmo exemplo duas vezes. Isso é especialmente útil ao treinar um modelo por várias épocas.

>>> ds['train'].set_transform(transforms)

Você pode verificar se a transformação funciona visualizando o 10º exemplo:

>>> 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']]
...     )
... )

Agora que você sabe como processar um conjunto de dados para detecção de objetos, aprenda como treinar um modelo de detecção de objetos e usá-lo para inferência.

Update on GitHub

←Segmentação semântica Carregar dados de vídeo→