[第十夜] 人臉偵測 : MTCNN 的實做 Part.2-3 -- 訓練&測試

15th鐵人賽

zivzhong

2023-09-25 23:59:31

369 瀏覽

分享至

前言

講完了資料的來源, 模型的建立以及怎麼處理 Model 的預測框之後，接著我們就要來開始訓練我們的 AI model 了！！我們需要做的事情有以下三者：

資料準備：我們一共須要不同的資料類型分別去訓練人臉分類器以及 landmark
跑訓練模型
使用訓練好的模型進行測試
那事不宜遲我們就開始今天的內容吧！

資料準備

我們已經下載下人臉資料了，內含足夠的資訊協助我們去訓練模型，但這樣就可以直接去訓練了嗎？我們舉一個例子，如果我們都只把有人臉的照片喂進分辨照片中是否含有人臉的分類器，那是不是我們的人臉分類器可能會傾向吳輪如何先猜有人臉，那這樣我們的 AI Model 就會失去作用了，導致的原因其實是因為 data umbalance ，亦即 Model 沒有看過沒有人臉的照片，那我們就需要額外提供這樣的照片才行！

資料分類

我們一共須要 4 種類行的資料：

Landmark 資料：有一張人臉照片搭配上對應的 Landmark label ，用以訓練 landmark 模型
Positive 人臉：擁有完整人臉的圖片配上對應的人臉框 label ，用以訓練有無人臉判斷模型
Part 人臉：擁有部份人臉的照片配上對應的人臉框 label，超出圖片的框記得上限要降回圖片範圍內，用以訓練有無人臉判斷模型
Negative 人臉：基本上沒有人臉的部份，人臉框全部標記 0 0 0 0即可，用以訓練有無人臉判斷模型
以上這四種資料中第一種跟第二種資料算好解決，那剩下兩種要怎麼辦？沒事其實我們可以多加利用原始的資料集。舉個例子，WiderFace 有一張資料跟人臉框這得向下圖這樣：

我們其實可 random 去切這張照片，例如下圖籃框部份就是 Part 而紅框部份就是 Negative，用以分別的基準恰好可以使用我們昨天介紹的 IOU function， ex. IOU < 0.3 是 Negative， 0.3~ 0.7 是 Part !

大家可以進行去切，後我們也會放上我們切好的給大家參考！

訓練

實現一個完整的MTCNN（Multi-task Cascaded Convolutional Networks）人臉偵測器包括三個部分：PNet、RNet和ONet。這些網絡是連接的，每個網絡都有不同的任務：檢測（detection）、對齊（alignment）和特徵提取（feature extraction）。在這個示例中，我們將使用之前介紹的 Model 定義去建立出 PNet，RNet以及 ONet，然後使用建立好的資料集去訓練和測試。

以下是完整的PNet的訓練和測試示例，包括訓練數據的準備和損失函數的定義：

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader

# 定義PNet模型，與前兩天示例相同
class PNet(nn.Module):
    def __init__(self):
        super(PNet, self).__init__()
        # ...（與前兩天示例相同的模型定義）

# 載入訓練數據集
transform = transforms.Compose([transforms.Resize((12, 12)), transforms.ToTensor()])
train_dataset = datasets.ImageFolder(root='path_to_train_data', transform=transform)
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)

# 初始化PNet模型
pnet = PNet()

# 定義損失函數和優化器
criterion_cls = nn.BCELoss()  # 二元交叉熵損失，用於人臉/非人臉分類
criterion_offset = nn.MSELoss()  # 均方誤差損失，用於人臉位置偏移預測
optimizer = optim.Adam(pnet.parameters(), lr=0.001)

# 訓練PNet模型
num_epochs = 10
for epoch in range(num_epochs):
    pnet.train()  # 設置為訓練模式
    running_loss = 0.0
    for inputs, labels in train_loader:
        optimizer.zero_grad()
        out_cls, out_offset = pnet(inputs)
        # 計算分類損失
        loss_cls = criterion_cls(out_cls, labels)
        # 計算位置偏移損失
        loss_offset = criterion_offset(out_offset, ground_truth_offset)
        # 總損失 = 分類損失 + 0.5 * 位置偏移損失（根據MTCNN論文）
        loss = loss_cls + 0.5 * loss_offset
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
    print(f'Epoch [{epoch + 1}/{num_epochs}], Loss: {running_loss / len(train_loader)}')

