基于pytorch与opencv简单做个人脸识别_pytorch 人脸识别

环境准备

pytorch 2.0.0+cu117  （这个版本你们自己根据实际情况设置）

opencv（注意不同版本的使用规则不同，一定要看官方文档，嫌麻烦跟我版本一致）

opencv-contrib-python             4.7.0.72
opencv-python                     4.7.0.72

数据集

105_classes_pins_dataset

请在数据集中新建文件夹命名为，被训练的人名字与照片，图片样式(无遮挡)

请使用 jupyter notebook 

以上都准备好了，那就开始代码之旅叭~

代码

各类包导入

import torch, torchvision
import torchvision.transforms as transforms
import torchvision.datasets as datasets
import torchvision.models as models
import torch.nn as nn
import torch.optim as optim
import torch.nn.init as init
import torch.nn.functional as F
import numpy as np
from PIL import Image
import matplotlib.pyplot as plt
import torch.utils.data as D

数据集的读取与处理

# 数据扩增方法（旋转、平移）一般都不会改变标签
transformations = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.RandomHorizontalFlip(p=0.5),
    transforms.RandomRotation(5),
    transforms.ColorJitter(),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], 
                         std=[0.229, 0.224, 0.225])
])
 
# 数据集获取并数据扩增
trainfolder = datasets.ImageFolder('data/105_classes_pins_dataset/', transform=transformations)
 
# 对数据集进行划分，每7个添加一次验证集
train_idx = []
val_idx = []
for idx in range(len(trainfolder)):
    if idx % 7 == 0:
        val_idx.append(idx)
    else:
        train_idx.append(idx)
 
# 图片的读取，（图片、类别）
train_dataset = D.Subset(trainfolder, train_idx)
val_dataset = D.Subset(trainfolder, val_idx)
 
# 对应参数意思 一次 16 张图片，每个epoch重新打乱数据，多进程5个运行
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=16, 
                                           shuffle=True, num_workers=5)
 
test_loader = torch.utils.data.DataLoader(val_dataset, batch_size=16,
                                         num_workers=5)
# 查看分了多少类
trainfolder.classes

模型构建

# res-net18做预训练模型
model = models.resnet18(pretrained=True)
# 将全连接层改为(512, cls_num) 根据预测类别个数实际填充,我的数据集裁剪过
model.fc = nn.Linear(512, 6)
 
# 启用gpu
device = torch.device("cuda")
# 将模型加载进gpu
model = model.to(device)
# res-net18模型采用交叉熵损失
criterion = nn.CrossEntropyLoss()
# 使用Adam优化器，lr学习率
optimizer = optim.Adam(model.parameters(), lr=0.001)

模型训练与验证

# 训练
def train(model, device, train_loader, optimizer, epoch):
    model.train() # 转为训练，dropout起作用
    for batch_idx, (data, target) in enumerate(train_loader):
        # 如果在gpu运行，则需要将数据传输到GPU
        data, target = data.to(device), target.to(device)
        
        optimizer.zero_grad() # 清空梯度
        
        # 所有计算，都需要在相同的设备进行计算
        # 如果model 在gpu，则data也要在gpu
        output = model(data) # 正向传播
        loss = criterion(output, target) # 计算损失
        loss.backward() # 梯度计算
        optimizer.step() # 参数更新
        
        # 训练100个打印一次
        if batch_idx % 100 == 0:
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch, 
                batch_idx * len(data), # 已经训练的样本个数
                len(train_loader.dataset), # 所有训练集样本个数
                100. * batch_idx / len(train_loader), # 训练batch / 所有batch计数
                loss.item() # 当前batch的损失
            ))
 
# 测试
def test(model, device, test_loader):
    model.eval() # 转为预测，drop不起作用
    test_loss = 0
    correct = 0
    
    # 下面不需要计算梯度，只需要正向传播
    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = model(data) # 原始输出的数值
            test_loss += criterion(output, target).sum(0).item()  
            output = F.softmax(output) # 转为概率
            pred = output.argmax(dim=1, keepdim=True)  # 获得最大值
            correct += pred.eq(target.view_as(pred)).sum().item() # 准确率
    
    # 每个样本的分类损失
    test_loss /= len(test_loader.dataset)
 
    print('Test set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
        test_loss, # 样本分类损失
        correct, # 分类正确的样本个数
        len(test_loader.dataset), # 测试集总共的样本的个数
        100. * correct / len(test_loader.dataset) # 测试集样本的分类准确率
    ))
 
# 学习率调整策略
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.3)
 
# 开始训练模型
for epoch in range(1, 10):
    train(model, device, train_loader, optimizer, epoch)
    test(model, device, test_loader)
    scheduler.step()

最后一步预测！（这里放个单张图片预测）

# 读入图像(自己改)
img = Image.open('data/test.jpg')
# 读取类别
trainfolder = datasets.ImageFolder('data/105_classes_pins_dataset/',transform=transformations)
cls_img = trainfolder.classes
print(cls_img)
 
 
# 转换图像
transformations = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], 
                         std=[0.229, 0.224, 0.225])
])
img = transformations(img)
img = img.unsqueeze(0) # 将数据增加一维，成为 1 x C x H x W 的格式
img = img.to(device)
 
# 预测结果输出
output = model(img)
_, predicted = torch.max(output.data, 1)
# 类别打印
print(cls_img[predicted.item()])

视频流识别

 
# 读取类别
trainfolder = datasets.ImageFolder('data/105_classes_pins_dataset/', transform=transformations)
cls_img = trainfolder.classes
print(cls_img)
 
# 转换图像
transformations = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], 
                         std=[0.229, 0.224, 0.225])
])
 
# cv4.7 废弃：cv2.CascadeClassifier('./haarcascade_frontalface_alt2.xml')
# 加载人脸识别模型
face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + "haarcascade_frontalface_default.xml")
 
# 打开摄像头
cap = cv2.VideoCapture(r"data/face_test.mp4")
 
while True:
    # 读取摄像头中的图像
    ret, frame = cap.read()
    if ret:
        # 对图像进行人脸检测
        height, width = frame.shape[:2]
        # 预测
        img = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
 
        img = transformations(img)
        img = img.unsqueeze(0) # 将数据增加一维，成为 1 x C x H x W 的格式
        img = img.to(device)
        output = model(img)
        _, predicted = torch.max(output.data, 1)
        # 将显示图像大小缩小一半
        new_height, new_width = int(height / 2), int(width / 2)
        frame = cv2.resize(frame, (new_width, new_height), interpolation=cv2.INTER_LINEAR)
        
        # 识别
        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        faces = face_cascade.detectMultiScale(gray, 1.3, 5)
        
        # 类别文本
        text = ""+str(cls_img[predicted])
        # 在图像中标记人脸区域
        for (x, y, w, h) in faces:
            cv2.rectangle(frame, (x, y), (x + w, y + h), (0, 0, 255), 2)
            cv2.putText(frame, text, (x, y), cv2.FONT_HERSHEY_SIMPLEX, 1.5, (0, 0, 255), 2)
 
        # 显示处理后的图像
        cv2.imshow('frame', frame)
 
        # 如果按下 q 键，则退出程序
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break
    else:
        break
# 释放摄像头资源
cap.release()
 
# 关闭所有窗口
cv2.destroyAllWindows()