Detecção de objetos

A detecção de objetos é a tarefa de visão computacional que consiste em detectar instâncias (como pessoas, prédios ou carros) em uma imagem. Os modelos de detecção de objetos recebem uma imagem como entrada e produzem como saída as coordenadas das bounding boxes e os rótulos associados aos objetos detectados. Uma imagem pode conter múltiplos objetos, cada um com sua própria bounding box e um rótulo (por exemplo, pode haver um carro e um prédio), e cada objeto pode estar presente em diferentes partes da imagem (por exemplo, a imagem pode ter vários carros). Essa tarefa é comumente usada em direção autônoma para detectar elementos como pedestres, placas de trânsito e semáforos. Outras aplicações incluem contagem de objetos em imagens, busca de imagens e mais.

Neste guia, você aprenderá como:

Ajuste fino de RF-DETR no conjunto de dados mobile-ui-design para detectar elementos de interface em capturas de tela de aplicativos móveis.
Use seu modelo ajustado para inferência.

Para ver todas as arquiteturas e checkpoints compatíveis com esta tarefa, recomendamos consultar a página da tarefa

Antes de começar, verifique se você tem todas as bibliotecas necessárias instaladas:

pip install -q datasets transformers accelerate timm trackio torchmetrics pycocotools

Você usará o 🤗 Datasets para carregar um dataset do Hugging Face Hub e o 🤗 Transformers para treinar seu modelo.

Incentivamos você a compartilhar seu modelo com a comunidade. Faça login na sua conta do Hugging Face para enviá-lo ao Hub. Quando solicitado, insira seu token para fazer login:

>>> from huggingface_hub import notebook_login

>>> notebook_login()

Defina as constantes globais, a saber, o nome do modelo e o tamanho da imagem. Este tutorial usa RF-DETR, mas você pode selecionar qualquer modelo de detecção de objetos em Transformers.

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

Carregue o conjunto de dados mobile-ui-design

O conjunto de dados mobile-ui-design contém capturas de tela de aplicativos móveis com anotações para detectar elementos de interface, como texto, imagens, retângulos e grupos.

Comece carregando o conjunto de dados e extraindo os rótulos de categoria. O conjunto de dados já possui divisões.

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

O conjunto de dados usa nomes de categorias em string e caixas delimitadoras no formato COCO (x, y, w, h). Converta as categorias em ids inteiros, calcule as áreas e filtre as caixas delimitadoras degeneradas antes do treinamento:

>>> 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é-processar os dados

AutoImageProcessor cuida do processamento dos dados de imagem para criar pixel_values, pixel_mask e labels com os quais o modelo pode ser treinado. O processador de imagem trata redimensionamento, preenchimento e normalização. Além disso, você pode adicionar opcionalmente aumentações de dados aleatórias (veja abaixo) para melhorar a generalização.

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

>>> image_processor = AutoImageProcessor.from_pretrained(MODEL_NAME)

O image_processor espera anotações no formato COCO: {'image_id': int, 'annotations': list[Dict]}. Formate as anotações de cada exemplo e deixe o processador cuidar do 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 dados

A transformação acima apenas redimensiona e normaliza as imagens. Aumentações aleatórias aplicadas ao split de treinamento geralmente melhoram a generalização, enquanto o split de validação deve permanecer sem aumentações para que a avaliação permaneça determinística. Uma escolha comum é o Albumentations, que aumenta a imagem e suas caixas delimitadoras juntas. Defina um pipeline com bbox_params para que as caixas sejam transformadas de forma consistente com a imagem e, em seguida, recalcule as áreas a partir das caixas 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

Aplique a transformação de aumento apenas ao conjunto de treinamento e mantenha o transform_fn simples para a validação:

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

Crie um collate_fn personalizado para agrupar as imagens em 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 a função para calcular o mAP

