Hugging Face 的 LOGO

社区计算机视觉课程文档

用于对象检测的 Vision Transformers

Hugging Face's logo
加入 Hugging Face 社区

并获得增强的文档体验

开始使用

用于对象检测的 Vision Transformers

本节将介绍如何使用 Vision Transformers 完成对象检测任务。我们将了解如何为我们的用例微调现有的预训练对象检测模型。开始之前，请查看此 HuggingFace Space，您可以在其中玩转最终输出。

引言

Object detection example

对象检测是一项计算机视觉任务，涉及识别图像或视频中的对象并确定其位置。它包括两个主要步骤：

  • 首先，识别存在的对象类型（例如汽车、人或动物）。
  • 其次，通过在它们周围绘制边界框来确定它们的精确位置。

这些模型通常接收图像（静态或视频帧）作为输入，每个图像中都存在多个对象。例如，考虑一个包含汽车、人、自行车等多个对象的图像。处理输入后，这些模型会生成一组数字，传达以下信息：

  • 对象的位置（边界框的 XY 坐标）。
  • 对象的类别。

对象检测有很多应用。其中最重要的例子之一是自动驾驶领域，其中对象检测用于检测汽车周围的不同对象（如行人、路标、交通信号灯等），这些对象成为决策的输入之一。

为了加深您对对象检测内部运作的理解，请查看我们关于对象检测的专用章节🤗。

对象检测中微调模型的必要性 🤔

你应该构建一个新模型，还是修改现有模型？这是一个很棒的问题。从头开始训练一个对象检测模型意味着：

  • 一遍又一遍地做已经完成的研究。
  • 编写重复的模型代码，训练它们，并为不同的用例维护不同的存储库。
  • 大量的实验和资源的浪费。

与其做所有这些，不如采用一个表现良好的预训练模型（一个在识别通用特征方面做得出色的模型），然后调整或重新调整其权重（或其部分权重），使其适应您的用例。我们相信或假设预训练模型已经学到了足够的知识来提取图像中的重要特征以定位和分类对象。因此，如果引入新对象，那么相同的模型可以在短时间内进行训练，并利用已学习和新特征来开始检测这些新对象。

在本教程结束时，您应该能够为对象检测用例构建一个完整的管道（从加载数据集、微调模型到进行推理）。

安装必要的库

让我们从安装开始。只需执行以下单元格即可安装必要的包。在本教程中，我们将使用 Hugging Face Transformers 和 PyTorch。

!pip install -U -q datasets transformers[torch] evaluate timm albumentations accelerate

场景

为了让本教程更有趣，我们来看一个真实的例子。考虑以下场景：建筑工人在施工区域工作时需要 utmost safety。基本的安全协议要求每次都佩戴头盔。由于建筑工人很多，很难时刻盯着每个人。

但是，如果我们能有一个摄像头系统，可以实时检测人员以及他们是否佩戴头盔，那真是太棒了，对吗？

因此，我们将微调一个轻量级对象检测模型来完成这项任务。让我们深入了解一下。

数据集

对于上述场景，我们将使用由中国东北大学提供的hardhat数据集。我们可以使用 🤗 datasets 下载和加载此数据集。

from datasets import load_dataset

dataset = load_dataset("anindya64/hardhat")
dataset

这将为您提供以下数据结构

DatasetDict({
    train: Dataset({
        features: ['image', 'image_id', 'width', 'height', 'objects'],
        num_rows: 5297
    })
    test: Dataset({
        features: ['image', 'image_id', 'width', 'height', 'objects'],
        num_rows: 1766
    })
})

上面是一个DatasetDict，它是一个高效的字典式结构，包含训练和测试拆分中的整个数据集。如您所见，在每个拆分（训练和测试）下，我们都有features和num_rows。在 features 下，我们有image（一个Pillow 对象）、图像的 ID、高度和宽度以及对象。现在让我们看看每个数据点（在训练/测试集中）是什么样子的。为此，运行以下行：

