Détection d'objets

La détection d'objets est la tâche de vision par ordinateur qui consiste à détecter des instances (tels que des humains, des bâtiments ou des voitures) dans une image. Les modèles de détection d'objets reçoivent une image en entrée et produisent en sortie les coordonnées des boîtes englobantes et les étiquettes associées aux objets détectés. Une image peut contenir plusieurs objets, chacun avec sa propre boîte englobante et son étiquette (par exemple, elle peut contenir une voiture et un bâtiment), et chaque objet peut être présent dans différentes parties de l'image (par exemple, l'image peut contenir plusieurs voitures). Cette tâche est couramment utilisée dans la conduite autonome pour détecter des éléments tels que des piétons, des panneaux de signalisation et des feux de circulation. D'autres applications incluent le comptage d'objets dans les images, la recherche d'images, et bien plus encore.

Dans ce guide, vous apprendrez à :

Effectuer le fine-tuning de RF-DETR sur le jeu de données mobile-ui-design pour détecter les éléments d'interface dans les captures d'écran d'applications mobiles.
Utilisez votre modèle affiné pour l'inférence.

Pour consulter toutes les architectures et points de contrôle compatibles avec cette tâche, nous vous recommandons de consulter la page de la tâche

Avant de commencer, assurez-vous d'avoir installé toutes les bibliothèques nécessaires :

pip install -q datasets transformers accelerate timm trackio torchmetrics pycocotools

Vous utiliserez 🤗 Datasets pour charger un jeu de données depuis le Hugging Face Hub et 🤗 Transformers pour entraîner votre modèle.

Nous vous encourageons à partager votre modèle avec la communauté. Connectez-vous à votre compte Hugging Face pour le téléverser sur le Hub. Lorsqu’on vous le demande, saisissez votre jeton pour vous connecter :

>>> from huggingface_hub import notebook_login

>>> notebook_login()

Définissez les constantes globales, à savoir le nom du modèle et la taille de l'image. Ce tutoriel utilise RF-DETR, mais vous pouvez sélectionner n'importe quel modèle de détection d'objets dans Transformers.

>>> MODEL_NAME = "Roboflow/rf-detr-medium"

Charger le jeu de données mobile-ui-design

Le jeu de données mobile-ui-design contient des captures d'écran d'applications mobiles avec des annotations pour la détection d'éléments d'interface utilisateur tels que le texte, les images, les rectangles et les groupes.

Commencez par charger le jeu de données et extraire les étiquettes de catégories. Le jeu de données comporte déjà des sous-ensembles.

>>> from datasets import load_dataset

>>> ds = load_dataset("merve/mobile-ui-design")

>>> CATEGORIES = sorted(set(
...     cat for split in ds.values() for example in split for cat in example["objects"]["category"]
... ))
>>> label2id = {label: i for i, label in enumerate(CATEGORIES)}
>>> id2label = {i: label for label, i in label2id.items()}
>>> print(f"Categories ({len(CATEGORIES)}): {CATEGORIES}")
Categories (4): ['group', 'image', 'rectangle', 'text']

Le jeu de données utilise des noms de catégories sous forme de chaînes et des boîtes englobantes au format COCO (x, y, w, h). Convertissez les catégories en identifiants entiers, calculez les aires et filtrez les boîtes englobantes dégénérées avant l'entraînement :

>>> def prepare_example(example, idx):
...     objects = example["objects"]
...     in_bboxes = objects["bbox"]
...     categories = objects["category"]
...     img_w, img_h = example["width"], example["height"]
...     bboxes, cats, areas, ids = [], [], [], []
...     for i, (bbox, cat) in enumerate(zip(in_bboxes, categories)):
...         x, y, w, h = bbox
...         if w <= 0 or h <= 0:
...             continue
...         x = max(0.0, min(x, img_w))
...         y = max(0.0, min(y, img_h))
...         w = min(w, img_w - x)
...         h = min(h, img_h - y)
...         if w <= 0 or h <= 0:
...             continue
...         bboxes.append([x, y, w, h])
...         cats.append(label2id[cat])
...         areas.append(w * h)
...         ids.append(i)
...     return {
...         "image_id": idx, "image": example["image"],
...         "width": example["width"], "height": example["height"],
...         "objects": {"id": ids, "bbox": bboxes, "category": cats, "area": areas},
...     }

