Datasets documentation
Detección de objetos
Detección de objetos
Los modelos de detección de objetos identifican algo en una imagen, y los conjuntos de datos de detección de objetos se utilizan para aplicaciones como la conducción autónoma y la detección de peligros naturales, como los incendios forestales. Esta guía le mostrará cómo aplicar transformaciones a un conjunto de datos de detección de objetos siguiendo el tutorial de Albumentations.
Para ejecutar estos ejemplos, asegúrate de tener instaladas las versiones más recientes de albumentations y cv2:
pip install -U albumentations opencv-python
En este ejemplo, utilizarás el conjunto de datos cppe-5 para identificar el equipo de protección personal (EPP) médico en el contexto de la pandemia de COVID-19.
Cargue el conjunto de datos y eche un vistazo a un ejemplo:
>>> 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}El conjunto de datos tiene los siguientes campos:
image: objeto PIL.Image.Image que contiene la imagen.image_id: El ID de la imagen.height: La altura de la imagen.width: El ancho de la imagen.objects: Un diccionario que contiene los metadatos de las cajas delimitadoras de los objetos en la imagen:id: El identificador de la anotación.area: El área del cuadro delimitador.bbox: La caja delimitadora del objeto (en el formato coco).category: La categoría del objeto, con valores posibles que incluyenCoverall (0),Face_Shield (1),Gloves (2),Goggles (3)yMask (4).
Puedes visualizar los bboxes en la imagen utilizando algunas utilidades internas de torch. Para ello, necesitarás hacer referencia a la función ClassLabel asociada con los IDs de categoría para poder buscar las etiquetas de texto:
>>> 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,
... )
... )
Con albumentations, puedes aplicar transformaciones que afectarán a la imagen y, al mismo tiempo, actualizarán los bboxes de forma correspondiente. En este caso, la imagen se redimensiona a (480, 480), se invierte horizontalmente y se aclara.
>>> 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'],
... )Ahora, al visualizar el resultado, la imagen debería estar invertida, pero los bboxes deberían seguir en las posiciones correctas.
>>> 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
... )
... )
Cree una función para aplicar la transformación a un lote de ejemplos:
>>> 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 la función set_transform() para aplicar la transformación al vuelo, lo que consume menos espacio en disco. La aleatoriedad de la aumentación de datos puede devolver una imagen diferente si accede al mismo ejemplo dos veces. Es especialmente útil al entrenar un modelo durante varias épocas.
>>> ds['train'].set_transform(transforms)Puede verificar que la transformación funciona visualizando el ejemplo 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 GitHubAhora que sabes cómo procesar un conjunto de datos para la detección de objetos, aprende cómo entrenar un modelo de detección de objetos y usarlo para la inferencia.