dataset["train"][0]

这将为您提供以下结构：

{'image': <PIL.JpegImagePlugin.JpegImageFile image mode=RGB size=500x375>,
 'image_id': 1,
 'width': 500,
 'height': 375,
 'objects': {'id': [1, 1],
  'area': [3068.0, 690.0],
  'bbox': [[178.0, 84.0, 52.0, 59.0], [111.0, 144.0, 23.0, 30.0]],
  'category': ['helmet', 'helmet']}}

如您所见，`objects` 是另一个字典，其中包含对象 ID（此处为类别 ID）、对象的区域以及边界框坐标（`bbox`）和类别（或标签）。以下是数据元素中每个键和值的更详细解释。

  • image：这是一个 Pillow Image 对象，有助于在从路径加载之前直接查看图像。
  • image_id：表示来自训练文件的图像编号。
  • width：图像的宽度。
  • height：图像的高度。
  • objects：另一个包含注释信息的字典。它包含以下内容：
    • id：一个列表，列表的长度表示对象的数量，每个值表示类别索引。
    • area：对象的面积。
    • bbox：表示对象的边界框坐标。
    • category：对象的类别（字符串）。

现在，让我们正确地提取训练和测试样本。对于本教程，我们大约有 5000 个训练样本和 1700 个测试样本。

# First, extract out the train and test set

train_dataset = dataset["train"]
test_dataset = dataset["test"]

既然我们知道了样本数据点包含什么，让我们开始绘制该样本。在这里，我们将首先绘制图像，然后绘制相应的边界框。

这是我们要做的事情：

  1. 获取图像及其对应的高度和宽度。
  2. 创建一个可以轻松在图像上绘制文本和线条的绘制对象。
  3. 从样本中获取注释字典。
  4. 遍历它。
  5. 对于每个注释，获取边界框坐标，即 x（边界框水平开始的位置）、y（边界框垂直开始的位置）、w（边界框的宽度）、h（边界框的高度）。
  6. 如果边界框度量已标准化，则进行缩放，否则保持不变。
  7. 最后绘制矩形和类别的文本。
import numpy as np
from PIL import Image, ImageDraw


def draw_image_from_idx(dataset, idx):
    sample = dataset[idx]
    image = sample["image"]
    annotations = sample["objects"]
    draw = ImageDraw.Draw(image)
    width, height = sample["width"], sample["height"]

    for i in range(len(annotations["id"])):
        box = annotations["bbox"][i]
        class_idx = annotations["id"][i]
        x, y, w, h = tuple(box)
        if max(box) > 1.0:
            x1, y1 = int(x), int(y)
            x2, y2 = int(x + w), int(y + h)
        else:
            x1 = int(x * width)
            y1 = int(y * height)
            x2 = int((x + w) * width)
            y2 = int((y + h) * height)
        draw.rectangle((x1, y1, x2, y2), outline="red", width=1)
        draw.text((x1, y1), annotations["category"][i], fill="white")
    return image


draw_image_from_idx(dataset=train_dataset, idx=10)

我们有一个函数可以绘制单张图像，现在我们来写一个简单的函数，利用上述功能绘制多张图像。这将有助于我们进行一些分析。

import matplotlib.pyplot as plt


def plot_images(dataset, indices):
    """
    Plot images and their annotations.
    """
    num_rows = len(indices) // 3
    num_cols = 3
    fig, axes = plt.subplots(num_rows, num_cols, figsize=(15, 10))

    for i, idx in enumerate(indices):
        row = i // num_cols
        col = i % num_cols

        # Draw image
        image = draw_image_from_idx(dataset, idx)

        # Display image on the corresponding subplot
        axes[row, col].imshow(image)
        axes[row, col].axis("off")

    plt.tight_layout()
    plt.show()


# Now use the function to plot images

plot_images(train_dataset, range(9))

运行该函数将给我们一个下面所示的精美拼贴画。

