每當(dāng)應(yīng)用機(jī)器學(xué)習(xí)來(lái)解決問(wèn)題時(shí)，目標(biāo)在某種程度上都是將模型擬合到某些數(shù)據(jù)上。為了使模型表現(xiàn)良好并泛化到未見(jiàn)過(guò)的數(shù)據(jù)，你需要確保使用高質(zhì)量的數(shù)據(jù)集進(jìn)行訓(xùn)練。特別是在監(jiān)督學(xué)習(xí)環(huán)境中，你需要確保數(shù)據(jù)被準(zhǔn)確標(biāo)注。

無(wú)論你構(gòu)建的模型有多大，投入了多少億個(gè)參數(shù)，或者對(duì)數(shù)據(jù)集進(jìn)行了多少增強(qiáng)，低質(zhì)量的輸入不會(huì)神奇地變成高質(zhì)量的輸出。

根據(jù)你試圖解決的任務(wù)，并不總是有合適的公開(kāi)數(shù)據(jù)集可用。在這種情況下，你可能需要構(gòu)建自己的數(shù)據(jù)集。然而，最初你的數(shù)據(jù)很可能沒(méi)有被標(biāo)注。讓我向你展示如何構(gòu)建一個(gè)簡(jiǎn)單、快速的標(biāo)注工具，從未標(biāo)注的數(shù)據(jù)集中分類你的圖像數(shù)據(jù)。

演示

圖像數(shù)據(jù)集

數(shù)據(jù)集中的樣本

為了演示這個(gè)標(biāo)注工具，我將使用手機(jī)錄制的圖像數(shù)據(jù)集，目標(biāo)是對(duì)三種不同的USB接口類型進(jìn)行分類：USB-A、USB-C、Micro USB和Mini USB。最初，所有圖像都未標(biāo)注，存放在輸入目錄中。我們的標(biāo)注工具應(yīng)該逐個(gè)展示圖像，并在指定類別后將其移動(dòng)到相應(yīng)的目錄。

標(biāo)注工具在運(yùn)行中

指南

1. 先決條件

如果你想跟著操作，你應(yīng)該安裝opencv-python。你可以在項(xiàng)目倉(cāng)庫(kù)的示例文件夾中找到一些示例圖像。

2. 數(shù)據(jù)加載

首先，我們從輸入文件夾中加載圖像。我們可以使用pathlib中的glob函數(shù)查找所有jpg圖像擴(kuò)展名的文件。將結(jié)果傳遞給sorted函數(shù)，確保圖像按順序處理。

from pathlib import Path


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

我們還需要準(zhǔn)備輸出目錄，確保它存在。

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

我們可以遍歷圖像列表，并使用cv2.imread將圖像加載到數(shù)組中。然后顯示圖像并等待按鍵。通過(guò)在cv2.waitKey函數(shù)中將延遲設(shè)置為0，我們無(wú)限期地等待，直到按下某個(gè)鍵。然后確保按下Q可以退出應(yīng)用程序，最后關(guān)閉所有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被激活，數(shù)字仍然與ord函數(shù)的數(shù)字相同。

3. 標(biāo)注

讓我們?cè)谧址斜碇卸x任務(wù)的標(biāo)簽。在我的例子中，我有四個(gè)不同接口的標(biāo)簽：

...


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

現(xiàn)在我們希望數(shù)字鍵0、1、2和3將圖像分類到相應(yīng)的標(biāo)簽文件夾中。waitKey函數(shù)中的key變量是一個(gè)整數(shù)，表示按下字符的Unicode代碼。要檢查按鍵是否是數(shù)字之一，我們需要使用ord函數(shù)將數(shù)字轉(zhuǎn)換為Unicode，類似于我們檢查按鍵q以關(guān)閉窗口的方式。該函數(shù)期望長(zhǎng)度為1的字符串，因此我們需要在將索引傳遞給函數(shù)之前將其轉(zhuǎn)換為字符串。

...


  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

要將圖像移動(dòng)到輸出路徑中的分類標(biāo)簽文件夾，我們可以使用pathlib中的/操作來(lái)連接路徑，然后使用rename函數(shù)將文件移動(dòng)到目標(biāo)位置。

...


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

在我們這樣做之前，我們需要確保目標(biāo)文件夾存在。所以在循環(huán)之前，我們遍歷所有標(biāo)簽并創(chuàng)建相應(yīng)的文件夾。

...


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


while True:
    ...

另一種更Pythonic的標(biāo)簽鍵檢查方法是在循環(huán)之前創(chuàng)建鍵Unicode到標(biāo)簽的映射。這樣，我們不需要在循環(huán)的每一步中遍歷所有鍵。

# mapping from key to label
labels_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

讓我們還為鍵到標(biāo)簽的映射添加一個(gè)小幫助文本。

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

結(jié)論

在本教程中，你學(xué)習(xí)了如何為圖像分類任務(wù)創(chuàng)建一個(gè)簡(jiǎn)單的標(biāo)注工具。我們可以對(duì)這個(gè)工具進(jìn)行很多改進(jìn)。我想進(jìn)一步探索的一件事是添加不僅分類圖像的功能，還可以分割圖像并創(chuàng)建分割掩碼。

當(dāng)然，有很多更復(fù)雜的工具可以簡(jiǎn)化你的標(biāo)注過(guò)程。然而，有時(shí)一個(gè)非常簡(jiǎn)單的工具就是你所需要的，特別是在項(xiàng)目的早期階段進(jìn)行探索性數(shù)據(jù)分析時(shí)，你需要一個(gè)快速的概念驗(yàn)證。

【參考資料】

• 代碼鏈接：https://github.com/trflorian/annotation-tool?source=post_page-----c0549b644d15--------------------------------
• 文章來(lái)源：https://towardsdatascience.com/build-your-own-annotation-tool-for-image-classification-in-5-minutes-c0549b644d15