>>> ds_prepared = ds["train"].map(prepare_example, with_indices=True, remove_columns=ds["train"].column_names)
>>> ds_prepared = ds_prepared.filter(lambda x: len(x["objects"]["bbox"]) > 0)

>>> split = ds_prepared.train_test_split(test_size=0.15, seed=1337)
>>> train_ds = split["train"]
>>> val_ds = split["test"]
>>> print(f"Train: {len(train_ds)}, Validation: {len(val_ds)}")
Train: 6669, Validation: 1177

Prétraiter les données

AutoImageProcessor se charge de traiter les données d’images pour générer pixel_values, pixel_mask et labels, que le modèle peut utiliser pour l’entraînement. Le processeur d’images gère le redimensionnement, le bourrage (padding) et la normalisation. En plus, vous pouvez ajouter des augmentations de données aléatoires (voir ci-dessous) pour améliorer la généralisation.

>>> import numpy as np
>>> from functools import partial
>>> from transformers import AutoImageProcessor

>>> image_processor = AutoImageProcessor.from_pretrained(MODEL_NAME)

Le image_processor s’attend à des annotations au format COCO : {'image_id': int, 'annotations': list[Dict]}. Formatez les annotations de chaque exemple et laissez le processeur s’occuper du reste :

>>> def format_image_annotations_as_coco(image_id, categories, areas, bboxes):
...     annotations = []
...     for category, area, bbox in zip(categories, areas, bboxes):
...         annotations.append({
...             "image_id": image_id,
...             "category_id": category,
...             "iscrowd": 0,
...             "area": area,
...             "bbox": list(bbox),
...         })
...     return {"image_id": image_id, "annotations": annotations}

>>> def transform_batch(examples, image_processor):
...     images = []
...     annotations = []
...     for image_id, image, objects in zip(examples["image_id"], examples["image"], examples["objects"]):
...         images.append(np.array(image.convert("RGB")))
...         formatted = format_image_annotations_as_coco(
...             image_id, objects["category"], objects["area"], objects["bbox"]
...         )
...         annotations.append(formatted)
...     result = image_processor(images=images, annotations=annotations, return_tensors="pt")
...     result.pop("pixel_mask", None)
...     return result

>>> transform_fn = partial(transform_batch, image_processor=image_processor)
>>> train_ds = train_ds.with_transform(transform_fn)
>>> val_ds = val_ds.with_transform(transform_fn)

Augmentation de données

La transformation ci-dessus ne redimensionne et normalise que les images. Les augmentations aléatoires appliquées à la partition training améliorent généralement la généralisation, tandis que la partition de validation doit rester sans augmentation afin que l’évaluation reste déterministe. Un choix courant est Albumentations, qui augmente l’image et ses boîtes englobantes de manière cohérente. Définissez un pipeline avec bbox_params pour que les boîtes soient transformées de façon cohérente avec l’image, puis recalculez les surfaces à partir des boîtes augmentées :

>>> import albumentations as A

>>> train_augment = A.Compose(
...     [
...         A.Perspective(p=0.1),
...         A.HorizontalFlip(p=0.5),
...         A.RandomBrightnessContrast(p=0.5),
...         A.HueSaturationValue(p=0.1),
...     ],
...     bbox_params=A.BboxParams(format="coco", label_fields=["category"], clip=True, min_area=25),
... )