input-image-plot

AutoImageProcessor

在微调模型之前，我们必须对数据进行预处理，使其与预训练时使用的方法完全匹配。HuggingFace AutoImageProcessor 负责处理图像数据，以创建 DETR 模型可以训练的 pixel_values、pixel_mask 和 labels。

现在，让我们从与我们想要微调的模型相同的检查点实例化图像处理器。

from transformers import AutoImageProcessor

checkpoint = "facebook/detr-resnet-50-dc5"
image_processor = AutoImageProcessor.from_pretrained(checkpoint)

预处理数据集

在将图像传递给 image_processor 之前，我们还要对图像及其相应的边界框应用不同类型的增强。

简单来说，增强是一些随机变换的集合，如旋转、缩放等。这些变换用于获取更多样本，并使视觉模型对不同的图像条件更具鲁棒性。我们将使用 albumentations 库来实现这一点。它允许您创建图像的随机变换，从而增加训练的样本量。

import albumentations
import numpy as np
import torch

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_processor 期望注释采用以下格式：{'image_id': int, 'annotations': List[Dict]}，其中每个字典都是一个 COCO 对象注释。

def formatted_anns(image_id, category, area, bbox):
    annotations = []
    for i in range(0, len(category)):
        new_ann = {
            "image_id": image_id,
            "category_id": category[i],
            "isCrowd": 0,
            "area": area[i],
            "bbox": list(bbox[i]),
        }
        annotations.append(new_ann)

    return annotations

最后，我们将图像和注释转换结合起来，对整个数据集批次进行转换。

这是最终的代码：

# transforming a batch


def transform_aug_ann(examples):
    image_ids = examples["image_id"]
    images, bboxes, area, categories = [], [], [], []
    for image, objects in zip(examples["image"], examples["objects"]):
        image = np.array(image.convert("RGB"))[:, :, ::-1]
        out = transform(image=image, bboxes=objects["bbox"], category=objects["id"])

        area.append(objects["area"])
        images.append(out["image"])
        bboxes.append(out["bboxes"])
        categories.append(out["category"])

    targets = [
        {"image_id": id_, "annotations": formatted_anns(id_, cat_, ar_, box_)}
        for id_, cat_, ar_, box_ in zip(image_ids, categories, area, bboxes)
    ]

    return image_processor(images=images, annotations=targets, return_tensors="pt")

最后，您所要做的就是将此预处理函数应用于整个数据集。您可以通过使用 HuggingFace 🤗 Datasets with transform 方法来实现。

# Apply transformations for both train and test dataset

train_dataset_transformed = train_dataset.with_transform(transform_aug_ann)
test_dataset_transformed = test_dataset.with_transform(transform_aug_ann)

现在让我们看看转换后的训练数据集样本是什么样子

train_dataset_transformed[0]

这将返回一个张量字典。我们这里主要需要的是表示图像的 pixel_values，表示注意力遮罩的 pixel_mask 以及 labels。这是一个数据点的样子：

