人脸识别和MTCNN模型 原创

import torch
from torch import nn

"""
    P-Net
"""

classPNet(nn.Module):
def__init__(self):
super().__init__()
        self.features_extractor = nn.Sequential(
# 第一层卷积
            nn.Conv2d(in_channels=3, out_channels=10, kernel_size=3, stride=1, padding=0),
            nn.BatchNorm2d(num_features=10),
            nn.ReLU(),

# 第一层池化
            nn.MaxPool2d(kernel_size=3,stride=2, padding=1),

# 第二层卷积
            nn.Conv2d(in_channels=10, out_channels=16, kernel_size=3, stride=1, padding=0),
            nn.BatchNorm2d(num_features=16),
            nn.ReLU(),

# 第三层卷积
            nn.Conv2d(in_channels=16, out_channels=32, kernel_size=3, stride=1, padding=0),
            nn.BatchNorm2d(num_features=32),
            nn.ReLU()
)

# 概率输出
        self.cls_out = nn.Conv2d(in_channels=32, out_channels=2, kernel_size=1, stride=1, padding=0)
# 回归量输出
        self.reg_out = nn.Conv2d(in_channels=32, out_channels=4, kernel_size=1, stride=1, padding=0)


defforward(self, x):
print(x.shape)
        x = self.features_extractor(x)
        cls_out = self.cls_out(x)
        reg_out = self.reg_out(x)

return cls_out, reg_out

import torch
from torch import nn

classRNet(nn.Module):
def__init__(self):
super().__init__()

        self.feature_extractor = nn.Sequential(
# 第一层卷积 24 x 24
            nn.Conv2d(in_channels=3, out_channels=28, kernel_size=3, stride=1, padding=0),
            nn.BatchNorm2d(num_features=28),
            nn.ReLU(),

# 第一层池化 11 x 11
            nn.MaxPool2d(kernel_size=3, stride=2, padding=1, ceil_mode=False),

# 第二层卷积 9 x 9
            nn.Conv2d(in_channels=28, out_channels=48, kernel_size=3, stride=1, padding=0),
            nn.BatchNorm2d(num_features=48),
            nn.ReLU(),

# 第二层池化 (没有补零) 4 x 4
            nn.MaxPool2d(kernel_size=3, stride=2, padding=0, ceil_mode=False),

# 第三层卷积 3 x 3
            nn.Conv2d(in_channels=48, out_channels=64, kernel_size=2, stride=1, padding=0),
            nn.BatchNorm2d(num_features=64),
            nn.ReLU(),

# 展平
            nn.Flatten(),

# 全连接层 [batch_size, 128]
            nn.Linear(in_features=3*3*64, out_features=128)
)

# 概率输出
        self.cls_out = nn.Linear(in_features=128, out_features=1)

# 回归量输出
        self.reg_out = nn.Linear(in_features=128, out_features=4)

defforward(self, x):
        x = self.feature_extractor(x)
        cls = self.cls_out(x)
        reg = self.reg_out(x)
return cls, reg

import torch
from torch import nn


classONet(nn.Module):
def__init__(self):
super().__init__()
        self.feature_extractor = nn.Sequential(
# 第1层卷积 48 x 48
            nn.Conv2d(in_channels=3, out_channels=32, kernel_size=3, stride=1, padding=0),
            nn.BatchNorm2d(num_features=32),
            nn.ReLU(),

# 第1层池化 11 x 11
            nn.MaxPool2d(kernel_size=3, stride=2, padding=1, ceil_mode=False),

# 第2层卷积 9 x 9
            nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, stride=1, padding=0),
            nn.BatchNorm2d(num_features=64),
            nn.ReLU(),

# 第2层池化 
            nn.MaxPool2d(kernel_size=3, stride=2, padding=0, ceil_mode=False),

# 第3层卷积 
            nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, stride=1, padding=0),
            nn.BatchNorm2d(num_features=64),
            nn.ReLU(),

# 第3层池化 
            nn.MaxPool2d(kernel_size=2, stride=2, padding=0, ceil_mode=False),

# 第4层卷积
            nn.Conv2d(in_channels=64, out_channels=128, kernel_size=2, stride=1, padding=0),
            nn.BatchNorm2d(num_features=128),
            nn.ReLU(),


