MTCNN网络之O网络——pytorch代码实现

MTCNN网络是多任务级联卷积神经网络，基于级联的特定目标检测器，在人脸识别中有着出色的表现。由P、R、O三个网络构成。常用于目标检测的算法还有SSD（单激多发探测器）、RCNN、YOLO等

如果对MTCNN的基础知识不太清楚，可以参考我的文章：

一、网络结构

MTCNN由三个子网络构成：P网络、R网络、O网络
P网络的输入是12*12，R网络输入是24*24，O网络输入是48*48。P网络注重从大量数据中筛选有效信息，R网络对P网络的筛选结果做二次筛选，O网络注重识别的精度。

（1）O网络（5层）

是三个网络中参数最多，网络层数最深的，因为到了最后一步，我们需要对目标进行更加精确的识别，所以相对神经网络的能力要求更高。

• 第一层：3*3卷积核，3*3最大池化层，输出特征图为：23*23*32
• 第二层：3*3卷积核，3*3最大池化层，输出特征图为：10*10*64
• 第三层：3*3卷积核，2*2最大池化层，输出特征图为：4*4*64
• 第四层：2*2卷积核，输出特征图为：3*3*128（这里我们仅需要置信度和四个坐标）
• 第五层：全连接层

代码实现：

import torch.nn as nn
import torch

class ONet(nn.Module):
    def __init__(self):
        super(ONet, self).__init__()
        self.layer = nn.Sequential(
            nn.Conv2d(3,32,3,1,0,bias=False),nn.BatchNorm2d(32),nn.ReLU(),
            nn.MaxPool2d(3,stride=2,padding=1),
            nn.Conv2d(32,64,3,1,0,bias=False),nn.BatchNorm2d(64),nn.ReLU(),
            nn.MaxPool2d(3,stride=2),
            nn.Conv2d(64,64,3,1,0,bias=False),nn.BatchNorm2d(64),nn.ReLU(),
            nn.MaxPool2d(2,stride=2),
            nn.Conv2d(64,128,2,1,bias=False),nn.BatchNorm2d(128),nn.ReLU()
        )
        self.fc1 = nn.Sequential(
            nn.Linear(128*3*3, 256),nn.ReLU(),
            nn.Linear(256,1)
        )

        self.fc2 = nn.Sequential(
            nn.Linear(128*3*3, 256),nn.ReLU(),
            nn.Linear(256,4)
        )

    def forward(self, x):
        y = self.layer(x)
        y = y.reshape(-1, 128*3*3)
        class_out = torch.sigmoid(self.fc1(y))
        bbox_out = self.fc2(y)
        return class_out, bbox_out
    
if __name__ == "__main__":
    net = ONet()
    x = torch.randn((1,3,48,48))
    y1, y2 = net(x)
    print(y1)
    print(y2)

O网络数据准备：

数据准备时应注意：

• 我们在准备数据集的时候，一定要让负样本的数量远多于正样本的数量和部分样本的数量，
• IOU阈值不能设置过小，也不能设置过大，否则会导致损失降不下去，准确率不够高。
• 数据量一定要够，数据量不够的话，一些不是人脸的物体容易被误识别
• 如果缺失亚洲人人脸，可以手动加入一些数据

"""
wider-face 生成48*48的训练数据
"""
import os
import sys
sys.path.append("./")
# print(os.getcwd())
import cv2
from PIL import Image
import numpy as np
from tqdm import tqdm
from utils.utils import IOU2 
import torch

TARGET_PATH = "./dataset/"
IMG_PATH = "./dataset/WIDER_train/images/"
DST_PATH = "./dataset/"
face_size = 48

label_file_path = TARGET_PATH + "wider_face_train.txt"

# 首先先生成对应的文件夹
if not os.path.exists(DST_PATH+f"{face_size}"):
    os.mkdir(DST_PATH+f"{face_size}")