Modelos de detecção de objetos são comumente avaliados com um conjunto de métricas no estilo COCO. Vamos usar o torchmetrics para calcular as métricas mAP (precisão média) e mAR (recall médio) e encapsulá-las na função compute_metrics para uso no Trainer durante a avaliação.

O formato intermediário das caixas usado para o treinamento é YOLO (normalizado), mas calcularemos as métricas para caixas no formato Pascal VOC (absoluto) a fim de lidar corretamente com as áreas das caixas. Vamos definir uma função que converte caixas delimitadoras para o 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

Em seguida, na função compute_metrics, coletamos as caixas delimitadoras predicted e target, as pontuações e os rótulos dos resultados do loop de avaliação e os passamos para a função de pontuação.

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

Treinando o modelo de detecção

Você já fez a maior parte do trabalho pesado nas seções anteriores, então agora está pronto para treinar seu modelo! As imagens neste conjunto de dados ainda são bastante grandes, mesmo após o redimensionamento. Isso significa que o ajuste fino deste modelo exigirá pelo menos uma GPU.

O treinamento envolve as seguintes etapas:

Carregue o modelo com AutoModelForObjectDetection usando o mesmo checkpoint do pré-processamento.
Defina os hiperparâmetros de treinamento em TrainingArguments.
Passe os argumentos de treinamento para Trainer junto com o modelo, o conjunto de dados, o processador de imagem e o data collator.
Chame train() para fazer o ajuste fino do seu modelo.

Ao carregar o modelo a partir do mesmo checkpoint usado no pré-processamento, lembre-se de passar os mapas label2id e id2label que você criou anteriormente a partir dos metadados do conjunto de dados. Além disso, especificamos ignore_mismatched_sizes=True para substituir a cabeça de classificação existente por uma nova.

>>> from transformers import AutoModelForObjectDetection

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

No TrainingArguments, use output_dir para especificar onde salvar seu modelo e, em seguida, configure os hiperparâmetros conforme desejar. Para num_train_epochs=5, o treinamento levará cerca de 35 minutos em uma GPU A100; aumente o número de épocas para obter melhores resultados.

Observações importantes:

Não remova colunas não utilizadas, pois isso removerá a coluna de imagem. Sem a coluna de imagem, você não consegue criar pixel_values. Por esse motivo, defina remove_unused_columns como False.
Defina eval_do_concat_batches=False para obter resultados de avaliação adequados. As imagens possuem um número diferente de caixas de destino; se os lotes forem concatenados, não será possível determinar a quais imagens pertencem as caixas específicas.

Se você deseja compartilhar seu modelo enviando-o para o Hub, defina push_to_hub como True (você deve estar conectado ao Hugging Face para enviar seu 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 fim, junte tudo e chame 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, Epoch 5/5]

Época	Perda de Treinamento	Perda de Validação	Map	Map 50	Map 75	Map Small	Map Medium	Map Large	1 de março	10 de março	Mar 100	Mar Small	Mar Média	Mar Grande	Grupo Map	Mar 100 Grupo	Mapa de Imagem	Mar 100 Imagem	Mapa Retangular	Mar 100 Retângulo	Mapear texto	Mar 100 Texto
1	No log	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	No log	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	No log	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	No log	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

Se você configurou push_to_hub como True em training_args, os checkpoints de treinamento são enviados para o Hugging Face Hub. Ao concluir o treinamento, envie também o modelo final para o Hub chamando o método push_to_hub().

>>> trainer.push_to_hub()

Avaliar

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

Esses resultados podem ser ainda mais aprimorados aumentando o número de epochs ou ajustando outros hiperparâmetros em TrainingArguments. Experimente!

Inferência

Agora que você ajustou um modelo, o avaliou e o enviou para o Hugging Face Hub, você pode usá-lo para inferência.

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

Carregue o modelo e o processador de imagem do Hugging Face Hub (pule para usar os já treinados nesta sessão):

>>> model_repo = "merve/rf_detr_finetuned_mobile_ui"

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

E detecte caixas 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 plotar o 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