{'pixel_values': tensor([[[-0.1657, -0.1657, -0.1657,  ..., -0.3369, -0.4739, -0.5767],
          [-0.1657, -0.1657, -0.1657,  ..., -0.3369, -0.4739, -0.5767],
          [-0.1657, -0.1657, -0.1828,  ..., -0.3541, -0.4911, -0.5938],
          ...,
          [-0.4911, -0.5596, -0.6623,  ..., -0.7137, -0.7650, -0.7993],
          [-0.4911, -0.5596, -0.6794,  ..., -0.7308, -0.7993, -0.8335],
          [-0.4911, -0.5596, -0.6794,  ..., -0.7479, -0.8164, -0.8507]],
 
         [[-0.0924, -0.0924, -0.0924,  ...,  0.0651, -0.0749, -0.1800],
          [-0.0924, -0.0924, -0.0924,  ...,  0.0651, -0.0924, -0.2150],
          [-0.0924, -0.0924, -0.1099,  ...,  0.0476, -0.1275, -0.2500],
          ...,
          [-0.0924, -0.1800, -0.3200,  ..., -0.4426, -0.4951, -0.5301],
          [-0.0924, -0.1800, -0.3200,  ..., -0.4601, -0.5126, -0.5651],
          [-0.0924, -0.1800, -0.3200,  ..., -0.4601, -0.5301, -0.5826]],
 
         [[ 0.1999,  0.1999,  0.1999,  ...,  0.6705,  0.5136,  0.4091],
          [ 0.1999,  0.1999,  0.1999,  ...,  0.6531,  0.4962,  0.3916],
          [ 0.1999,  0.1999,  0.1825,  ...,  0.6356,  0.4614,  0.3568],
          ...,
          [ 0.4788,  0.3916,  0.2696,  ...,  0.1825,  0.1302,  0.0953],
          [ 0.4788,  0.3916,  0.2696,  ...,  0.1651,  0.0953,  0.0605],
          [ 0.4788,  0.3916,  0.2696,  ...,  0.1476,  0.0779,  0.0431]]]),
 'pixel_mask': tensor([[1, 1, 1,  ..., 1, 1, 1],
         [1, 1, 1,  ..., 1, 1, 1],
         [1, 1, 1,  ..., 1, 1, 1],
         ...,
         [1, 1, 1,  ..., 1, 1, 1],
         [1, 1, 1,  ..., 1, 1, 1],
         [1, 1, 1,  ..., 1, 1, 1]]),
 'labels': {'size': tensor([800, 800]), 'image_id': tensor([1]), 'class_labels': tensor([1, 1]), 'boxes': tensor([[0.5920, 0.3027, 0.1040, 0.1573],
         [0.7550, 0.4240, 0.0460, 0.0800]]), 'area': tensor([8522.2217, 1916.6666]), 'iscrowd': tensor([0, 0]), 'orig_size': tensor([480, 480])}}

我们差不多了 🚀。作为最后一个预处理步骤，我们需要编写一个自定义的 collate_fn。那么什么是 collate_fn 呢？

collate_fn 负责从数据集中获取样本列表，并将它们转换为适合模型输入格式的批次。

通常，`DataCollator` 通常执行填充、截断等任务。在自定义 collate 函数中，我们通常定义我们想要如何将数据分组到批次中，或者简单地说，如何表示每个批次。

数据整理器主要将数据整合在一起，然后对其进行预处理。让我们制作我们的整理函数。

def collate_fn(batch):
    pixel_values = [item["pixel_values"] for item in batch]
    encoding = image_processor.pad(pixel_values, return_tensors="pt")
    labels = [item["labels"] for item in batch]
    batch = {}
    batch["pixel_values"] = encoding["pixel_values"]
    batch["pixel_mask"] = encoding["pixel_mask"]
    batch["labels"] = labels
    return batch

训练 DETR 模型。

到目前为止，所有繁重的工作都已完成。现在，剩下的就是将拼图的各个部分一一组装起来。开始吧！

训练过程包括以下步骤：

  1. 使用AutoModelForObjectDetection加载基础（预训练）模型，使用与预处理相同的检查点。

  2. 在TrainingArguments中定义所有超参数和附加参数。

  3. 将训练参数传递给HuggingFace Trainer，以及模型、数据集和图像。

  4. 调用 train() 方法并微调您的模型。

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

from transformers import AutoModelForObjectDetection

id2label = {0: "head", 1: "helmet", 2: "person"}
label2id = {v: k for k, v in id2label.items()}


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

在继续之前，请登录 Hugging Face Hub，以便在训练时即时上传您的模型。这样，您就不需要处理检查点并将其保存到其他地方。

from huggingface_hub import notebook_login

notebook_login()

完成后，我们开始训练模型。我们首先定义训练参数，然后定义一个使用这些参数进行训练的训练器对象，如下所示：

