Detección de objetos

La detección de objetos es la tarea de visión por computadora que consiste en detectar instancias (como personas, edificios o automóviles) en una imagen. Los modelos de detección de objetos reciben una imagen como entrada y producen como salida las coordenadas de las cajas delimitadoras y las etiquetas asociadas de los objetos detectados. Una imagen puede contener múltiples objetos, cada uno con su propia caja delimitadora y una etiqueta (por ejemplo, puede tener un automóvil y un edificio), y cada objeto puede estar presente en diferentes partes de la imagen (por ejemplo, la imagen puede tener varios automóviles). Esta tarea se utiliza comúnmente en la conducción autónoma para detectar elementos como peatones, señales de tráfico y semáforos. Otras aplicaciones incluyen el conteo de objetos en imágenes, la búsqueda de imágenes y más.

En esta guía aprenderás a:

Ajusta RF-DETR en el conjunto de datos mobile-ui-design para detectar elementos de la interfaz de usuario en capturas de pantalla de aplicaciones móviles.
Usa tu modelo afinado para la inferencia.

Para ver todas las arquitecturas y puntos de control compatibles con esta tarea, recomendamos consultar la página de la tarea

Antes de comenzar, asegúrate de tener instaladas todas las bibliotecas necesarias:

pip install -q datasets transformers accelerate timm trackio torchmetrics pycocotools

Utilizarás 🤗 Datasets para cargar un conjunto de datos desde el Hugging Face Hub y 🤗 Transformers para entrenar tu modelo.

Te animamos a compartir tu modelo con la comunidad. Inicia sesión en tu cuenta de Hugging Face para subirlo al Hub. Cuando se te solicite, introduce tu token para iniciar sesión:

>>> from huggingface_hub import notebook_login

>>> notebook_login()

Defina las constantes globales, a saber, el nombre del modelo y el tamaño de la imagen. Este tutorial utiliza RF-DETR, pero puede seleccionar cualquier modelo de detección de objetos en Transformers.

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

Cargar el conjunto de datos mobile-ui-design

El conjunto de datos mobile-ui-design contiene capturas de pantalla de aplicaciones móviles con anotaciones para detectar elementos de la interfaz de usuario, como texto, imágenes, rectángulos y grupos.

Comience cargando el conjunto de datos y extrayendo las etiquetas de categoría. El conjunto de datos ya tiene divisiones.

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

El conjunto de datos utiliza nombres de categorías de cadena y cajas delimitadoras en formato COCO (x, y, w, h). Convierta las categorías en ids enteros, calcule las áreas y filtre las cajas delimitadoras degeneradas antes de entrenar:

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

Preprocesar los datos

AutoImageProcessor se encarga de procesar los datos de imagen para crear pixel_values, pixel_mask y labels con los que el modelo puede entrenar. El procesador de imágenes gestiona el redimensionado, el relleno y la normalización. Además, puedes añadir opcionalmente aumentos de datos aleatorios (ver a continuación) para mejorar la generalización.

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

>>> image_processor = AutoImageProcessor.from_pretrained(MODEL_NAME)

El image_processor espera anotaciones en el formato COCO: {'image_id': int, 'annotations': list[Dict]}. Formatea las anotaciones de cada ejemplo y deja que el procesador se encargue del resto:

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

Aumento de datos

La transformación anterior solo redimensiona y normaliza las imágenes. Las aumentaciones aleatorias aplicadas al conjunto de entrenamiento suelen mejorar la generalización, mientras que el conjunto de validación debe permanecer sin aumentaciones para que la evaluación sea determinista. Una opción común es Albumentations, que aumenta la imagen y sus cajas delimitadoras de forma conjunta. Defina un pipeline con bbox_params para que las cajas se transformen de manera coherente con la imagen, y luego recalcule las áreas a partir de las cajas aumentadas:

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

Aplica la transformación de aumento solo a la división de entrenamiento y conserva la transform_fn simple para la validación:

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

Cree una collate_fn personalizada para agrupar imágenes en lotes:

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

Preparando la función para calcular mAP

