用OpenCV和Python构建自己的图像分类标注工具

每当应用机器学习来解决问题时，目标在某种程度上都是将模型拟合到某些数据上。为了使模型表现良好并泛化到未见过的数据，你需要确保使用高质量的数据集进行训练。特别是在监督学习环境中，你需要确保数据被准确标注。

无论你构建的模型有多大，投入了多少亿个参数，或者对数据集进行了多少增强，低质量的输入不会神奇地变成高质量的输出。

根据你试图解决的任务，并不总是有合适的公开数据集可用。在这种情况下，你可能需要构建自己的数据集。然而，最初你的数据很可能没有被标注。让我向你展示如何构建一个简单、快速的标注工具，从未标注的数据集中分类你的图像数据。

演示

图像数据集

数据集中的样本

为了演示这个标注工具，我将使用手机录制的图像数据集，目标是对三种不同的USB接口类型进行分类：USB-A、USB-C、Micro USB和Mini USB。最初，所有图像都未标注，存放在输入目录中。我们的标注工具应该逐个展示图像，并在指定类别后将其移动到相应的目录。

标注工具在运行中

指南
先决条件
如果你想跟着操作，你应该安装opencv-python。你可以在项目仓库的示例文件夹中找到一些示例图像。

数据加载
首先，我们从输入文件夹中加载图像。我们可以使用pathlib中的glob函数查找所有jpg图像扩展名的文件。将结果传递给sorted函数，确保图像按顺序处理。

from pathlib import Path
input_path = Path("input")input_img_paths = sorted(input_path.glob("*.jpg"))

我们还需要准备输出目录，确保它存在。

output_path = Path("output")output_path.mkdir(parents=True, exist_ok=True)

我们可以遍历图像列表，并使用cv2.imread将图像加载到数组中。然后显示图像并等待按键。通过在cv2.waitKey函数中将延迟设置为0，我们无限期地等待，直到按下某个键。然后确保按下Q可以退出应用程序，最后关闭所有opencv窗口。

import cv2
...
def annotate_images(    input_img_paths: list[Path],    output_path: Path,)-> None:
    for img_path in input_img_paths:        img = cv2.imread(str(img_path))
        cv2.imshow("Image", img)
        while True:            key = cv2.waitKey(0)
            # Quit Annotation Tool            if key == ord("q"):                return
        cv2.destroyAllWindows()

注意：使用按位与（&）和0xFF，我们只查看按键的最后几位。这确保即使NumLock被激活，数字仍然与ord函数的数字相同。

标注
让我们在字符串列表中定义任务的标签。在我的例子中，我有四个不同接口的标签：

...
def annotate_images(    input_img_paths: list[Path],    output_path: Path,    labels: list[str],) -> None:    ...
annotate_images(    input_img_paths=input_img_paths,    output_path=output_path,    labels=["usb_a", "usb_c", "usb_mini", "usb_micro"],)

现在我们希望数字键0、1、2和3将图像分类到相应的标签文件夹中。waitKey函数中的key变量是一个整数，表示按下字符的Unicode代码。要检查按键是否是数字之一，我们需要使用ord函数将数字转换为Unicode，类似于我们检查按键q以关闭窗口的方式。该函数期望长度为1的字符串，因此我们需要在将索引传递给函数之前将其转换为字符串。

  ...
  while True:
    ...
    for i in range(len(labels)):        if key == ord(str(i)):            label = labels[i]            print(f"Classified as {label}")
            # TODO: move to correct label folder
            break

要将图像移动到输出路径中的分类标签文件夹，我们可以使用pathlib中的/操作来连接路径，然后使用rename函数将文件移动到目标位置。

...
if key == ord(str(i)):    label = labels[i]    print(f"Classified as {label}")
    output_img_path = output_path / label / img_path.name    img_path.rename(output_img_path)
    break

在我们这样做之前，我们需要确保目标文件夹存在。所以在循环之前，我们遍历所有标签并创建相应的文件夹。

...
# create all classification foldersfor label in labels:    label_dir = output_path / label    label_dir.mkdir(parents=True, exist_ok=True)
while True:    ...

另一种更Pythonic的标签键检查方法是在循环之前创建键Unicode到标签的映射。这样，我们不需要在循环的每一步中遍历所有键。

# mapping from key to labellabels_key_dict = {ord(str(i)): label for i, label in enumerate(labels)}
while True:    ...
    if key in labels_key_dict:        label = labels_key_dict[key]        print(f"Classified as {label}")
        output_img_path = output_path / label / img_path.name        img_path.rename(output_img_path)
        break

让我们还为键到标签的映射添加一个小帮助文本。

for i, label in enumerate(labels):    cv2.putText(        img,        f"{i}: {label}",        (10, 30 + 30 * i),        cv2.FONT_HERSHEY_SIMPLEX,        1,        (255, 255, 255),        2,        cv2.LINE_AA,    )

结论
在本教程中，你学习了如何为图像分类任务创建一个简单的标注工具。我们可以对这个工具进行很多改进。我想进一步探索的一件事是添加不仅分类图像的功能，还可以分割图像并创建分割掩码。

当然，有很多更复杂的工具可以简化你的标注过程。然而，有时一个非常简单的工具就是你所需要的，特别是在项目的早期阶段进行探索性数据分析时，你需要一个快速的概念验证。

【参考资料】

代码链接：https://github.com/trflorian/annotation-tool?source=post_page-----c0549b644d15--------------------------------

文章来源：https://towardsdatascience.com/build-your-own-annotation-tool-for-image-classification-in-5-minutes-c0549b644d15