幾行程式碼構建全功能的對象檢測模型，他是如何做到的？

• 2020 年 2 月 24 日
• 筆記

以下文章來源於AI科技大本營 ，作者Alan Bi

作者 | Alan Bi

譯者 | 武明利，責編 | Carol

出品 | AI科技大本營（ID:rgznai100）

如今，機器學習和電腦視覺已成為一種熱潮。我們都看過關於自動駕駛汽車和面部識別的新聞，可能會想像建立自己的電腦視覺模型有多酷。然而，進入這個領域並不總是那麼容易，尤其是在沒有很強的數學背景的情況下。如果你只想做一些小的實驗，像PyTorch和TensorFlow這樣的庫可能會很枯燥。

在本教程中，作者提供了一種簡單的方法，任何人都可以使用幾行程式碼構建全功能的對象檢測模型。更具體地說，我們將使用Detecto，這是一個在PyTorch之上構建的Python軟體包，可簡化該過程並向所有級別的程式設計師開放。

快速簡單的例子

為了演示如何簡單地使Detecto，讓我們載入一個預先訓練的模型，並對以下影像進行推斷：

首先，使用pip下載Detecto軟體包：

pip3 install detecto

然後，將上面的影像另存為「fruit.jpg」，並在與影像相同的文件夾中創建一個Python文件。在Python文件中，編寫以下5行程式碼：

from detectoimport core, utils, visualize

image = utils.read_image('fruit.jpg')

model = core.Model()

labels, boxes, scores = model.predict_top(image)

visualize.show_labeled_image(image, boxes, labels)

運行此文件後（如果你的電腦上沒有啟用CUDA的GPU，可能會花費幾秒鐘；稍後再進行介紹），你應該會看到類似下面的圖：

作者僅用了5行程式碼就完成了所有工作，真的是太棒了。下面是我們每步中分別做的：

1）導入Detecto模組

2）讀入影像

3）初始化預訓練模型

4）在影像上生成最高預測

5）為預測繪圖

繪製我們的預測

Detecto使用來自PyTorch模型動物園中的Faster R-CNN ResNet-50 FPN，它能夠檢測大約80種不同的物體，例如動物，車輛，廚房用具等。但是，如果你想要檢測自定義對象，例如可口可樂與百事可樂罐，斑馬與長頸鹿，該怎麼辦呢？

這時你會發現，在自定義數據集上訓練探測器模型同樣簡單; 同樣，你只需要5行程式碼，以及現有的數據集或花一些時間標記影像。

構建自定義數據集

在本教程中，作者將從頭開始構建自己的數據集。建議你也這樣做，但是如果你想跳過這一步，你可以在這裡下載一個示例數據集(從斯坦福的Dog數據集修改)。

對於我們的數據集，我們將訓練我們的模型來檢測來自RoboSub競賽的水下外星人，蝙蝠和女巫，如下所示：

理想情況下，每個類至少需要100張影像。好在每張影像中可以有多個對象，所以理論上，如果每張影像包含你想要檢測的每類對象，那麼你可以總共獲得100張影像。另外，如果你有影片素材，Detico可以輕鬆地將這些影片素材分割成可用於數據集的影像：

from detecto.utilsimport split_video

split_video('video.mp4','frames/', step_size=4)

上面的程式碼在「video.mp4」中每第4幀拍攝一次，並將其另存為JPEG文件存在「frames」文件夾中。

生成訓練數據集後，應該具有一個類似於以下內容的文件夾：

images/

|   image0.jpg

|   image1.jpg

|   image2.jpg

|   ...

如果需要的話，你還可以使用另一個文件夾，其中包含一組驗證影像。

現在是耗時的部分：標記。Detecto支援PASCAL VOC格式，其中具有XML文件，其中包含影像中每個對象的標籤和位置數據。要創建這些XML文件，可以使用開源LabelImg工具，如下所示：

pip3 install labelImg   # Download LabelImg using pip

labelImg                # Launch the application

現在，你應該會看到一個彈出窗口。單擊左側「打開目錄」按鈕，然後選擇想要標記的影像文件夾。如果一切正常，你應該會看到類似以下內容：

要繪製邊界框，請單擊左側菜單欄中的圖標（或使用鍵盤快捷鍵「w」）。然後，你可以在對象周圍拖動一個框並編寫/選擇標籤：

標記完影像後，請使用CTRL+S或CMD+S保存XML文件（為簡便起見，你可以使用自動填充的默認文件位置和名稱）。要標記下一張影像，請單擊「下一張影像」（或使用鍵盤快捷鍵「d」）。

整個數據集處理完畢之後，你的文件夾應如下所示：

images/

|   image0.jpg

|   image0.xml

|   image1.jpg

|   image1.xml

|   ...

我們已經準備好開始訓練我們的對象檢測模型了!

訪問GPU

首先，檢查你的電腦是否具有啟用CUDA的GPU。由於深度學習需要大量處理能力，因此在通常的CPU上進行訓練可能會非常緩慢。值得慶幸的是，大多數現代深度學習框架（例如PyTorch和Tensorflow）都可以在GPU上運行，從而使處理速度更快。確保已經下載了PyTorch（如果你安裝了Detecto，應該已經下載了），然後運行以下兩行程式碼：

import torch

print(torch.cuda.is_available())

如果列印True，那你可以跳到下一部分。如果顯示False，不要擔心。請按照以下步驟創建Google Colaboratory筆記型電腦，這是一個在線編碼環境，帶有免費可用的GPU。對於本教程，你將只在Google Drive文件夾中工作，而不是在電腦上工作。

1）登錄到Google Drive

2）創建一個名為「Detecto Tutorial」的文件夾並導航到該文件夾