# 測試PNet模型
pnet.eval()  # 設置為評估模式
# 在測試數據集上進行人臉檢測，並繪製結果
# 這部分需要另外的代碼來實現，包括後處理和圖像可視化

上述代碼展示了PNet的訓練過程，包括訓練數據的載入、模型的初始化、損失函數的定義、優化器的設置以及每個epoch的訓練循環。而 RNet 跟 ONet 的過程如下：

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader

# 定義RNet模型，與前兩天示例相同
class RNet(nn.Module):
    def __init__(self):
        super(RNet, self).__init__()
        # ...（與前兩天示例相同的模型定義）

# 載入訓練數據集
transform = transforms.Compose([transforms.Resize((24, 24)), transforms.ToTensor()])
train_dataset = datasets.ImageFolder(root='path_to_train_data', transform=transform)
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)

# 初始化RNet模型
rnet = RNet()

# 定義損失函數和優化器
criterion_cls = nn.BCELoss()  # 二元交叉熵損失，用於人臉/非人臉分類
criterion_offset = nn.MSELoss()  # 均方誤差損失，用於人臉位置偏移預測
optimizer = optim.Adam(rnet.parameters(), lr=0.001)

# 訓練RNet模型
num_epochs = 10
for epoch in range(num_epochs):
    rnet.train()  # 設置為訓練模式
    running_loss = 0.0
    for inputs, labels in train_loader:
        optimizer.zero_grad()
        out_cls, out_offset = rnet(inputs)
        # 計算分類損失
        loss_cls = criterion_cls(out_cls, labels)
        # 計算位置偏移損失
        loss_offset = criterion_offset(out_offset, ground_truth_offset)
        # 總損失 = 分類損失 + 0.5 * 位置偏移損失（根據MTCNN論文）
        loss = loss_cls + 0.5 * loss_offset
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
    print(f'Epoch [{epoch + 1}/{num_epochs}], Loss: {running_loss / len(train_loader)}')

# 測試RNet模型
rnet.eval()  # 設置為評估模式
# 在測試數據集上進行人臉檢測，並繪製結果
# 這部分需要另外的代碼來實現，包括後處理和圖像可視化

ONet 如下：

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader

# 定義ONet模型，與前兩天示例相同
class ONet(nn.Module):
    def __init__(self):
        super(ONet, self).__init__()
        # ...（與前兩天示例相同的模型定義）

# 載入訓練數據集
transform = transforms.Compose([transforms.Resize((48, 48)), transforms.ToTensor()])
train_dataset = datasets.ImageFolder(root='path_to_train_data', transform=transform)
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)

# 初始化ONet模型
onet = ONet()

# 定義損失函數和優化器
criterion_cls = nn.BCELoss()  # 二元交叉熵損失，用於人臉/非人臉分類
criterion_offset = nn.MSELoss()  # 均方誤差損失，用於人臉位置偏移預測
criterion_landmarks = nn.MSELoss()  # 均方誤差損失，用於特徵點位置預測
optimizer = optim.Adam(onet.parameters(), lr=0.001)

# 訓練ONet模型
num_epochs = 10
for epoch in range(num_epochs):
    onet.train()  # 設置為訓練模式
    running_loss = 0.0
    for inputs, labels in train_loader:
        optimizer.zero_grad()
        out_cls, out_offset, out_landmarks = onet(inputs)
        # 計算分類損失
        loss_cls = criterion_cls(out_cls, labels)
        # 計算位置偏移損失
        loss_offset = criterion_offset(out_offset, ground_truth_offset)
        # 計算特徵點位置損失
        loss_landmarks = criterion_landmarks(out_landmarks, ground_truth_landmarks)
        # 總損失 = 分類損失 + 0.5 * 位置偏移損失 + 0.5 * 特徵點位置損失（根據MTCNN論文）
        loss = loss_cls + 0.5 * loss_offset + 0.5 * loss_landmarks
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
    print(f'Epoch [{epoch + 1}/{num_epochs}], Loss: {running_loss / len(train_loader)}')