# pos，part,neg裁剪图片放置位置
pos_save_dir = os.path.join(DST_PATH, f'{face_size}/positive')
part_save_dir = os.path.join(DST_PATH, f'{face_size}/part')
neg_save_dir = os.path.join(DST_PATH, f'{face_size}/negative')
# PNet数据地址
save_dir = os.path.join(DST_PATH, f'{face_size}/')

# 创建文件夹
if not os.path.exists(save_dir):
    os.mkdir(save_dir)
if not os.path.exists(pos_save_dir):
    os.mkdir(pos_save_dir)
if not os.path.exists(part_save_dir):
    os.mkdir(part_save_dir)
if not os.path.exists(neg_save_dir):
    os.mkdir(neg_save_dir)

# 生成后的数据列表文件
positive_file = open(os.path.join(save_dir, 'positive.txt'), 'w')
negative_file = open(os.path.join(save_dir, 'negative.txt'), 'w')
part_file = open(os.path.join(save_dir, 'part.txt'), 'w')

# 原数据集的列表文件
with open(label_file_path, 'r') as f:
    annotations = f.readlines()
num = len(annotations)
print('总共的图片数： %d' % num)

# 记录pos,neg,part三类生成数
positive_counter = 0
negative_counter = 0
part_counter = 0

# 记录读取的图片数
idx = 0
for anno in tqdm(annotations):
    anno = anno.strip().split(' ')
    img_name = anno[0]
    img_path = IMG_PATH+img_name+".jpg"
    # 获取人脸的box所有坐标
    box = list(map(float, anno[1:]))
    # 把所有box的坐标按照4分割，里面是左上角和右下角的信息
    boxes = np.array(box, dtype=np.float32).reshape(-1, 4)
    w = box[2] - box[0]
    h = box[3] - box[1]
    
    # 已经获取了坐标信息下面对图像进行操作
    with Image.open(img_path) as img:
        idx += 1
        img_w, img_h = img.size

        # 过滤字段
        if max(w,h)<40 or box[0]<0 or box[1]<0 or w<0 or h<0:  continue 
        # 标注不标准，给框适当的偏移量
        x1 = int(box[0] + w*0.12)
        y1 = int(box[1] + h*0.1)
        x2 = int(box[0] + w*0.9)
        y2 = int(box[1] + h*0.85)
        w = int(x2-x1)
        h = int(y2-y1)

        boxes = [[x1, y1, x2, y2]]
        cx = x1 + w/2
        cy = y1 + h/2

        # 查看框是否合适
        # img = cv2.imread(img_path)
        # img = cv2.rectangle(img, (x1,y1), (x2,y2),(0,0,255),2)
        # cv2.imshow("img",img)
        # cv2.waitKey(0)

        # 让正样本和部分样本随着图片中心进行随机偏移
        for _ in range(5):
            w_ = np.random.randint(-w*0.2, w*0.2)
            h_ = np.random.randint(-h*0.2, h*0.2)
            cx_ = cx + w_
            cy_ = cy + h_

            side_len = np.random.randint(round(min(w,h)*0.8), round(max(w,h)*1.25))
            x1_ = np.max(cx_-side_len/2, 0)
            y1_ = np.max(cy_-side_len/2, 0)
            x2_ = x1_ + side_len
            y2_ = y1_ + side_len

            crop_box = np.array([x1_, y1_, x2_, y2_])

            # 计算坐标偏移量
            offset_x1 = (x1 - x1_) / side_len # ▲δ = (x1-x1_)/side_len
            offset_y1 = (y1 - y1_) / side_len
            offset_x2 = (x2 - x2_) / side_len
            offset_y2 = (y2 - y2_) / side_len

            face_crop = img.crop(crop_box)
            face_resize = face_crop.resize((face_size, face_size), Image.Resampling.LANCZOS)

            iou = IOU2(torch.tensor(boxes), torch.tensor(crop_box)).item()

            if iou > 0.8: # 正样本
                positive_file.write(f"positive/{ positive_counter}.jpg 1 {offset_x1} {offset_y1} {offset_x2} {offset_y2}\n")
                positive_file.flush()
                face_resize.save(pos_save_dir+f"/{positive_counter}.jpg")
                positive_counter += 1
            elif iou>0.5: # 部分样本
                part_file.write(f"part/{ part_counter}.jpg 2 {offset_x1} {offset_y1} {offset_x2} {offset_y2}\n")
                part_file.flush()
                face_resize.save(part_save_dir+f"/{ part_counter}.jpg")
                part_counter += 1
            # elif iou < 0.1: # 负样本
            #     negative_file.write(f"negative/{ negative_counter}.jpg 0 0 0 0 0\n")
            #     negative_file.flush()
            #     face_resize.save(neg_save_dir+f"/{ negative_counter}.jpg")
            #     negative_counter += 1
        
        _boxes = torch.tensor(boxes[0])

        for _ in range(10): # 这里是为了能够生成足够的负样本
            side_len = np.random.randint(face_size, min(img_w, img_h)/2)
            x_ = np.random.randint(0,img_w - side_len)
            y_ = np.random.randint(0,img_h - side_len)
            crop_box = np.array([x_, y_, x_+side_len, y_+side_len])

            if IOU2(torch.tensor(crop_box),_boxes).item() < 0.05:
                face_crop = img.crop(crop_box)
                face_resize = face_crop.resize((face_size, face_size), Image.Resampling.LANCZOS)
                
                negative_file.write(f"negative/{ negative_counter}.jpg 0 0 0 0 0\n")
                negative_file.flush()
                face_resize.save(neg_save_dir+f"/{ negative_counter}.jpg")
                negative_counter += 1

