Press "Enter" to skip to content

我如何使用Grounding DINO进行自动图像标注

我很高兴地分享最近计算机视觉领域的进展,例如突破性的零样本目标检测器(如Grounding DINO)的出现,已经彻底改变了图像标注的过程

作为一名机器学习开发者,我个人认为图像标注是一项乏味、耗时且昂贵的任务。但是幸运的是,随着计算机视觉领域的最新发展,特别是强大的零样本目标检测器,如Grounding DINO的引入,我们实际上可以自动化大部分用例的图像标注过程。我们可以编写一个Python脚本,它会为我们完成95%的工作。我们唯一的任务是在最后审查这些注释,并可能添加或删除一些边界框。

我如何使用Grounding DINO进行自动图像标注 计算科学 第1张

在进入自动图像标注之前,我们应该知道Grounding DINO是什么?我们为什么要使用它?

Grounding DINO可以检测给定提示输入(如类别名称或指示性表达式)中的主要对象。开放集对象检测的主要解决方案是将语言引入到封闭集检测器中。DINO用于开放集概念泛化:为了有效地融合语言和视觉模态,我们在概念上将封闭集检测器分为三个阶段:骨干、颈部和头部。然后,我们通过在颈部查询初始化和头部中融合语言信息来提出一种紧密的融合解决方案。Grounding DINO包括特征增强器、语言引导查询选择和跨模态融合的跨模态解码器。

在COCO数据集检测零样本转移基准测试中,Grounding DINO实现了52.5%的平均精度(AP),这是没有任何来自COCO数据集的训练数据的。在对COCO数据集进行微调后,它达到了63.0 AP。它在OdinW零样本基准测试中平均达到26.1 AP,创造了新的记录。我们还探讨了如何利用预训练的DINO仅通过训练语言和融合模块来提高性能。Grounding DINO通过DINO快速收敛比基线模型更快。

我们的Grounding DINO还可以与稳定的扩散一起协作进行图像编辑,例如,我们可以检测图像中的Green Mountain,并生成一个文本提示Red Mountain的新图像,还可以通过首先检测面部来修改人物的背景。我们还可以使用GLIGEN进行更详细的控制,如为每个框分配一个对象,这是我们的Grounding DINO模型,用于开放集对象检测。

好的,进入自动图像标注部分,这里我正在使用高计算能力的Google Colab

让我们开始吧,

让我们确保我们有GPU访问权限。我们可以使用nvidia-smi命令检查GPU是否连接。如果遇到任何问题,请导航到“编辑”->“笔记本设置”->“硬件加速器”,将其设置为GPU,然后单击“保存”。这将大大缩短自动标注完成所需的时间。

nvidia-smi

安装Grounding DINO模型

我们的项目将使用创新性的设计——Grounding DINO进行零样本检测。我们必须先安装它。

!git clone https://github.com/IDEA-Research/GroundingDINO.git
%cd GroundingDINO
!git checkout -q 57535c5a79791cb76e36fdb64975271354f10251
!pip install -q -e .

监督Python索引包将帮助我们处理、过滤和可视化我们的检测结果,以及保存我们的数据集,它将是我们演示所有部分的粘合剂。使用Grounding DINO安装了一个较小版本的“supervision”。但是对于这个演示,我们需要最近迭代中添加的新功能。为了安装版本“0.6.0”,我们首先卸载当前的“supervision”版本。

!pip uninstall -y supervision
!pip install -q supervision==0.6.0

import supervision as svn
print(svn.__version__)

下载Grounding DINO模型权重

我们需要配置文件和模型权重文件才能运行Grounding DINO。我们已经克隆了Grounding DINO存储库,其中包含配置文件。另一方面,我们必须下载权重文件。在将两个文件的路径写入变量GROUNDING_DINO_CONFIG_PATH和GROUNDING_DINO_CHECKPOINT_PATH后,我们检查路径是否正确,并且文件是否存在于磁盘上。

import os

GROUNDING_DINO_CONFIG_PATH = os.path.join("groundingdino/config/GroundingDINO_SwinT_OGC.py")
print(GROUNDING_DINO_CONFIG_PATH, "; 存在:", os.path.isfile(GROUNDING_DINO_CONFIG_PATH))
!mkdir -p weights
%cd weights

!wget -q https://github.com/IDEA-Research/GroundingDINO/releases/download/v0.1.0-alpha/groundingdino_swint_ogc.pth
import os
%cd /content/GroundingDINO
GROUNDING_DINO_CHECKPOINT_PATH = os.path.join("weights/groundingdino_swint_ogc.pth")
print(GROUNDING_DINO_CHECKPOINT_PATH, "; 存在:", os.path.isfile(GROUNDING_DINO_CHECKPOINT_PATH))

假设您已经安装了 PyTorch,您可以使用以下命令行导入 torch 并设置用于计算的设备:

import torch

DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

加载 Grounding DINO 模型

from groundingdino.util.inference import Model

grounding_dino_model = Model(model_config_path=GROUNDING_DINO_CONFIG_PATH, model_checkpoint_path=GROUNDING_DINO_CHECKPOINT_PATH)

