目标检测

目标检测是计算机视觉中的一项任务，旨在检测图像中的实例（例如人、建筑物或汽车）。目标检测模型以图像作为输入，输出检测到的对象的边界框坐标及其对应的标签。一张图像可以包含多个对象，每个对象都有各自的边界框和标签（例如，图像中可能同时包含一辆汽车和一座建筑物），并且每个对象可以出现在图像的不同位置（例如，图像中可能有多辆汽车）。该任务常用于自动驾驶领域，以检测行人、交通标志和交通信号灯等。其他应用包括图像中的对象计数、图像搜索等。

在本指南中，您将学习如何：

在 mobile-ui-design 数据集上微调 RF-DETR，以检测移动应用截图中的 UI 元素。
使用微调后的模型进行推理。

要查看与此任务兼容的所有架构和检查点，我们建议查阅任务页面

在开始之前，请确保已安装所有必要的库：

pip install -q datasets transformers accelerate timm trackio torchmetrics pycocotools

您将使用 🤗 Datasets 从 Hugging Face Hub 加载数据集，并使用 🤗 Transformers 训练您的模型。

我们鼓励您与社区分享您的模型。请登录您的 Hugging Face 账户，将其上传到 Hub。当系统提示时，请输入您的 token 以完成登录：

>>> from huggingface_hub import notebook_login

>>> notebook_login()

定义全局常量，即模型名称和图像尺寸。本教程使用 RF-DETR，但您可以选择 Transformers 中的任意目标检测模型。

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

加载 mobile-ui-design 数据集

mobile-ui-design 数据集包含带有注释的移动应用截图，用于检测文本、图像、矩形和组等 UI 元素。

首先加载数据集并提取类别标签。该数据集已经包含数据划分。

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

该数据集使用字符串类别名称和 COCO 格式的边界框 (x, y, w, h)。在训练之前，请将类别转换为整数 ID，计算面积，并过滤掉退化的边界框：

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

预处理数据

AutoImageProcessor 负责处理图像数据，生成模型可用于训练的 pixel_values、pixel_mask 和 labels。图像处理器会处理缩放、填充和归一化。此外，你还可以选择性地添加随机数据增强（参见下文）以提升泛化能力。

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

>>> image_processor = AutoImageProcessor.from_pretrained(MODEL_NAME)

image_processor 期望以 COCO 格式提供标注：{'image_id': int, 'annotations': list[Dict]}。请格式化每个示例的标注，其余部分交由处理器处理：

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

数据增强

上述转换仅对图像进行缩放和归一化。通常，对训练集应用随机数据增强可以提升模型的泛化能力，而验证集应保持无增强，以确保评估结果的确定性。一个常见的选择是 Albumentations，它可以同时对图像及其边界框进行增强。使用 bbox_params 定义一个 pipeline，使边界框与图像保持一致的变换，然后根据增强后的边界框重新计算面积：

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

仅对训练集应用增强转换，验证集保留原始的 transform_fn：

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

创建一个自定义的 collate_fn 以将图像进行批量处理：

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

准备用于计算 mAP 的函数

目标检测模型通常使用一组 COCO 风格的指标进行评估。我们将使用 torchmetrics 来计算 mAP（平均精度均值）和 mAR（平均召回率均值）指标，并将其封装为 compute_metrics 函数，以便在 Trainer 中用于评估。

用于训练的框的中间格式为 YOLO（归一化），但为了正确处理框的面积，我们将以 Pascal VOC（绝对值）格式计算框的指标。让我们定义一个函数，将边界框转换为 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

然后，在 compute_metrics 函数中，我们从评估循环的结果中收集 predicted 和 target 的边界框、分数和标签，并将其传递给评分函数。

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

训练检测模型

在之前的章节中，您已经完成了大部分繁重的工作，现在可以开始训练您的模型了！该数据集中的图像即使经过调整大小后仍然相当大。这意味着微调该模型至少需要一块 GPU。

训练包括以下步骤：

使用与预处理中相同的检查点，通过 AutoModelForObjectDetection 加载模型。
在 TrainingArguments 中定义你的训练超参数。
将训练参数与模型、数据集、图像处理器和数据整理器一起传递给 Trainer。
调用 train() 来微调你的模型。

从用于预处理的同一检查点加载模型时，请记得传入之前从数据集元数据中创建的 label2id 和 id2label 映射。此外，我们指定 ignore_mismatched_sizes=True，以用新的分类头替换现有的分类头。

>>> from transformers import AutoModelForObjectDetection

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

在 TrainingArguments 中，使用 output_dir 指定模型的保存位置，然后按需配置超参数。对于 num_train_epochs=5，在 A100 GPU 上训练大约需要 35 分钟，增加训练轮数可获得更好的效果。

重要说明：

不要删除未使用的列，因为这会丢弃图像列。如果没有图像列，就无法创建 pixel_values。因此，请将 remove_unused_columns 设置为 False。
将 eval_do_concat_batches=False 设置为 False 以获得正确的评估结果。由于不同图像的目标框数量不同，如果将批次进行拼接，将无法确定哪些框属于特定图像。

如果您希望通过推送到 Hub 来共享您的模型，请将 push_to_hub 设置为 True（您必须已登录 Hugging Face 才能上传模型）。

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

最后，将所有内容整合起来，并调用 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]

Epoch	训练损失	验证损失	映射	Map 50	Map 75	映射小型	中等映射	映射大型	3 月 1 日	3 月 10 日	3月100日	3月小	3月中等	Mar Large	映射组	3月100组	图像映射	Mar 100 图像	映射矩形	Mar 100 矩形	文本映射	3月100 文本
1	无日志	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	无日志	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	无日志	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	无日志	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

如果在 training_args 中将 push_to_hub 设置为 True，训练检查点将被推送到 Hugging Face Hub。训练完成后，调用 push_to_hub() 方法将最终模型也推送到 Hub。

>>> trainer.push_to_hub()

评估

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

通过增加训练轮数或调整 TrainingArguments 中的其他超参数，可以进一步提升这些结果。试试看吧！

推理

现在，您已经微调了模型、对其进行了评估，并将其上传到了 Hugging Face Hub，因此您可以将其用于推理。

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

从 Hugging Face Hub 加载模型和图像处理器（如果本次会话中已经训练过，可跳过）：

>>> model_repo = "merve/rf_detr_finetuned_mobile_ui"

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

并检测边界框：

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

让我们绘制结果：

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