O网络训练代码实现：

import torch 
import os 
import sys
sys.path.append("./")
import torch.nn as nn
from models.ONet import ONet
from utils.trainer import Trainer 
from utils.data import MTCNN_Dataset
from torch.utils.data import DataLoader, Dataset 

model_save_path = "./param/onet.pt"
data = MTCNN_Dataset("./dataset/48")
data_loader = DataLoader(data, shuffle=True, batch_size=128)
net = ONet()

trainer = Trainer(net, model_save_path,data_loader,50)
trainer.train()

import torch 
import torch.nn as nn
import os 
import sys
sys.path.append("./")
from utils.utils import accuracy

class Trainer():
    def __init__(self,net, param_path, data_loader, epoch):
        self.net = net
        self.param_path = param_path
        self.data_loader = data_loader 
        self.epoch = epoch
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        # self.classLoss = nn.BCELoss(reduction="mean")
        self.offsetLoss = nn.MSELoss(reduction="mean") 
        # self.optim = torch.optim.Adam(self.net.parameters(),lr=0.001,weight_decay=0.0001)
        self.optim = torch.optim.SGD(self.net.parameters(),lr=1e-5,weight_decay=0.0001) # 使用SGD进行微调
        self.scheduler = torch.optim.lr_scheduler.MultiStepLR(self.optim, milestones=[6, 14, 20], gamma=0.1)

    def train(self):
        if os.path.exists(self.param_path):
            self.net.load_state_dict(torch.load(self.param_path))
        self.net.to(self.device)
        self.net.train()

        for e in range(self.epoch):
            for b_i, (img, cat, offset) in enumerate(self.data_loader):
                if self.device == "cuda":
                    img = img.to(self.device)
                    cat = cat.to(self.device)
                    offset = offset.to(self.device)
                cat_out, offset_out = self.net(img)
                cat_out = cat_out.reshape(-1,1) # [1024,1]
                offset_out = offset_out.reshape(-1, 4) # [1024, 4]

                # 在计算损失的时候，置信度只需要正样本和负样本
                # 偏移量只需要部分样本和正样本
                category_mask = cat < 2
                category_select = torch.masked_select(cat, category_mask)
                category_out_select = torch.masked_select(cat_out, category_mask)
                pred = torch.clamp(category_out_select, min=1e-7, max=1-1e-7)
                cls_loss = - category_select * torch.log(pred) - (1-category_select) * torch.log(1-pred)
                cls_loss = torch.mean(cls_loss)
                # cls_loss = self.classLoss(category_select, category_out_select) # 这里会出现nan的大坑

                offset_mask = cat > 0 
                offset_index = torch.nonzero(offset_mask)[:,0]
                offset_select = offset[offset_index]
                offset_out_select = offset_out[offset_index]
                offset_loss = self.offsetLoss(offset_select, offset_out_select)
                loss = cls_loss+offset_loss

                # print(cls_loss, offset_loss)
                # 更新参数
                self.optim.zero_grad()
                loss.backward()
                self.optim.step()

                if b_i%100 == 0:
                    acc = accuracy(class_out=cat_out, label=cat)
                    print('Train epoch %d, batch %d, total_loss: %f, cls_loss: %f, box_loss: %f,'
                  'accuracy：%.2f' % (e, b_i, loss.item(), cls_loss.item(), offset_loss.item(), acc*100))
                    torch.save(self.net.state_dict(), self.param_path)

            self.scheduler.step()