数据集准备

创建一个名为 data 的文件夹,并将未标记的图像移动到该文件夹中。

!mkdir -p data

单张图像蒙版自动注释

在自动注释整个数据集之前,让我们暂时专注于一张图像。

SOURCE_IMAGE_PATH = "/content/GroundingDINO/data/example_image_3.png"
CLASSES = ['person','dog'] #添加需要自动标记的类名
BOX_TRESHOLD = 0.35
TEXT_TRESHOLD = 0.15

使用 Grounding DINO 进行零样本物体检测

我们将使用下面描述的 enhance_class_name 函数,使用一些提示工程来获得更好的 Grounding DINO 检测。

from typing import List

def enhance_class_name(class_names: List[str]) -> List[str]:
   return [
       f"all {class_name}s"
       for class_name
       in class_names
   ]
import cv2
import supervision as sv

# 加载图像
image = cv2.imread(SOURCE_IMAGE_PATH)

# 检测物体
detections = grounding_dino_model.predict_with_classes(
   image=image,
   classes=enhance_class_name(class_names=CLASSES),
   box_threshold=BOX_TRESHOLD,
   text_threshold=TEXT_TRESHOLD
)

# 使用检测结果注释图像
box_annotator = svn.BoxAnnotator()
labels = [
   f"{CLASSES[class_id]} {confidence:0.2f}"
   for _, _, confidence, class_id, _
   in detections]
annotated_frame = box_annotator.annotate(scene=image.copy(), detections=detections, labels=labels)

%matplotlib inline
svn.plot_image(annotated_frame, (16, 16))

我如何使用Grounding DINO进行自动图像标注 计算科学 第2张

完整数据集蒙版自动标注

import os

IMAGES_DIRECTORY = "./data"
IMAGES_EXTENSIONS = ['jpg', 'jpeg', 'png']

CLASSES = ['person','dog']
BOX_TRESHOLD = 0.35
TEXT_TRESHOLD = 0.15

从图像中提取标签

import cv2
from tqdm.notebook import tqdm

images = {}
annotations = {}

image_paths = svn.list_files_with_extensions(
   directory=IMAGES_DIRECTORY,
   extensions=IMAGES_EXTENSIONS)

for image_path in tqdm(image_paths):
   image_name = image_path.name
   image_path = str(image_path)
   image = cv2.imread(image_path)

   detections = grounding_dino_model.predict_with_classes(
       image=image,
       classes=enhance_class_name(class_names=CLASSES),
       box_threshold=BOX_TRESHOLD,
       text_threshold=TEXT_TRESHOLD
   )
   detections = detections[detections.class_id != None]
   images[image_name] = image
   annotations[image_name] = detections

绘制结果

plot_images = []
plot_titles = []

box_annotator = svn.BoxAnnotator()
mask_annotator = svn.MaskAnnotator()

for image_name, detections in annotations.items():
   image = images[image_name]
   plot_images.append(image)
   plot_titles.append(image_name)

   labels = [
       f"{CLASSES[class_id]} {confidence:0.2f}"
       for _, _, confidence, class_id, _
       in detections]
   annotated_image = mask_annotator.annotate(scene=image.copy(), detections=detections)
   annotated_image = box_annotator.annotate(scene=annotated_image, detections=detections, labels=labels)
   plot_images.append(annotated_image)
   title = " ".join(set([
       CLASSES[class_id]
       for class_id
       in detections.class_id
   ]))
   plot_titles.append(title)

svn.plot_images_grid(
   images=plot_images,
   titles=plot_titles,
   grid_size=(len(annotations), 2),
   size=(2 * 4, len(annotations) * 4)

我如何使用Grounding DINO进行自动图像标注 计算科学 第3张

将标签保存在Pascal VOC XML中

%cd /content/GroundingDINO
!mkdir annotations
ANNOTATIONS_DIRECTORY = "/content/GroundingDINO/annotations"

MIN_IMAGE_AREA_PERCENTAGE = 0.002
MAX_IMAGE_AREA_PERCENTAGE = 0.80
APPROXIMATION_PERCENTAGE = 0.75
svn.Dataset(
   classes=CLASSES,
   images=images,
   annotations=annotations
).as_pascal_voc(
   annotations_directory_path=ANNOTATIONS_DIRECTORY,
   min_image_area_percentage=MIN_IMAGE_AREA_PERCENTAGE,
   max_image_area_percentage=MAX_IMAGE_AREA_PERCENTAGE,
   approximation_percentage=APPROXIMATION_PERCENTAGE
)

感谢阅读!

这里有整个colab文件的链接。

参考文献:https://arxiv.org/abs/2303.05499 & https://github.com/IDEA-Research/GroundingDINO

Parthiban M目前居住在印度金奈,就职于SeeWise。他是一位机器学习开发人员,具有广泛的经验,能够使用计算机视觉、TensorFlow和深度学习理解问题并提供解决方案。

Leave a Reply

Your email address will not be published. Required fields are marked *