3）將你的訓練影像（和/或驗證影像）上傳到此文件夾

4）右鍵單擊，轉到「更多」，然後單擊「Google Colaboratory」：

你現在應該看到這樣的介面：

5）根據需要給筆記型電腦起個名字，然後轉到「編輯」->「筆記型電腦設置」->「硬體加速器」，然後選擇「GPU」

6）輸入以下程式碼以「裝入」你的雲端硬碟，將目錄更改為當前文件夾，然後安裝Detecto：

import os

from google.colabimport drive

drive.mount('/content/drive')

os.chdir('/content/drive/My Drive/Detecto Tutorial')

!pip install detecto

為了確保一切正常，你可以創建一個新的程式碼單元，然後輸入!ls以檢查你是否處於正確的目錄中。

訓練自定義模型

最後，我們現在可以在自定義數據集上訓練模型了。如前所述，這是容易的部分。它只需要4行程式碼：

from detectoimport core, utils, visualize

dataset = core.Dataset('images/')

model = core.Model(['alien','bat','witch'])

model.fit(dataset)

讓我們再次分解一下我們每行程式碼所做的工作：

1、導入的Detecto模組

2、從「images」文件夾（包含我們的JPEG和XML文件）創建了一個數據集

3、初始化模型檢測自定義對象（外星人，蝙蝠和女巫）

4、在數據集上訓練我們的模型

根據數據集的大小，這可能需要10分鐘到1個小時以上的時間來運行，因此請確保你的程式在完成上述語句後不會立即退出（例如：你使用的是Jupyter / Colab筆記型電腦，它在活動時保留狀態）。

使用訓練好的模型

現在你已經有了訓練好的模型，讓我們在一些影像上對其進行測試。要從文件路徑讀取影像，可以使用detecto.utils模組中的read_image函數（也可以使用上面創建的數據集中的影像）：

# Specify the path to your image

image = utils.read_image('images/image0.jpg')

predictions = model.predict(image)

# predictions format: (labels, boxes, scores)

labels, boxes, scores = predictions

# ['alien', 'bat', 'bat']

print(labels)

#           xmin       ymin       xmax       ymax

# tensor([[ 569.2125,  203.6702, 1003.4383,  658.1044],

#         [ 276.2478,  144.0074,  579.6044,  508.7444],

#         [ 277.2929,  162.6719,  627.9399,  511.9841]])

print(boxes)

# tensor([0.9952, 0.9837, 0.5153])

print(scores)

正像你看到的，模型的預測方法返回一個由3個元素組成的元組：標籤，方框和分數。在上面的示例中，此模型在坐標[569、204、1003、658]（框[0]）處預測了一個外星人（標籤[0]），其置信度為0.995（得分[0]）。

根據這些預測，我們可以使用detecto.visualize模組繪製結果。例如：

visualize.show_labeled_image(image, boxes, labels)

將上面的程式碼與收到的影像和預測一起運行將產生如下所示的內容：

如果你有一個影片，你可以在它上面運行對象檢測：

visualize.detect_video(model,'input.mp4','output.avi')

這將獲取一個名為「input.mp4」的影片文件，並根據給定模型的預測結果生成一個「output.avi」文件。如果你使用VLC或其他影片播放器打開此文件，應該會看到一些希望看到的結果！

最後，你可以從文件中保存和載入模型，從而可以保存進度並稍後返回：

model.save('model_weights.pth')

# ... Later ...

model = core.Model.load('model_weights.pth', ['alien','bat','witch'])

高級用法

你會發現Detecto不僅限於5行程式碼。舉例來說，這個模型沒有你希望的那麼好。我們可以嘗試通過使用Torchvision轉換來擴展我們的數據集並定義一個自定義數據載入器來提高其性能：

from torchvisionimport transforms

augmentations = transforms.Compose([

    transforms.ToPILImage(),

    transforms.RandomHorizontalFlip(0.5),

    transforms.ColorJitter(saturation=0.5),

    transforms.ToTensor(),

    utils.normalize_transform(),

])

dataset = core.Dataset('images/', transform=augmentations)

loader = core.DataLoader(dataset, batch_size=2, shuffle=True)

此程式碼對數據集中的影像應用了隨機的水平翻轉和飽和效果，從而增加了數據的多樣性。然後，我們使用batch_size = 2定義一個數據載入對象；我們將其傳遞給model.fit而不是Dataset，這樣來告訴我們的模型是對2張影像進行批量訓練，而不是默認的1張。

如果你之前創建了單獨的驗證數據集，那麼現在是在訓練期間載入它的時候了。通過提供驗證數據集，fit方法將返回每個時期的損失列表，如果verbose = True，則會在訓練過程中將其列印出來。以下程式碼塊演示了這一點，並自定義了其他幾個訓練參數：

import matplotlib.pyplotas plt

val_dataset = core.Dataset('validation_images/')

losses = model.fit(loader, val_dataset, epochs=10, learning_rate=0.001,

                   lr_step_size=5, verbose=True)

plt.plot(losses)

plt.show()

損失的結果圖應或多或少地減少：

為了更具有靈活性和對模型的控制，你可以完全繞過Detecto。你可以根據需要隨意調整model.get_internal_model方法返回使用的基礎模型。

結論

在本教程中，作者展示了打造電腦視覺和對象檢測並沒有多大的挑戰性。你所需要的是一點時間和耐心來處理標記的的數集。

如果你對進一步探索感興趣的話，請查看Detecto on GitHub或訪問文檔以獲取更多教程和用例！

原文：

https://hackernoon.com/build-a-custom-trained-object-detection-model-with-5-lines-of-code-y08n33vi