Pytorch Mask RCNN訓練自定義數據集

文章目錄

前言（必讀）

最近做目標檢測，然後記錄一下 Faster RCNN、Mask RCNN來做目標檢測踩得那些坑。

首先，本文並不是利用Pytorch從頭去實現Faster RCNN、Mask RCNN這兩個結構的文章。如果有意向去從頭實現並瞭解每一步細節可以看看下面這些視頻和博客：

來自B站的兩位大佬講解

大佬一：視頻博客 GitHub
大佬二：視頻博客 GitHub

上面都是利用pytorch從原理到具體實現的步驟。

不過本文主要還是利用Pytorch中的Torchvision.model中的Faster RCNN、Mask RCNN來實現遷移學習。

關於如何利用遷移學習來訓練自己的數據集，這裏也給出兩個超讚的教程：

看完以上三個教程，基本上利用Pytorch中的Torchvision.model中的Faster RCNN、Mask RCNN來實現遷移學習也基本上沒問題了。

下面介紹採坑：

1. 訓練自己的數據集

這裏講一下如何製作自己的數據集訓練這兩個網絡：

1.1 準備圖片

首先你得準備圖片數據，這個數據可能是別人給你提供的，也可能是你自己下載的，反正你得先準備好你的圖片並放在一個指定的文件夾裏。
這裏推薦一個批量下載網上圖片的工具：IMAGE CYBORG

1.2 數據標註工具（labelme）

現在一般用於目標檢測、實力分割的數據集製作的有兩個工具：labelme和labeling。至於這兩個有什麼區別和聯繫呢，有如下解釋：labelimg用於做目標檢測，佔內存小，labelme用於圖像語義/實例分割，佔內存大。

	Labelme	LabelImg
打開文件	Open	Open
打開文件夾	OpenDir	OpenDir
前後幀變換	Next & Prev Image	Next & Prev Image
標註方式	聚點分割&矩形框&圓&線	矩形框
保存方式	JSON	VOC&YOLO格式
標註大小	內存佔比大	忽略不計
優點	標註類型多（分割和檢測）標註文件另存了原始圖像	存儲簡單，RCNN、SSD、YOLO指定標註格式，對象，明確，模式靈活
適應場景	大部分2D分割任務,少數據量2D檢測任務	單目標的2D檢測任務,適用於VOC格式的數據

1.3 如何安裝labelme

windows安裝

conda install pyqt  #conda現在已經自帶pyqt5了，所以會提示你已安裝
pip install labelme

使用
在終端中執行以下命令：labelme就可以使用了，具體如何使用百度吧。

1.4 製作數據集

將new_json_to_dataset.py，draw.py兩個文件複製到下面labelme的路徑下C:\Users\lee\Anaconda3\Lib\site-packages\labelme\cli
修改new_json_to_dataset.py中的你自己的類別名(一定得修改，不然會報錯)：

NAME_LABEL_MAP = {
    '_background_': 0,
    "che": 1,
    "biao": 2

}

LABEL_NAME_MAP = ['0: _background_',
                  '1: che',
                  '2: biao']

在當前路徑下打開命令行窗口，輸入：python ./new_json_to_dataset.py C:\TorchVision_Maskrcnn\Maskrcnn_LPR\labelme
其中C:\TorchVision_Maskrcnn\Maskrcnn_LPR\labelme是你存放圖片的路徑
將copy.py複製到剛剛的C:\TorchVision_Maskrcnn\Maskrcnn_LPR\labelme存放圖片的路徑，運行 python ./copy.py將生成gts（mask）和png（原始圖片）兩個文件夾
將gts（mask）裏面的所有文件複製到PedMasks文件夾，png（原始圖片）裏面的所有文件複製到PNGImages文件夾

1.6 安裝 pycocotools

在整個過程中你要安裝 pycocotools，主要用到其中的IOU計算的庫來評價模型的性能。但是這個安裝並不簡單，看了很多文章相對來說只有這篇文章最符合。（Windows下安裝 pycocotools）

2. 定義 Faster RCNN、Mask RCNN 模型

前言

爲什麼要說這個呢？因爲如果你不是很有錢，或者公司有點摳買不起一張8G以上的顯卡，不改動就訓練這兩個網絡你基本上不可能成功。懂？財大氣粗可以忽略……

因爲本人就用的普通顯卡（GTX1660，6G內存），訓練Faster RCNN、Mask RCNN 這兩個網絡不要想着使用多GPU運行，我看了GitHub說了在windows上Faster RCNN、Mask RCNN暫時不支持多GPU運行。

幸運的是，在改動一些參數之後就可以完美運行。

Mask R-CNN是基於Faster R-CNN改造而來的。Faster R-CNN用於預測圖像中潛在的目標框和分類得分.