# 展平 [batch_size, n_features]
            nn.Flatten(),

# 全连接 [batch_size, 128]
            nn.Linear(in_features=3*3*128, out_features=256)

)

# 概率输出
        self.cls_out = nn.Linear(in_features=256, out_features=1)

# 回归量输出
        self.reg_out = nn.Linear(in_features=256, out_features=4)

# 关键点输出
        self.landmark_out = nn.Linear(in_features=256, out_features=10)


defforward(self, x):
        x = self.feature_extractor(x)
        cls = self.cls_out(x)
        reg = self.reg_out(x)
        landmark = self.landmark_out(x)
return cls, reg, landmark

CelebA
  |-Anno
|-identity_CelebA.txt # 图片标注的身份信息
|-list_attr_celeba.txt # 图片标注的属性信息
|-list_bbox_celeba.txt # 图片标注的人脸框
|-list_landmarks_align_celeba.txt # 图片标注的人脸关键点(对齐图片)
|-list_landmarks_celeba.txt # 图片标注的人脸关键点
|-Eval
|-list_eval_partition.txt # 图片划分训练集、验证集、测试集
|-Img
|-img_align_celeba # 图片经过对齐后的
|-img_celeba # 原始图片(未做对齐处理)
|-img_celeba_png # PNG格式的图片

代码根目录
  |-datasets
|-celeba
|-identity_CelebA.txt
|-list_attr_celeba.txt
|-list_bbox_celeba.txt
|-list_landmarks_align_celeba.txt
|-list_landmarks_celeba.txt
|-Img
|-img_celeba
|-*.py 代码文件

import numpy as np
from torchvision.datasets importCelebA
import matplotlib.pyplot as plt
import cv2
from PIL importImage


root_dir ='datasets'# jupyter notebook的文件放在与datasets同一级目录下
celeba =CelebA(root=root_dir, split='train',
                target_type=['attr','identity','bbox','landmarks'],
                download=False)
attr_names = celeba.attr_names
attr_names.pop()
attr_names = np.array(attr_names)
fig = plt.figure(figsize=(14,7))

n =4# 显示的图片数量
for idx inrange(n):
    img,(attr, identity, bbox, landmarks)= celeba[idx]# 读取图片和相关标签值
    ax = fig.add_subplot(1,4, idx +1)
# 不显示刻度标签和边框
    ax.set_xticks([])
    ax.set_yticks([])
    ax.set_frame_on(b=False)

# 将CelebA数据集读取到的PIL图片格式转换成OprnCV所需格式
    img_cv2 = cv2.cvtColor(src=np.asanyarray(img), code=cv2.COLOR_RGB2BGR)

# 绘制特征点
    landmarks = landmarks.numpy()
for idx, point inenumerate(landmarks):
if idx %2==0:
            cv2.circle(img=img_cv2, center=(point, landmarks[idx +1]),
                        radius=1, color=(255,0,0), thickness=2)
    attr_list = attr.numpy()
    attrs = attr_names[attr_list==1]
    label =''
# 属性标签
for att in attrs:
        label = label + att +'\n'
    ax.set_xlabel(label)
# 身份ID
    cele_id = identity.numpy()
    ax.set_title(f'ID: {cele_id}')
# 将OpenCV图片再次转成成Pillow图片格式
    img_pil =Image.fromarray( cv2.cvtColor(src=img_cv2,
                                            code=cv2.COLOR_BGR2RGB,))
    ax.imshow(img_pil)
plt.show()

# 显示标注框
fig = plt.figure(figsize=(14,7))
for idx inrange(n):
    img,(attr, identity, bbox, landmarks)= celeba[idx]
    ax = fig.add_subplot(1,4, idx +1)
    ax.set_xticks([])
    ax.set_yticks([])
    ax.set_frame_on(b=False)

# 图片的读取改为使用原图像
    file_path = os.path.join(celeba.root, celeba.base_folder,'img_celeba', celeba.filename[idx])

    img_cv2 = cv2.imread(file_path)
    bbox = bbox.numpy()
    cv2.rectangle(img=img_cv2, pt1=(bbox[0], bbox[1]),
                  pt2=(bbox[0]+ bbox[2], bbox[1]+ bbox[3]),
                  color=(255,0,0), thickness=2)
    cele_id = identity.numpy()
    ax.set_title(f'ID: {cele_id}')
    img_pil =Image.fromarray(cv2.cvtColor(src=img_cv2,
                                           code=cv2.COLOR_BGR2RGB,))
    ax.imshow(img_pil)