# 測試ONet模型
onet.eval()  # 設置為評估模式
# 在測試數據集上進行人臉檢測，並繪製結果
# 這部分需要另外的代碼來實現，包括後處理和圖像可視化

測試

那訓練完這三個模型之後，我們就可以來測試看看我們模型囉!那這三個模型是級聯關系，意即 PNet 執行完的結果是可以接到 R-Net，R-Net 的結果可以接到 O-Net 的。五們其實之前已經給大家看過是如何接的囉，在[第七夜] 人臉偵測 : MTCNN 的實做 Part.1 簡單用 OpenCV 體驗。這裡如果熟悉 Python 的人可以直接去把 OpenCV 示例中的 model 部分換成 Pytorch 即可! 那我們這裡還是大致釋例一下 Pytroch model 的測試環節，我們以 P-Net 為例:

import torch
import torch.nn as nn
import torch.nn.functional as F
import cv2
import numpy as np

# 定義P-Net模型，詳見兩天前的定義模型!
class PNet(nn.Module):
    def __init__(self):
        super(PNet, self).__init__()
        # 定義網絡結構，這裡省略了具體結構
        # 通常包括多個卷積層和全連接層
        # ...

    def forward(self, x):
        # 定義前向傳播過程
        # ...

# 加載已經訓練好的P-Net權重
pnet = PNet()
pnet.load_state_dict(torch.load('pnet_weights.pth'))
pnet.eval()  # 設置為評估模式

# 讀取輸入圖像
image = cv2.imread('test_image.jpg')

# 預處理輸入圖像
# 這裡需要根據P-Net的預處理要求來進行圖像預處理
# 通常包括尺度變換、正規化等操作
# 預處理後的圖像應該是一個Tensor

# 將預處理後的圖像轉換為PyTorch的Tensor
image_tensor = torch.from_numpy(image.transpose((2, 0, 1))).float()

# 進行推理
with torch.no_grad():
    output = pnet(image_tensor.unsqueeze(0))  # 添加批次維度

# 解析檢測結果
# 通常需要根據模型的輸出格式來解析檢測結果
# 例如，P-Net可能輸出人臉的位置和信心分數
# 解析過程需要根據模型的具體輸出結構來實現

# 獲取人臉位置和信心分數
# 這裡需要根據P-Net的輸出結構來獲取位置和分數
# 假設output包含'tensor_boxes'和'tensor_scores'兩個Tensor
boxes = output['tensor_boxes']
scores = output['tensor_scores']

# 後處理檢測結果
# 通常需要進行非極大值抑制（NMS）等後處理操作
# 後處理的具體步驟取決於模型的應用和任務
# 在這個示例中，我們使用Soft NMS來獲取最終的檢測結果

def soft_nms(boxes, scores, threshold=0.5, sigma=0.5):
    # 實現Soft NMS的代碼，這裡需要根據具體需求實現，詳可見昨天的 NMS 那一章
    # ...

# 使用Soft NMS獲取最終的檢測結果
selected_boxes, selected_scores = soft_nms(boxes, scores)

# 將選擇的框繪製到原始圖像上
for box in selected_boxes:
    x1, y1, x2, y2 = box
    cv2.rectangle(image, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 2)

# 保存結果圖像
cv2.imwrite('result_image.jpg', image)

這個示例程式中，我們首先定義了P-Net的網絡結構，然後加載了已經訓練好的權重。接著，我們讀取輸入圖像並進行預處理，然後使用P-Net進行推理，獲得人臉位置和信心分數。最後，我們進行後處理（Soft NMS）以獲得最終的檢測結果，並將選擇的框繪製到原始圖像上，最終保存結果圖像。請注意，這只是一個簡單的示例，MTCNN 完整流程會將 P-Net 的結果餵進 R-Net 去做調整，然後 R-Net 的結果會餵進 O-Net 做最終預測!