>>> def augment_and_transform_batch(examples, image_processor, transform):
...     images = []
...     annotations = []
...     for image_id, image, objects in zip(examples["image_id"], examples["image"], examples["objects"]):
...         image = np.array(image.convert("RGB"))
...         output = transform(image=image, bboxes=objects["bbox"], category=objects["category"])
...         images.append(output["image"])
...         areas = [w * h for (_, _, w, h) in output["bboxes"]]
...         formatted = format_image_annotations_as_coco(
...             image_id, output["category"], areas, output["bboxes"]
...         )
...         annotations.append(formatted)
...     result = image_processor(images=images, annotations=annotations, return_tensors="pt")
...     result.pop("pixel_mask", None)
...     return result

Appliquez la transformation d’augmentation uniquement à la partition d’entraînement, et conservez la transform_fn simple pour la validation :

>>> train_augment_fn = partial(augment_and_transform_batch, image_processor=image_processor, transform=train_augment)
>>> train_ds = train_ds.with_transform(train_augment_fn)
>>> val_ds = val_ds.with_transform(transform_fn)

Créez une collate_fn personnalisée pour regrouper les images en lots :

>>> import torch

>>> def collate_fn(batch):
...     data = {}
...     data["pixel_values"] = torch.stack([x["pixel_values"] for x in batch])
...     data["labels"] = [x["labels"] for x in batch]
...     if "pixel_mask" in batch[0]:
...         data["pixel_mask"] = torch.stack([x["pixel_mask"] for x in batch])
...     return data

Préparation de la fonction de calcul de mAP

Les modèles de détection d’objets sont couramment évalués à l’aide d’un ensemble de métriques de type COCO. Nous allons utiliser torchmetrics pour calculer les métriques mAP (mean average precision) et mAR (mean average recall), et nous les encapsulerons dans la fonction compute_metrics afin de les utiliser dans Trainer pour l’évaluation.

Le format intermédiaire des boîtes utilisé pour l’entraînement est YOLO (normalisé), mais nous calculerons les métriques pour les boîtes au format Pascal VOC (absolu) afin de traiter correctement les surfaces des boîtes. Définissons une fonction qui convertit les boîtes englobantes au format Pascal VOC :

>>> from transformers.image_transforms import center_to_corners_format

>>> def convert_bbox_yolo_to_pascal(boxes, image_size):
...     """
...     Convert bounding boxes from YOLO format (x_center, y_center, width, height) in range [0, 1]
...     to Pascal VOC format (x_min, y_min, x_max, y_max) in absolute coordinates.

...     Args:
...         boxes (torch.Tensor): Bounding boxes in YOLO format
...         image_size (tuple[int, int]): Image size in format (height, width)

...     Returns:
...         torch.Tensor: Bounding boxes in Pascal VOC format (x_min, y_min, x_max, y_max)
...     """
...     # convert center to corners format
...     boxes = center_to_corners_format(boxes)

...     # convert to absolute coordinates
...     height, width = image_size
...     boxes = boxes * torch.tensor([[width, height, width, height]])

...     return boxes

Ensuite, dans la fonction compute_metrics, nous collectons les boîtes englobantes predicted et target, ainsi que les scores et les étiquettes, à partir des résultats de la boucle d'évaluation, puis nous les transmettons à la fonction de scoring.

>>> import numpy as np
>>> from dataclasses import dataclass
>>> from torchmetrics.detection.mean_ap import MeanAveragePrecision


>>> @dataclass
>>> class ModelOutput:
...     logits: torch.Tensor
...     pred_boxes: torch.Tensor


>>> def _get_orig_size(image_target):
...     """Robust orig_size extraction - Trainer serialization can truncate to 1 element."""
...     orig = np.atleast_1d(np.asarray(image_target["orig_size"])).flatten()
...     if len(orig) >= 2:
...         return (int(orig[0]), int(orig[1]))
...     return (int(orig[0]), int(orig[0]))