而Mask R-CNN在此基礎上加了一個額外的分支，用於預測每個實例的分割mask。

有兩種方式來修改torchvision modelzoo中的模型，以達到預期的目的。第一種，採用預訓練的模型，在修改網絡最後一層後finetune。第二種，根據需要替換掉模型中的骨幹網絡，如將ResNet替換成MobileNet等。

2.1 微調一個預訓練好的Faster RCNN模型

假設你想從一個在COCO上預先訓練過的模型開始，並想針對你的特定類對它進行微調。下面有一種可行的方法：

import torchvision
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
 
# 在COCO上加載經過預訓練的預訓練模型
model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
#這個操作你是真的要有固定參數
for param in model.parameters():
    param.requires_grad = False
    
# 將分類器替換爲具有用戶定義的 num_classes的新分類器
num_classes = 2  # 1 class (person) + background

# 獲取分類器的輸入參數的數量
in_features = model.roi_heads.box_predictor.cls_score.in_features
# 用新的頭部替換預先訓練好的頭部
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)

2.2 替換 Faster RCNN 模型的骨幹網絡

場景：替換掉模型的骨幹網絡。舉例來說，默認的骨幹網絡（ResNet-50）對於某些應用來說可能參數過多不易部署，可以考慮將其替換成更輕量的網絡（如MobileNet）。

import torchvision
from torchvision.models.detection import FasterRCNN
from torchvision.models.detection.rpn import AnchorGenerator
 
# 加載用於分類的預先訓練的模型並僅返回features
backbone = torchvision.models.mobilenet_v2(pretrained=True).features
# FasterRCNN需要知道骨幹網中的輸出通道數量。對於mobilenet_v2，它是1280，所以我們需要在這裏添加它
backbone.out_channels = 1280
 
# 我們讓RPN在每個空間位置生成5 x 3個Anchors(具有5種不同的大小和3種不同的寬高比)
# 我們有一個元組[元組[int]],因爲每個特徵映射可能具有不同的大小和寬高比
anchor_generator = AnchorGenerator(sizes=((32, 64, 128, 256, 512),),
                                   aspect_ratios=((0.5, 1.0, 2.0),))
 
# 定義一下我們將用於執行感興趣區域裁剪的特徵映射，以及重新縮放後裁剪的大小。
# 如果您的主幹返回Tensor，則featmap_names應爲[0]。
# 更一般地，主幹應該返回OrderedDict [Tensor]
# 並且在featmap_names中，您可以選擇要使用的功能映射。
roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=[0],#featmap_names=['0']
                                                output_size=7,
                                                sampling_ratio=2)
# 將這些pieces放在FasterRCNN模型中
model = FasterRCNN(backbone,
                   num_classes=2,
                   rpn_anchor_generator=anchor_generator,
                   box_roi_pool=roi_pooler)

2.3 微調一個預訓練好的Mask RCNN模型（本文使用）

import torchvision
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
from torchvision.models.detection.mask_rcnn import MaskRCNNPredictor
 
# load an instance segmentation model pre-trained on COCO
model = torchvision.models.detection.maskrcnn_resnet50_fpn(pretrained=True)
for param in model.parameters():
    param.requires_grad = False

num_classes = 3
# get the number of input features for the classifier
in_features = model.roi_heads.box_predictor.cls_score.in_features

# replace the pre-trained head with a new one
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)

# now get the number of input features for the mask classifier
in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels
hidden_layer = 256

# and replace the mask predictor with a new one
model.roi_heads.mask_predictor = MaskRCNNPredictor(in_features_mask,
                                                   hidden_layer,
                                                   num_classes)

2.4 替換 Mask RCNN 模型的骨幹網絡

import torchvision
from torchvision.models.detection import FasterRCNN
from torchvision.models.detection.rpn import AnchorGenerator
 
# 加載用於分類的預先訓練的模型並僅返回features
backbone = torchvision.models.mobilenet_v2(pretrained=True).features
# FasterRCNN需要知道骨幹網中的輸出通道數量。對於mobilenet_v2，它是1280，所以我們需要在這裏添加它
backbone.out_channels = 1280
 
# 我們讓RPN在每個空間位置生成5 x 3個Anchors(具有5種不同的大小和3種不同的寬高比)
# 我們有一個元組[元組[int]],因爲每個特徵映射可能具有不同的大小和寬高比
anchor_generator = AnchorGenerator(sizes=((32, 64, 128, 256, 512),),
                                   aspect_ratios=((0.5, 1.0, 2.0),))
 