plt.show()

import os
import random
import numpy as np
import torch
from PIL importImage
from utils.tool import iou as IOU

current_path = os.path.dirname(os.path.abspath(__file__))
BASE_PATH = os.path.join(current_path,"datasets")
TARGET_PATH = os.path.join(BASE_PATH,"celeba")
IMG_PATH = os.path.join(BASE_PATH,"celeba/img_celeba")
DST_PATH = os.path.join(BASE_PATH,"train")
LABEL_PATH = os.path.join(TARGET_PATH,"list_bbox_celeba.txt")
LANMARKS_PATH = os.path.join(TARGET_PATH,"list_landmarks_celeba.txt")

# 测试样本个数限制,设置为 -1 表示全部
TEST_SAMPLE_LIMIT =100

# 为随机数种子做准备，使正样本，部分样本，负样本的比例为1：1：3
float_num =[0.1,0.1,0.3,0.5,0.95,0.95,0.99,0.99,0.99,0.99]

defcreate_directories(base_path, face_size):
    paths ={}
    base_path = os.path.join(base_path,f"{face_size}")
ifnot os.path.exists(base_path):
        os.makedirs(base_path)

    paths['positive']= os.path.join(base_path,"positive")
    paths['negative']= os.path.join(base_path,"negative")
    paths['part']= os.path.join(base_path,"part")

for path in paths.values():
ifnot os.path.exists(path):
            os.makedirs(path)

return paths, base_path

defopen_label_files(base_path):
    files ={}
    files['positive']=open(os.path.join(base_path,"positive.txt"),"w")
    files['negative']=open(os.path.join(base_path,"negative.txt"),"w")
    files['part']=open(os.path.join(base_path,"part.txt"),"w")
return files

defparse_annotation_line(line):
    strs = line.strip().split()
return strs

defadjust_bbox(x1, y1, w, h):
# 标注不标准，给框适当的偏移量
    x1 =int(x1 + w *0.12)
    y1 =int(y1 + h *0.1)

    x2 =int(x1 + w *0.9)
    y2 =int(y1 + h *0.85)

    w =int(x2 - x1)
    h =int(y2 - y1)

return x1, y1, x2, y2, w, h

defgenerate_crop_boxes(cx, cy, max_side, img_w, img_h):
"""
    根据给定的人脸中心点坐标和尺寸,生成5个候选的裁剪框。
    
    参数:
    cx (float): 人脸中心点的 x 坐标
    cy (float): 人脸中心点的 y 坐标
    max_side (int): 人脸框的最大边长
    img_w (int): 图像宽度
    img_h (int): 图像高度
    
    返回:
    crop_boxes (list): 一个包含5个裁剪框坐标的列表,每个裁剪框的格式为 [x1, y1, x2, y2]
    """

    crop_boxes =[]
for _ inrange(5):
# 随机偏移中心点坐标以及边长
        seed = float_num[np.random.randint(0,len(float_num))]

# 最大边长随机偏移
        _max_side = max_side + np.random.randint(int(-max_side * seed),int(max_side * seed))

# 中心点x坐标随机偏移
        _cx = cx + np.random.randint(int(-cx * seed),int(cx * seed))

# 中心点y坐标随机偏移
        _cy = cy + np.random.randint(int(-cy * seed),int(cy * seed))

# 得到偏移后的坐标值（方框）
        _x1 = _cx - _max_side /2
        _y1 = _cy - _max_side /2
        _x2 = _x1 + _max_side
        _y2 = _y1 + _max_side

# 偏移过大，偏出图像了，此时，不能用，应该再次尝试偏移
if _x1 <0or _y1 <0or _x2 > img_w or _y2 > img_h:
continue

# 添加裁剪框坐标到列表中
        crop_boxes.append(np.array([_x1, _y1, _x2, _y2]))

return crop_boxes

