Datasets documentation

Détection d'objets

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).
Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

Détection d'objets

Les modèles de détection d'objets identifient des éléments dans une image, et les jeux de données de détection d'objets sont utilisés pour des applications telles que la conduite autonome et la détection de dangers naturels comme les feux de forêt. Ce guide vous montre comment appliquer des transformations à un jeu de données de détection d'objets en suivant le tutorial de Albumentations.

Pour exécuter ces exemples, assurez-vous d'avoir des versions à jour de albumentations et cv2 installées :

pip install -U albumentations opencv-python

Dans cet exemple, vous utiliserez le jeu de données cppe-5 pour identifier les équipements de protection individuelle (EPI) médicaux dans le contexte de la pandémie de COVID-19.

Chargez le jeu de données et examinez un exemple :

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

Le jeu de données contient les champs suivants :

  • image : objet PIL.Image.Image contenant l'image.
  • image_id : L'identifiant de l'image.
  • height : La hauteur de l'image.
  • width : La largeur de l'image.
  • objects: Un dictionnaire contenant les métadonnées des boîtes englobantes des objets de l'image :
    • id : l'identifiant de l'annotation.
    • area : La surface de la boîte englobante.
    • bbox : la boîte englobante de l’objet (au format coco).
    • category : la catégorie de l’objet, avec des valeurs possibles incluant Coverall (0), Face_Shield (1), Gloves (2), Goggles (3) et Mask (4).

Vous pouvez visualiser les bboxes sur l'image à l'aide de quelques utilitaires internes de torch. Pour ce faire, vous devez référencer la fonctionnalité ClassLabel associée aux identifiants de catégories afin de pouvoir rechercher les étiquettes en chaîne :

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

Avec albumentations, vous pouvez appliquer des transformations qui modifient l'image tout en mettant à jour les bboxes en conséquence. Dans ce cas, l'image est redimensionnée à (480, 480), retournée horizontalement et éclaircie.

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

Maintenant, lorsque vous visualisez le résultat, l'image doit être retournée, mais les bboxes doivent toujours être aux bons endroits.

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

Créez une fonction pour appliquer la transformation à un lot d'exemples :

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

Utilisez la fonction set_transform() pour appliquer la transformation à la volée, ce qui réduit l’espace disque consommé. L’aléa de l’augmentation de données peut renvoyer une image différente si vous accédez deux fois au même exemple. C’est particulièrement utile lors de l’entraînement d’un modèle sur plusieurs époques.

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

Vous pouvez vérifier que la transformation fonctionne en visualisant le 10e exemple :

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

Maintenant que vous savez comment traiter un jeu de données pour la détection d'objets, apprenez comment entraîner un modèle de détection d'objets et l'utiliser pour l'inférence.

Update on GitHub