>>> @torch.no_grad()
>>> def compute_metrics(evaluation_results, image_processor, threshold=0.0, id2label=None):
...     predictions, targets = evaluation_results.predictions, evaluation_results.label_ids
...     image_sizes = []
...     post_processed_targets = []
...     post_processed_predictions = []
...
...     for batch in targets:
...         batch_sizes = []
...         for image_target in batch:
...             h, w = _get_orig_size(image_target)
...             batch_sizes.append([h, w])
...             boxes = torch.tensor(image_target["boxes"])
...             boxes = convert_bbox_yolo_to_pascal(boxes, (h, w))
...             labels = torch.tensor(image_target["class_labels"])
...             post_processed_targets.append({"boxes": boxes, "labels": labels})
...         image_sizes.append(torch.tensor(batch_sizes))
...
...     for batch, target_sizes in zip(predictions, image_sizes):
...         batch_logits, batch_boxes = batch[1], batch[2]
...         output = ModelOutput(logits=torch.tensor(batch_logits), pred_boxes=torch.tensor(batch_boxes))
...         post_processed_output = image_processor.post_process_object_detection(
...             output, threshold=threshold, target_sizes=target_sizes
...         )
...         post_processed_predictions.extend(post_processed_output)
...
...     metric = MeanAveragePrecision(box_format="xyxy", class_metrics=True)
...     metric.update(post_processed_predictions, post_processed_targets)
...     metrics = metric.compute()
...
...     classes = metrics.pop("classes")
...     map_per_class = metrics.pop("map_per_class")
...     mar_100_per_class = metrics.pop("mar_100_per_class")
...     for class_id, class_map, class_mar in zip(classes, map_per_class, mar_100_per_class):
...         class_name = id2label[class_id.item()] if id2label is not None else class_id.item()
...         metrics[f"map_{class_name}"] = class_map
...         metrics[f"mar_100_{class_name}"] = class_mar
...
...     metrics = {k: round(v.item(), 4) for k, v in metrics.items()}
...     return metrics

>>> eval_compute_metrics_fn = partial(
...     compute_metrics, image_processor=image_processor, id2label=id2label, threshold=0.0
... )

Entraînement du modèle de détection

Vous avez effectué l'essentiel du travail dans les sections précédentes, vous êtes donc prêt à entraîner votre modèle ! Les images de ce jeu de données restent assez volumineuses, même après redimensionnement. Cela signifie que l'ajustement fin de ce modèle nécessitera au moins une GPU.

L'entraînement comprend les étapes suivantes :

Chargez le modèle avec AutoModelForObjectDetection en utilisant le même point de contrôle que lors de la prétraitement.
Définissez vos hyperparamètres d’entraînement dans TrainingArguments.
Transmettez les arguments d'entraînement à Trainer ainsi que le modèle, le jeu de données, le processeur d'images et le data collator.
Appel de train() pour l’ajustement fin de votre modèle.

Lors du chargement du modèle depuis le même point de contrôle que celui utilisé pour le prétraitement, n’oubliez pas de transmettre les correspondances label2id et id2label que vous avez créées plus tôt à partir des métadonnées du jeu de données. De plus, nous spécifions ignore_mismatched_sizes=True pour remplacer la tête de classification existante par une nouvelle.

>>> from transformers import AutoModelForObjectDetection

>>> model = AutoModelForObjectDetection.from_pretrained(
...     MODEL_NAME,
...     id2label=id2label,
...     label2id=label2id,
...     ignore_mismatched_sizes=True,
... )

Dans TrainingArguments, utilisez output_dir pour indiquer l’emplacement de sauvegarde de votre modèle, puis configurez les hyperparamètres selon vos besoins. Pour num_train_epochs=5, l’entraînement prend environ 35 minutes sur un GPU A100 ; augmentez le nombre d’époques pour obtenir de meilleurs résultats.

Notes importantes :