defprocess_crop_box(img, face_size, max_side, crop_box, boxes, landmarks):
"""
    处理单个裁剪框,生成正负样本。
    
    参数:
    img (Image): 原始图像
    crop_box (list): 裁剪框坐标 [x1, y1, x2, y2]
    boxes (list): 人脸框坐标列表
    face_size (int): 生成的人脸图像尺寸
    
    返回:
    sample (dict): 样本信息 {'image': image, 'label': label, 'bbox_offsets': offsets, 'landmark_offsets': landmark_offsets}
    """
    x1, y1, x2, y2 = boxes[0][:4]
    _x1, _y1, _x2, _y2 = crop_box[:4]
    px1, py1, px2, py2, px3, py3, px4, py4, px5, py5 = landmarks
    _max_side = max_side


    offset_x1 =(x1 - _x1)/ _max_side
    offset_y1 =(y1 - _y1)/ _max_side
    offset_x2 =(x2 - _x2)/ _max_side
    offset_y2 =(y2 - _y2)/ _max_side

    offset_px1 =(px1 - _x1)/ _max_side
    offset_py1 =(py1 - _y1)/ _max_side
    offset_px2 =(px2 - _x1)/ _max_side
    offset_py2 =(py2 - _y1)/ _max_side
    offset_px3 =(px3 - _x1)/ _max_side
    offset_py3 =(py3 - _y1)/ _max_side
    offset_px4 =(px4 - _x1)/ _max_side
    offset_py4 =(py4 - _y1)/ _max_side
    offset_px5 =(px5 - _x1)/ _max_side
    offset_py5 =(py5 - _y1)/ _max_side



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

    iou = IOU(torch.tensor([x1, y1, x2, y2]), torch.tensor([crop_box[:4]]))

if iou >0.7:# 正样本
        label =1
elif0.4< iou <0.6:# 部分样本
        label =2
elif iou <0.2:# 负样本
        label =0
else:
returnNone# 不符合任何条件的样本不处理

return{
'image': face_resize,
'label': label,
'bbox_offsets':(offset_x1, offset_y1, offset_x2, offset_y2),
'landmark_offsets':(offset_px1, offset_py1, offset_px2, offset_py2, offset_px3, offset_py3, offset_px4, offset_py4, offset_px5, offset_py5)
}
defprocess_annotation(face_size, anno_line, landmarks):
"""
    处理单行注释信息,生成正负样本。
    
    参数:
    anno_line (str): 一行注释信息,格式为 "image_filename x1 y1 w h"
    face_size (int): 生成的人脸图像尺寸
    landmarks (str): 关键点标注字符串
    
    返回:
    samples (list): 生成的样本列表
    """
# 5个关键点
    _landmarks = landmarks.split()

# 使用列表解析和解包一次性获取所有关键点的坐标
    landmarks =[float(x)for x in _landmarks[1:11]]

# 解析注释行,获取图像文件名和人脸位置信息
    strs = parse_annotation_line(anno_line)
    image_filename = strs[0].strip()
    x1, y1, w, h =map(int, strs[1:])

# 标签矫正
    x1, y1, x2, y2, w, h = adjust_bbox(x1, y1, w, h)
    boxes =[[x1, y1, x2, y2]]

# 计算人脸中心点坐标
    cx = w /2+ x1
    cy = h /2+ y1

# 最大边长
    max_side =max(w, h)

# 打开图像文件
    image_filepath = os.path.join(IMG_PATH, image_filename)
withImage.open(image_filepath)as img:
# 解析出宽度和高度
        img_w, img_h = img.size
# 生成候选的裁剪框
        samples =[]
for crop_box in generate_crop_boxes(cx, cy, max_side, img_w, img_h):
# 处理每个候选裁剪框,生成正负样本
            sample = process_crop_box(img, face_size, max_side, crop_box, boxes, landmarks )
if sample:
                samples.append(sample)

return samples

defsave_samples(samples, files, base_path, counters):
"""
    保存正负样本到文件中。
    
    参数:
    samples (list): 样本列表, 每个元素为一个字典, 包含 'image', 'label', 'bbox_offsets', 'landmark_offsets'
    files (dict): 包含正负样本输出文件的字典
    base_path (str): 输出文件的基础路径
    counters (dict): 样本计数器字典
    """