import pickle
import shutil
import numpy as np
import random
import os
import cv2
from tqdm import tqdm
import torch

def IOU(box, boxes):
    """裁剪的box和图片所有人脸box的iou值
    参数：
      box：裁剪的box,当box维度为4时表示box左上右下坐标，维度为5时，最后一维为box的置信度
      boxes：图片所有人脸box,[n,4]
    返回值：
      iou值，[n,]
    """
    # box面积
    box_area = (box[2] - box[0] + 1) * (box[3] - box[1] + 1)
    # boxes面积,[n,]
    area = (boxes[:, 2] - boxes[:, 0] + 1) * (boxes[:, 3] - boxes[:, 1] + 1)
    # 重叠部分左上右下坐标
    xx1 = np.maximum(box[0], boxes[:, 0])
    yy1 = np.maximum(box[1], boxes[:, 1])
    xx2 = np.minimum(box[2], boxes[:, 2])
    yy2 = np.minimum(box[3], boxes[:, 3])

    # 重叠部分长宽
    w = np.maximum(0, xx2 - xx1 + 1)
    h = np.maximum(0, yy2 - yy1 + 1)
    # 重叠部分面积
    inter = w * h
    return inter / (box_area + area - inter + 1e-10)

def IOU2(box, other_boxes):
    #box: [x1,y1,x2,y2] 分别表示方框的左上角的点和右下角的点
    #other_boxs: N个box，多了一个维度(代表box的数量)
    if len(box.shape)==2: box = box.squeeze()
    if len(other_boxes.shape)==1: other_boxes = other_boxes.unsqueeze(dim=0)
    box_area = (box[2]-box[0])*(box[3]-box[1])
    other_boxes_area = (other_boxes[:,2]-other_boxes[:,0]) * (other_boxes[:,3]-other_boxes[:,1])

    #交集
    x1 = torch.max(box[0],other_boxes[:,0])
    y1 = torch.max(box[1],other_boxes[:,1])
    x2 = torch.min(box[2],other_boxes[:,2])
    y2 = torch.min(box[3],other_boxes[:,3])
    Min = torch.tensor([0])
    w,h = torch.max(Min,x2-x1),torch.max(Min,y2-y1) #如果没有相交的框，两者相减是负值，为防止此类事情出现
    # 交集的面积
    overlap_area = w*h

    # 交并比 = 交集 / 并集
    iou = overlap_area / (box_area+other_boxes_area-overlap_area)
    return iou