Ne supprimez pas les colonnes inutilisées, car cela supprimerait la colonne d’image. Sans la colonne d’image, vous ne pouvez pas créer pixel_values. Pour cette raison, définissez remove_unused_columns sur False.
Définissez eval_do_concat_batches=False pour obtenir des résultats d’évaluation corrects. Les images ont un nombre différent de boîtes cibles ; si les lots sont concaténés, il ne sera pas possible de déterminer à quelle image appartient chaque boîte.

Si vous souhaitez partager votre modèle en le publiant sur le Hub, définissez push_to_hub sur True (vous devez être connecté à Hugging Face pour téléverser votre modèle).

>>> from transformers import TrainingArguments

>>> training_args = TrainingArguments(
...     output_dir="rf_detr_finetuned_mobile_ui",
...     num_train_epochs=5,
...     bf16=True,
...     per_device_train_batch_size=8,
...     dataloader_num_workers=4,
...     learning_rate=5e-5,
...     lr_scheduler_type="cosine",
...     weight_decay=1e-4,
...     max_grad_norm=0.01,
...     metric_for_best_model="eval_map",
...     greater_is_better=True,
...     load_best_model_at_end=True,
...     eval_strategy="epoch",
...     save_strategy="epoch",
...     save_total_limit=2,
...     remove_unused_columns=False,
...     report_to="trackio",
...     run_name="mobile-ui-detection",
...     eval_do_concat_batches=False,
...     push_to_hub=True,
... )

Enfin, rassemblez le tout et appelez train() :

>>> from transformers import Trainer

>>> trainer = Trainer(
...     model=model,
...     args=training_args,
...     train_dataset=train_ds,
...     eval_dataset=val_ds,
...     processing_class=image_processor,
...     data_collator=collate_fn,
...     compute_metrics=eval_compute_metrics_fn,
... )

>>> trainer.train()

[2085/2085 38:39, Époque 5/5]

Époque	Perte d'entraînement	Perte de validation	Carte	Map 50	Carte 75	Cartographie de petite taille	Carte Medium	Carte Large	1er mars	10 mars	Mar 100	Mar Small	Mar Moyen	Mar Large	Groupe Map	Groupe Mar 100	Carte d'image	Mar 100 Image	Carte Rectangle	Mar 100 Rectangle	Carte de texte	Mar 100 Texte
1	Aucun journal	9.9234	0.1303	0.2236	0.1478	0.0909	0.2030	0.2524	0.0421	0.2520	0.4683	0.3113	0.5607	0.6782	0.1244	0.5122	0.0958	0.5035	0.1285	0.4328	0.1725	0.4413
2	Aucun journal	9.8472	0.1893	0.3017	0.2124	0.1347	0.2789	0.3038	0.0549	0.2961	0.5140	0.3433	0.5941	0.7406	0.1305	0.5423	0.1979	0.5578	0.1964	0.4648	0.2324	0.4437
3	Aucun journal	9.6401	0.2275	0.3547	0.2657	0.1698	0.3336	0.3892	0.0611	0.3204	0.5270	0.3625	0.6143	0.7496	0.1602	0.5684	0.2617	0.5763	0.2249	0.4684	0.2631	0.4692
4	Aucun journal	9.5770	0.2733	0.4068	0.3133	0.2100	0.3867	0.4343	0.0668	0.3456	0.5593	0.3875	0.6393	0.7725	0.2013	0.5941	0.3158	0.6065	0.2733	0.4998	0.3028	0.4756
5	10.3700	11.0500	0.2827	0.4193	0.2913	0.2021	0.2814	0.3763	0.0609	0.3403	0.5668	0.4138	0.5669	0.7317	0.2092	0.5979	0.3334	0.6295	0.2793	0.5245	0.3089	0.5151

Si vous avez défini push_to_hub sur True dans training_args, les points de contrôle d’entraînement sont poussés vers le Hugging Face Hub. Une fois l’entraînement terminé, poussez également le modèle final vers le Hub en appelant la méthode push_to_hub().

>>> trainer.push_to_hub()

Évaluer

>>> from pprint import pprint