for sample in samples:
        image = sample['image']
        label = sample['label']
        bbox_offsets = sample['bbox_offsets']
        landmark_offsets = sample['landmark_offsets']

if label ==1:
            category ='positive'
            counters['positive']+=1
elif label ==2:
            category ='part'
            counters['part']+=1
else:
            category ='negative'
            counters['negative']+=1

        filename =f"{category}/{counters[category]}.jpg"
        image.save(os.path.join(base_path, filename))

try:
            bbox_str =' '.join(map(str, bbox_offsets))
            landmark_str =' '.join(map(str, landmark_offsets))
            files[category].write(f"{filename} {label} {bbox_str} {landmark_str}\n")
exceptIOErroras e:
print(f"Error writing to file: {e}")

defgenerate_samples(face_size, max_samples=-1):
"""
    生成指定大小的人脸样本,并保存到文件中。
    
    参数:
    face_size (int): 生成的人脸图像尺寸
    max_samples (int): 最大生成样本数量,设置为 -1 表示不限制
    """
ifnot os.path.exists(DST_PATH):
        os.makedirs(DST_PATH)

    paths, base_path = create_directories(DST_PATH, face_size)
# 新建标注文件
    files = open_label_files(base_path)

# 样本计数
    counters ={'positive':0,'negative':0,'part':0}

# 读取标注信息
withopen(LANMARKS_PATH)as f:
        landmarks_list = f.readlines()
withopen(LABEL_PATH)as f:
        anno_list = f.readlines()

for i,(anno_line, landmarks)inenumerate(zip(anno_list, landmarks_list)):
print(f"positive:{counters['positive']}, \
                negative:{counters['negative']}, \
                part:{counters['part']}")

# 跳过前两行
if i <2:
continue

# 如果处理了指定数量的样本,则退出循环
if max_samples >0and i > max_samples:
break

# 处理单行标注信息,生成正负样本
        samples = process_annotation(
            face_size, anno_line, landmarks
)

# 保存正负样本到文件
        save_samples(
            samples,
            files, base_path, counters
)

for file in files.values():
        file.close()

defmain():
# 生成12×12的样本
    generate_samples(12,1000)

    generate_samples(24,1000)

    generate_samples(48,1000)

if __name__ =="__main__":
    main()

import torch
import os
from torch.utils.data importDataLoader
from train.FaceDatasetimportFaceDataset
import matplotlib.pyplot as plt


classTrainer:
def__init__(self, net, param_path, data_path):
# 检测是否有GPU
        self.device ='cuda:0'if torch.cuda.is_available()else"cpu"
# 把模型搬到device
        self.net = net.to(self.device)

        self.param_path = param_path

# 打包数据
        self.datasets =FaceDataset(data_path)

# 定义损失函数：类别判断（分类任务）
        self.cls_loss_func = torch.nn.BCELoss()

# 定义损失函数：框的偏置回归
        self.offset_loss_func = torch.nn.MSELoss()

# 定义损失函数：关键点的偏置回归
        self.point_loss_func = torch.nn.MSELoss()

# 定义优化器
        self.optimizer = torch.optim.Adam(params=self.net.parameters(), lr=1e-3)

defcompute_loss(self, out_cls, out_offset, out_point, cls, offset, point, landmark):
# 选取置信度为0，1的正负样本求置信度损失
        cls_mask = torch.lt(cls,2)
        cls_loss = self.cls_loss_func(torch.masked_select(out_cls, cls_mask),
                                      torch.masked_select(cls, cls_mask))

# 选取正样本和部分样本求偏移率的损失
        offset_mask = torch.gt(cls,0)
        offset_loss = self.offset_loss_func(torch.masked_select(out_offset, offset_mask),
                                            torch.masked_select(offset, offset_mask))

if landmark:
            point_loss = self.point_loss_func(torch.masked_select(out_point, offset_mask),
                                              torch.masked_select(point, offset_mask))
return cls_loss, offset_loss, point_loss
else:
return cls_loss, offset_loss,None

deftrain(self, epochs, landmark=False):
"""
            - 断点续传 --> 短点续训
            - transfer learning 迁移学习
            - pretrained model 预训练

        :param epochs: 训练的轮数
        :param landmark: 是否为landmark任务
        :return:
        """

# 加载上次训练的参数
if os.path.exists(self.param_path):
            self.net.load_state_dict(torch.load(self.param_path))
