幾行代碼構建全功能的對象檢測模型,他是如何做到的?
如今,機器學習和計算機視覺已成為一種熱潮。我們都看過關于自動駕駛汽車和面部識別的新聞,可能會想象建立自己的計算機視覺模型有多酷。然而,進入這個領域并不總是那幺容易,尤其是在沒有很強的數學背景的情況下。如果你只想做一些小的實驗,像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或訪問文檔以獲取更多教程和用例!