>>> metrics = trainer.evaluate(eval_dataset=val_ds, metric_key_prefix="test")
>>> pprint(metrics)
{'test_loss': 11.05,
 'test_map': 0.2827,
 'test_map_50': 0.4193,
 'test_map_75': 0.2913,
 'test_map_group': 0.2092,
 'test_map_image': 0.3334,
 'test_map_large': 0.3763,
 'test_map_medium': 0.2814,
 'test_map_rectangle': 0.2793,
 'test_map_small': 0.2021,
 'test_map_text': 0.3089,
 'test_mar_1': 0.0609,
 'test_mar_10': 0.3403,
 'test_mar_100': 0.5668,
 'test_mar_100_group': 0.5979,
 'test_mar_100_image': 0.6295,
 'test_mar_100_rectangle': 0.5245,
 'test_mar_100_text': 0.5151,
 'test_mar_large': 0.7317,
 'test_mar_medium': 0.5669,
 'test_mar_small': 0.4138}

Ces résultats peuvent être encore améliorés en augmentant le nombre d'époques ou en ajustant d'autres hyperparamètres dans TrainingArguments. À vous de jouer !

Inférence

Maintenant que vous avez affiné un modèle, l'avez évalué et l'avez téléversé sur le Hugging Face Hub, vous pouvez l'utiliser pour l'inférence.

>>> import torch
>>> from PIL import Image, ImageDraw
>>> from transformers import AutoImageProcessor, AutoModelForObjectDetection
>>> from datasets import load_dataset

>>> ds = load_dataset("merve/mobile-ui-design", split="train")
>>> image = ds[5]["image"].convert("RGB")

Charger le modèle et le processeur d’images depuis le Hugging Face Hub (à passer si vous avez déjà entraîné un modèle au cours de cette session) :

>>> model_repo = "merve/rf_detr_finetuned_mobile_ui"

>>> image_processor = AutoImageProcessor.from_pretrained(model_repo)
>>> model = AutoModelForObjectDetection.from_pretrained(model_repo)
>>> model.eval()

Et détecter les boîtes englobantes :

>>> with torch.no_grad():
...     inputs = image_processor(images=[image], return_tensors="pt")
...     outputs = model(**inputs)
...     target_sizes = torch.tensor([[image.size[1], image.size[0]]])
...     results = image_processor.post_process_object_detection(outputs, threshold=0.5, target_sizes=target_sizes)[0]

>>> for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
...     box = [round(i, 2) for i in box.tolist()]
...     print(
...         f"Detected {model.config.id2label[label.item()]} with confidence "
...         f"{round(score.item(), 3)} at location {box}"
...     )
Detected text with confidence 0.727 at location [324.02, 340.55, 339.52, 359.12]
Detected rectangle with confidence 0.717 at location [39.97, 705.14, 335.93, 753.54]
Detected text with confidence 0.702 at location [199.94, 473.66, 213.41, 490.6]
Detected text with confidence 0.678 at location [153.14, 474.81, 165.33, 491.0]
Detected text with confidence 0.675 at location [262.67, 718.28, 281.44, 740.81]
Detected rectangle with confidence 0.655 at location [143.57, 242.51, 214.32, 274.26]
Detected text with confidence 0.653 at location [298.68, 637.77, 345.68, 656.26]

Tracer le résultat :

>>> draw = ImageDraw.Draw(image)

>>> colors = {"group": "blue", "image": "green", "rectangle": "red", "text": "orange"}
>>> for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
...     box = [round(i, 2) for i in box.tolist()]
...     x, y, x2, y2 = tuple(box)
...     label_name = model.config.id2label[label.item()]
...     color = colors.get(label_name, "red")
...     draw.rectangle((x, y, x2, y2), outline=color, width=2)
...     draw.text((x, y), f"{label_name} {score:.2f}", fill=color)

>>> image

Object detection result on a cart screen

Object detection result on a followers screen

Update on GitHub