print("加载参数文件,继续训练 ...")
else:
print("没有参数文件,全新训练 ...")

# 封装数据加载器
        dataloader =DataLoader(self.datasets, batch_size=32, shuffle=True)

# 定义列表存储损失值
        cls_losses =[]
        offset_losses =[]
        point_losses =[]
        total_losses =[]

for epoch inrange(epochs):
# 训练一轮
for i,(img_data, _cls, _offset, _point)inenumerate(dataloader):
# 数据搬家 [32, 3, 12, 12]
                img_data = img_data.to(self.device)
                _cls = _cls.to(self.device)
                _offset = _offset.to(self.device)
                _point = _point.to(self.device)

if landmark:
# O-Net输出三个
                    out_cls, out_offset, out_point = self.net(img_data)
                    out_point = out_point.view(-1,10)
else:
# O-Net输出两个
                    out_cls, out_offset = self.net(img_data)
                    out_point =None

# [B, C, H, W] 转换为 [B, C]
                out_cls = out_cls.view(-1,1)
                out_offset = out_offset.view(-1,4)

if landmark:
                    out_point = out_point.view(-1,10)

# 计算损失
                cls_loss, offset_loss, point_loss = self.compute_loss(out_cls, out_offset, out_point,
                                                                    _cls, _offset, _point, landmark)

if landmark:
                    loss = cls_loss + offset_loss + point_loss
else:
                    loss = cls_loss + offset_loss

# 打印损失
if landmark:
print(f"Epoch [{epoch+1}/{epochs}], loss:{loss.item():.4f}, cls_loss:{cls_loss.item():.4f}, "
f"offset_loss:{offset_loss.item():.4f}, point_loss:{point_loss.item():.4f}")
else:
print(f"Epoch [{epoch+1}/{epochs}], loss:{loss.item():.4f}, cls_loss:{cls_loss.item():.4f}, "
f"offset_loss:{offset_loss.item():.4f}")

# 存储损失值
                cls_losses.append(cls_loss.item())
                offset_losses.append(offset_loss.item())
if landmark:
                    point_losses.append(point_loss.item())
                total_losses.append(loss.item())

# 清空梯度
                self.optimizer.zero_grad()

# 梯度回传
                loss.backward()

# 优化
                self.optimizer.step()

# 保存模型（参数）
            torch.save(self.net.state_dict(), self.param_path)

# 绘制损失曲线
        self.plot_losses(cls_losses, offset_losses, point_losses, total_losses, landmark)

print("训练完成!")

defplot_losses(self, cls_losses, offset_losses, point_losses, total_losses, landmark):
"""
        绘制训练过程中的损失曲线
        :param cls_losses: 分类损失列表
        :param offset_losses: 边界框偏移损失列表
        :param point_losses: 关键点偏移损失列表
        :param total_losses: 总损失列表
        :param landmark: 是否为landmark任务
        """
        plt.figure(figsize=(12,6))
        plt.subplot(1,2,1)
        plt.plot(cls_losses, label='Classification Loss')
        plt.plot(offset_losses, label='Offset Loss')
if landmark:
            plt.plot(point_losses, label='Point Loss')
        plt.legend()
        plt.xlabel('Iteration')
        plt.ylabel('Loss')
        plt.title('Training Losses')

        plt.subplot(1,2,2)
        plt.plot(total_losses)
        plt.xlabel('Iteration')
        plt.ylabel('Total Loss')
        plt.title('Total Training Loss')

        plt.savefig('training_losses.png')
        plt.close()

from train import model_mtcnn as nets
import os
import train.train as train


if __name__ =='__main__':

    current_path = os.path.dirname(os.path.abspath(__file__))
# 权重存放地址
    base_path = os.path.join(current_path,"model")
    model_path = os.path.join(base_path,"p_net.pt")

# 数据存放地址
    data_path = os.path.join(current_path,"datasets/train/12")

# 如果没有这个参数存放目录，则创建一个目录
ifnot os.path.exists(base_path):
        os.makedirs(base_path)

# 构建模型
    pnet = nets.PNet()

# 开始训练
    t = train.Trainer(pnet, model_path, data_path)

# t.train2(0.01)
    t.train(100)