使用 OpenCV 和 Python 構(gòu)建自己的圖像分類標注工具
每當應(yīng)用機器學(xué)習(xí)來解決問題時,目標在某種程度上都是將模型擬合到某些數(shù)據(jù)上。為了使模型表現(xiàn)良好并泛化到未見過的數(shù)據(jù),你需要確保使用高質(zhì)量的數(shù)據(jù)集進行訓(xùn)練。特別是在監(jiān)督學(xué)習(xí)環(huán)境中,你需要確保數(shù)據(jù)被準確標注。
無論你構(gòu)建的模型有多大,投入了多少億個參數(shù),或者對數(shù)據(jù)集進行了多少增強,低質(zhì)量的輸入不會神奇地變成高質(zhì)量的輸出。
根據(jù)你試圖解決的任務(wù),并不總是有合適的公開數(shù)據(jù)集可用。在這種情況下,你可能需要構(gòu)建自己的數(shù)據(jù)集。然而,最初你的數(shù)據(jù)很可能沒有被標注。讓我向你展示如何構(gòu)建一個簡單、快速的標注工具,從未標注的數(shù)據(jù)集中分類你的圖像數(shù)據(jù)。
演示
圖像數(shù)據(jù)集
數(shù)據(jù)集中的樣本
為了演示這個標注工具,我將使用手機錄制的圖像數(shù)據(jù)集,目標是對三種不同的USB接口類型進行分類:USB-A、USB-C、Micro USB和Mini USB。最初,所有圖像都未標注,存放在輸入目錄中。我們的標注工具應(yīng)該逐個展示圖像,并在指定類別后將其移動到相應(yīng)的目錄。
標注工具在運行中
指南
1. 先決條件
如果你想跟著操作,你應(yīng)該安裝opencv-python。你可以在項目倉庫的示例文件夾中找到一些示例圖像。
2. 數(shù)據(jù)加載
首先,我們從輸入文件夾中加載圖像。我們可以使用pathlib中的glob函數(shù)查找所有jpg圖像擴展名的文件。將結(jié)果傳遞給sorted函數(shù),確保圖像按順序處理。
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將圖像加載到數(shù)組中。然后顯示圖像并等待按鍵。通過在cv2.waitKey函數(shù)中將延遲設(shè)置為0,我們無限期地等待,直到按下某個鍵。然后確保按下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. 標注
讓我們在字符串列表中定義任務(wù)的標簽。在我的例子中,我有四個不同接口的標簽:
...
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)的標簽文件夾中。waitKey函數(shù)中的key變量是一個整數(shù),表示按下字符的Unicode代碼。要檢查按鍵是否是數(shù)字之一,我們需要使用ord函數(shù)將數(shù)字轉(zhuǎn)換為Unicode,類似于我們檢查按鍵q以關(guān)閉窗口的方式。該函數(shù)期望長度為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
要將圖像移動到輸出路徑中的分類標簽文件夾,我們可以使用pathlib中的/操作來連接路徑,然后使用rename函數(shù)將文件移動到目標位置。
...
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
在我們這樣做之前,我們需要確保目標文件夾存在。所以在循環(huán)之前,我們遍歷所有標簽并創(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的標簽鍵檢查方法是在循環(huán)之前創(chuàng)建鍵Unicode到標簽的映射。這樣,我們不需要在循環(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
讓我們還為鍵到標簽的映射添加一個小幫助文本。
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)建一個簡單的標注工具。我們可以對這個工具進行很多改進。我想進一步探索的一件事是添加不僅分類圖像的功能,還可以分割圖像并創(chuàng)建分割掩碼。
當然,有很多更復(fù)雜的工具可以簡化你的標注過程。然而,有時一個非常簡單的工具就是你所需要的,特別是在項目的早期階段進行探索性數(shù)據(jù)分析時,你需要一個快速的概念驗證。
【參考資料】
- 代碼鏈接: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