# 求训练时的准确率
def accuracy(class_out, label):
    # 查找neg 0 和pos 1所在的位置
    class_out = class_out.detach().cpu().numpy().round()
    label = label.detach().cpu().numpy()
    mask = class_out == label
    correct = np.sum(mask)
    return correct/len(class_out)

def IOU3(box, other_boxes, isMin=False):
    #box: [x1,y1,x2,y2] 分别表示方框的左上角的点和右下角的点
    #other_boxs: N个box，多了一个维度(代表box的数量)
    if len(box.shape)==2: box = box.squeeze()
    if len(other_boxes.shape)==1: other_boxes = other_boxes.unsqueeze(dim=0)
    box_area = (box[2]-box[0])*(box[3]-box[1])
    other_boxes_area = (other_boxes[:,2]-other_boxes[:,0]) * (other_boxes[:,3]-other_boxes[:,1])

    #交集
    x1 = torch.max(box[0],other_boxes[:,0])
    y1 = torch.max(box[1],other_boxes[:,1])
    x2 = torch.min(box[2],other_boxes[:,2])
    y2 = torch.min(box[3],other_boxes[:,3])
    Min = torch.tensor([0])
    w,h = torch.max(Min,x2-x1),torch.max(Min,y2-y1) #如果没有相交的框，两者相减是负值，为防止此类事情出现
    # 交集的面积
    overlap_area = w*h

    if isMin:
        # 交并比 = 交集 / 最小面积
        iou = overlap_area / torch.min(box_area,other_boxes_area)
    else:
        # 交并比 = 交集 / 并集
        iou = overlap_area / (box_area+other_boxes_area-overlap_area)
    return iou


def NMS(boxes, thre = 0.5, isMin=False):
    # boxes: [[置信度, x1,y1,x2,y2], [置信度, x1,y1,x2,y2], ...]
    # - 第一步：找出所有预测目标中得分最高的预测框作为基准框；
    # - 第二步：计算剩余的预测框与基准框的IOU，如果IOU>阈值t，则将这些预测框删除，因为，这些预测框与0.97的太相似了，所以是冗余框，需要抑制；
    # - 第三步：在去除冗余框后剩余预测框中，找到得分第二高的预测框0.95作为基准框，执行第二步操作；
    # - 第四步：重复上述过程，直至剩余预测框中的每个预测框都曾被用作过基准框。即剩余的预测框任意两个IOU均小于t，没有两个预测框过于相似，完成nms操作。
    res = []
    if len(boxes) == 0:
        return torch.tensor([])
    sort_boxes = boxes[boxes[:, 0].argsort(descending=True)] # 获得置信度降序排序的下标，并根据下标进行索引。

    # 防止传入错误的值
    if len(sort_boxes) == 0:
        print("ERROR: NO boxes! ")
        return []

    while len(sort_boxes) > 1:
        ref_box = sort_boxes[0]
        res.append(ref_box)

        other_boxes = sort_boxes[1:]
        sort_boxes = other_boxes[torch.where(IOU3(ref_box[1:], other_boxes[:,1:], isMin=isMin)<thre)]
    
    if len(sort_boxes)>0: res.append(sort_boxes[0])

    return torch.stack(res)

def convert2square(p_box):
    square_box = p_box.copy()
    if square_box.shape[0] == 0: return np.array([])
    h = p_box[:, 3] - p_box[:, 1]
    w = p_box[:, 2] - p_box[:, 0]
    maxside = np.maximum(h,w)
    square_box[:, 0] = p_box[:, 0] + w*0.5 - maxside*0.5
    square_box[:, 1] = p_box[:, 1] + w*0.5 - maxside*0.5
    square_box[:, 2] = square_box[:, 0] + maxside
    square_box[:, 3] = square_box[:, 1] + maxside
    return square_box

本文代码仅供参考，具体代码可参照：Pytorch-MTCNN: 基于Pytorch实现的MTCNN模型，人脸检测，人脸关键点检测。 (gitee.com)