# 定義一下我們將用於執行感興趣區域裁剪的特徵映射，以及重新縮放後裁剪的大小。
# 如果您的主幹返回Tensor，則featmap_names應爲[0]。
# 更一般地，主幹應該返回OrderedDict [Tensor]
# 並且在featmap_names中，您可以選擇要使用的功能映射。
roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=[0],
                                                output_size=7,
                                                sampling_ratio=2)
mask_roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=[0],
                                                     output_size=14,
                                                     sampling_ratio=2)

model = MaskRCNN(backbone,
                 num_classes=2,
                 rpn_anchor_generator=anchor_generator,
                 box_roi_pool=roi_pooler,
                 mask_roi_pool=mask_roi_pooler)

3. 訓練過程

具體代碼看我GitHub

訓練過程總的來說沒什麼，平平無奇，但還是總結一下幾個我遇到的坑。

batch_size不宜過大，過大GPU喫不消
num_workers設置爲0，我也不知道我設置成其他數會報錯
學習率lr不宜設置太小，可以循序漸進
不一定要使用lr_scheduler，可以一直保持一個固定的學習率（我的0.01）由於我內存不夠我給他禁了，不過相應engine.py的地方也得修改
Mask RCNN 好像暫時不支持多GPU運行，（會的小夥伴下方請留言）

此外補充一個我在訓練時發生的一個BUG：

TypeError: object of type <class 'numpy.float64'> cannot be safely interpreted as an integer.

解決辦法：

找到cocoeval.py文件，大概507行。修改self.iouThrs = np.linspace(.5, 0.95, np.round((0.95 - .5) / .05) + 1, endpoint=True)爲self.iouThrs = np.linspace(.5, 0.95, int(np.round((0.95 - .5) / .05) + 1), endpoint=True)修改self.recThrs = np.linspace(.0, 1.00, np.round((1.00 - .0) / .01) + 1, endpoint=True)爲self.recThrs = np.linspace(.0, 1.00, int(np.round((1.00 - .0) / .01) + 1), endpoint=True)

4. 測試

import os
import numpy as np
import torch
from PIL import Image

import torchvision
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
from torchvision.models.detection.mask_rcnn import MaskRCNNPredictor

import sys
sys.path.append("./detection")
from engine import train_one_epoch, evaluate
import utils
import transforms as T
import cv2
import cv2_util

import random
import time
import datetime

def random_color():
    b = random.randint(0,255)
    g = random.randint(0,255)
    r = random.randint(0,255)
 
    return (b,g,r)


def toTensor(img):
    assert type(img) == np.ndarray,'the img type is {}, but ndarry expected'.format(type(img))
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    img = torch.from_numpy(img.transpose((2, 0, 1)))
    return img.float().div(255)  # 255也可以改爲256

def PredictImg( image, model,device):
    #img, _ = dataset_test[0] 
    img = cv2.imread(image)
    result = img.copy()
    dst=img.copy()
    img=toTensor(img)

    names = {'0': 'background', '1': 'che', '2': 'biao'}
    # put the model in evaluati
    # on mode

    prediction = model([img.to(device)])

    boxes = prediction[0]['boxes']
    labels = prediction[0]['labels']
    scores = prediction[0]['scores']
    masks=prediction[0]['masks']

    m_bOK=False;
    for idx in range(boxes.shape[0]):
        if scores[idx] >= 0.8:
            m_bOK=True
            color=random_color()
            mask=masks[idx, 0].mul(255).byte().cpu().numpy()
            thresh = mask
            contours, hierarchy = cv2_util.findContours(
                thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE
            )
            cv2.drawContours(dst, contours, -1, color, -1)

            x1, y1, x2, y2 = boxes[idx][0], boxes[idx][1], boxes[idx][2], boxes[idx][3]
            name = names.get(str(labels[idx].item()))
            cv2.rectangle(result,(x1,y1),(x2,y2),color,thickness=2)
            cv2.putText(result, text=name, org=(x1, y1+10), fontFace=cv2.FONT_HERSHEY_SIMPLEX, 
                fontScale=0.5, thickness=1, lineType=cv2.LINE_AA, color=color)

            dst1=cv2.addWeighted(result,0.7,dst,0.5,0)

            
    if m_bOK:
        cv2.imshow('result',dst1)
        cv2.waitKey()
        cv2.destroyAllWindows()

if __name__ == "__main__":

    device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
    num_classes = 3
    model = torchvision.models.detection.maskrcnn_resnet50_fpn(pretrained=False,num_classes=num_classes)
    model.to(device)
    model.eval()
    save = torch.load('model.pth')
    model.load_state_dict(save['model'])
    start_time = time.time()
    PredictImg('test/11.jpg',model,device)
    total_time = time.time() - start_time
    total_time_str = str(datetime.timedelta(seconds=int(total_time)))
    print('Training time {}'.format(total_time_str))
    print(total_time)