from transformers import TrainingArguments
from transformers import Trainer

# Define the training arguments

training_args = TrainingArguments(
    output_dir="detr-resnet-50-hardhat-finetuned",
    per_device_train_batch_size=8,
    num_train_epochs=3,
    max_steps=1000,
    fp16=True,
    save_steps=10,
    logging_steps=30,
    learning_rate=1e-5,
    weight_decay=1e-4,
    save_total_limit=2,
    remove_unused_columns=False,
    push_to_hub=True,
)

# Define the trainer

trainer = Trainer(
    model=model,
    args=training_args,
    data_collator=collate_fn,
    train_dataset=train_dataset_transformed,
    eval_dataset=test_dataset_transformed,
    tokenizer=image_processor,
)

trainer.train()

训练完成后，您现在可以删除模型，因为检查点已上传到 HuggingFace Hub。

del model
torch.cuda.synchronize()

测试与推理

现在我们将尝试对新微调的模型进行推理。在本教程中，我们将测试此图像：

input-test-image

这里我们首先编写一个非常简单的代码，用于对一些新图像进行对象检测推理。我们从对单个图像进行推理开始，然后将所有内容组合起来并将其制成一个函数。

import requests
from transformers import pipeline

# download a sample image

url = "https://huggingface.co/datasets/hf-vision/course-assets/resolve/main/test-helmet-object-detection.jpg"
image = Image.open(requests.get(url, stream=True).raw)

# make the object detection pipeline

obj_detector = pipeline(
    "object-detection", model="anindya64/detr-resnet-50-dc5-hardhat-finetuned"
)
results = obj_detector(train_dataset[0]["image"])

print(results)

现在让我们编写一个非常简单的函数，用于在我们的图像上绘制结果。我们从结果中获取分数、标签和相应的边界框坐标，这些将用于在图像中绘制。

def plot_results(image, results, threshold=0.7):
    image = Image.fromarray(np.uint8(image))
    draw = ImageDraw.Draw(image)
    for result in results:
        score = result["score"]
        label = result["label"]
        box = list(result["box"].values())
        if score > threshold:
            x, y, x2, y2 = tuple(box)
            draw.rectangle((x, y, x2, y2), outline="red", width=1)
            draw.text((x, y), label, fill="white")
            draw.text(
                (x + 0.5, y - 0.5),
                text=str(score),
                fill="green" if score > 0.7 else "red",
            )
    return image

最后，将此函数用于我们使用的同一测试图像。

results = obj_detector(image)
plot_results(image, results)

这将绘制以下输出：

output-test-image-plot

现在，让我们把所有东西组合成一个简单的函数。

def predict(image, pipeline, threshold=0.7):
    results = pipeline(image)
    return plot_results(image, results, threshold)


# Let's test for another test image

img = test_dataset[0]["image"]
predict(img, obj_detector)

我们甚至可以使用我们的推理函数在少量测试样本上绘制多张图像。

from tqdm.auto import tqdm


def plot_images(dataset, indices):
    """
    Plot images and their annotations.
    """
    num_rows = len(indices) // 3
    num_cols = 3
    fig, axes = plt.subplots(num_rows, num_cols, figsize=(15, 10))

    for i, idx in tqdm(enumerate(indices), total=len(indices)):
        row = i // num_cols
        col = i % num_cols

        # Draw image
        image = predict(dataset[idx]["image"], obj_detector)

        # Display image on the corresponding subplot
        axes[row, col].imshow(image)
        axes[row, col].axis("off")

    plt.tight_layout()
    plt.show()


plot_images(test_dataset, range(6))

运行此函数将得到如下输出：

test-sample-output-plot

嗯，还不错。如果再进行微调，我们可以改进结果。您可以在此处找到这个微调的检查点。

< > 在 GitHub 上更新

←MobileViT v2 DEtection TRansformer (DETR)→
© . This site is unofficial and not affiliated with Hugging Face, Inc.