Los modelos de detección de objetos se evalúan comúnmente con un conjunto de métricas estilo COCO. Vamos a usar torchmetrics para calcular las métricas mAP (precisión media) y mAR (recall medio) y las envolveremos en la función compute_metrics para utilizarlas en Trainer durante la evaluación.

El formato intermedio de las cajas utilizado para el entrenamiento es YOLO (normalizado), pero calcularemos las métricas para las cajas en el formato Pascal VOC (absoluto) para manejar correctamente las áreas de las cajas. Definiamos una función que convierte las cajas delimitadoras al formato 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

Luego, en la función compute_metrics recopilamos las cajas delimitadoras, puntuaciones y etiquetas predicted y target a partir de los resultados del bucle de evaluación y las pasamos a la función de puntuación.

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

Entrenamiento del modelo de detección

Ya has realizado la mayor parte del trabajo en las secciones anteriores, así que ahora estás listo para entrenar tu modelo. Las imágenes de este conjunto de datos siguen siendo bastante grandes, incluso después de redimensionarlas. Esto significa que el ajuste fino de este modelo requerirá al menos una GPU.

El entrenamiento implica los siguientes pasos:

Cargue el modelo con AutoModelForObjectDetection utilizando el mismo punto de control que en el preprocesamiento.
Define los hiperparámetros de entrenamiento en TrainingArguments.
Pasa los argumentos de entrenamiento a Trainer junto con el modelo, el conjunto de datos, el procesador de imágenes y el data collator.
Llama a train() para ajustar tu modelo.

Al cargar el modelo desde el mismo punto de control que utilizaste para el preprocesamiento, recuerda pasar los mapas label2id y id2label que creaste anteriormente a partir de los metadatos del conjunto de datos. Además, especificamos ignore_mismatched_sizes=True para reemplazar la cabeza de clasificación existente por una nueva.

>>> from transformers import AutoModelForObjectDetection

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

En TrainingArguments, use output_dir para especificar dónde guardar su modelo y, a continuación, configure los hiperparámetros según lo considere oportuno. Con num_train_epochs=5, el entrenamiento tardará aproximadamente 35 minutos en una GPU A100; aumente el número de épocas para obtener mejores resultados.

Notas importantes:

No elimines las columnas no utilizadas, ya que esto eliminaría la columna de imagen. Sin la columna de imagen, no puedes crear pixel_values. Por esta razón, establece remove_unused_columns en False.
Establezca eval_do_concat_batches=False para obtener resultados de evaluación correctos. Las imágenes tienen un número diferente de cajas objetivo; si se concatenan los lotes, no será posible determinar a qué imagen pertenecen las cajas.

Si deseas compartir tu modelo publicándolo en el Hub, establece push_to_hub en True (debes iniciar sesión en Hugging Face para subir tu modelo).

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

Por último, reúne todo y llama a 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, Época 5/5]

Época	Pérdida de entrenamiento	Pérdida de validación	Map	Map 50	Map 75	Map Small	Map Medium	Map Large	1 de marzo	10 de marzo	Mar 100	Mar Small	Mar Medio	Mar Grande	Grupo Map	Mar 100 Group	Mapa de imagen	Mar 100 Imagen	Map Rectangle	Mar 100 Rectángulo	Mapa de texto	Mar 100 Texto
1	Sin registro	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	Sin registro	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	Sin registro	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	Sin registro	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 has establecido push_to_hub en True en training_args, los puntos de control del entrenamiento se envían al Hugging Face Hub. Al completar el entrenamiento, envía también el modelo final al Hub llamando al método push_to_hub().

>>> trainer.push_to_hub()

Evaluar

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

Estos resultados pueden mejorarse aún más aumentando el número de épocas o ajustando otros hiperparámetros en TrainingArguments. ¡Pruébalo!

Inferencia

Ahora que has afinado un modelo, lo has evaluado y lo has subido al Hugging Face Hub, puedes usarlo para inferencia.

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

Cargue el modelo y el procesador de imágenes desde el Hugging Face Hub (omitir para usar los ya entrenados en esta sesión):

>>> model_repo = "merve/rf_detr_finetuned_mobile_ui"

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

Y detectar cajas delimitadoras:

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

Vamos a graficar el resultado:

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