使用opencv库实现口罩识别、人脸检测数据库
之前做电子设计赛省赛,需要实现口罩识别,人脸检测,我们使用opencv自带库实现。最近无聊,重新在linux系统上完善一波,记录一下代码(待完善)。
文件结构如下:
- database文件夹:
用于存放图像数据,文件夹命名为id name,图像命名为数字.jpg
- xml文件夹
cv模型,github查查应该能找到
- data.py
主代码data.py如下:
import cv2
import os
import numpy as np
from time import sleep
import time
import shutil
videoIn = cv2.VideoCapture("/dev/video0")
ret, np_frame = videoIn.read()
#利用cv2内置的函数,初步识别人脸
face_cascade = cv2.CascadeClassifier('xml/haarcascade_frontalface_default.xml')
nose_cascade = cv2.CascadeClassifier('xml/haarcascade_nose.xml')
mouth_cascade = cv2.CascadeClassifier('xml/haarcascade_mouth.xml')
eye_cascade = cv2.CascadeClassifier('xml/haarcascade_eye.xml')
#利用cv2内置的人脸识别模块,可进行预训练和预测
face_recognizer = cv2.face.LBPHFaceRecognizer_create()
#全局使用
datapath = 'database'
database = {'name':[],'id':[]}
#读取数据库
def Read_Database():
dirs = os.listdir(datapath)
for dir in dirs:
id,name = dir.split()#文件夹命名格式为id name
database['name'].append(name)
database['id'].append(id)
#用cv2内置的人脸检测模块,判断图片中是否有人。如果有,则将人脸部分和人脸坐标返回
def detect_face(img):
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
faces = face_cascade.detectMultiScale(gray, scaleFactor=1.2, minNeighbors=5)
if (len(faces) == 0):
return None, None
(x, y, w, h) = faces[0]
return gray[y:y + w, x:x + h], faces[0]
#进度条打印
def progress(per,width=50):
text = ('\r[%%-%ds]'%width)%('#'*int(per*width))
text += '%3s%%'
text = text%(round(per*100))
print('\t'+text,end='')
#准备训练数据集:faces放置人脸图像,labels放置下表标签0、1、2、3...
def prepare_training_data(data_folder_path):
dirs = os.listdir(data_folder_path)
faces = []
labels = []
print('<------------------- Train Mode ------------------->')
for dir_name in dirs:
if not dir_name == '.ipynb_checkpoints':
id,name = dir_name.split()
for i in range(len(database['id'])):
if database['id'][i] == id:
label = int(i)#训练的标签是id name在database中的下标
break
subject_dir_path = data_folder_path + "/" + dir_name
subject_images_names = os.listdir(subject_dir_path)
print('\n[Name] {0} [Id] {1}'.format(name.ljust(5,' '),id.ljust(3,' ')))
now_num = 0#本数据已经训练好的图片数量
for image_name in subject_images_names:
if not image_name == '.ipynb_checkpoints':
image_path = subject_dir_path + "/" + image_name
image = cv2.imread(image_path)
face, rect = detect_face(image)
if face is not None:
faces.append(face)
labels.append(label)
now_num += 1
#打印训练进度:
per = now_num / len(subject_images_names)
progress(per)
progress(1)
print('\n\n<------------------- Complete ------------------->\n')
return faces, labels
#画方框
def draw_rectangle(img, rect):
(x, y, w, h) = rect
cv2.rectangle(img, (x, y), (x + w, y + h), (128, 128, 0), 2)
#写文字
def draw_text(img, text, x, y):
cv2.putText(img, text, (x, y), cv2.FONT_HERSHEY_COMPLEX, 1, (128, 128, 0), 2)
#图片预测
def predict(test_img):
img = test_img.copy()
face, rect = detect_face(img)#得到人脸部分
#如果图像中检测到人脸
if not face is None:
label = face_recognizer.predict(face)
if label[1]<100:#loss越小,可信度越高
label_text = database['name'][label[0]]
else:
label_text = 'Unknown'
draw_rectangle(img, rect)
draw_text(img, label_text, rect[0], rect[1] - 5)
# print("found")
else:
draw_text(img, 'nofound',0,20)
# print("nofound")
return img
#新图像保存
def imgSave(img,num,id,name):
newdir = datapath + '/' + id + ' ' +name#文件夹格式id name
if not os.path.exists(newdir):
os.mkdir(newdir)
face, rect = detect_face(img)
if not face is None:#识别到脸部才进行录入
cv2.imwrite(newdir+"/"+str(num)+".jpg", img[:,:,:])
num=num+1
return img, num
#显示当前数据库
def Dis_Database():
print('序号'.ljust(4,' ')+'id'.ljust(6,' ')+'name'.ljust(6,' '))
for i in range(len(database['name'])):
id = database['id'][i]
name = database['name'][i]
print(str(i).ljust(6,' ')+id.ljust(6,' ')+name.ljust(6,' '))
####################################################################################################
stage = 0
Read_Database()
faces, labels = prepare_training_data(datapath)#初始数据集
if (len(database['name']) != 0):
face_recognizer.train(faces, np.array(labels))#训练
while True:
# 检测按键,停止训练,保存当前模型和最佳模型
#0正常显示;1口罩识别;2人脸检测;3数据录入
k = cv2.waitKey(1) & 0xFF
if k == ord('0'):
stage = 0;
elif k == ord('1'):
stage = 1;
elif k == ord('2'):
stage = 2;
elif k == ord('3'):
stage = 3;
if stage == 0:
ret, np_frame = videoIn.read()
cv2.imshow("capture", np_frame)
elif stage == 1:
ret, np_frame = videoIn.read()
gray = cv2.cvtColor(np_frame, cv2.COLOR_BGR2GRAY)
noses = nose_cascade.detectMultiScale(gray)
if(len(noses)>=1):
draw_text(np_frame, 'Nomask',0,20)
else:
eyes = eye_cascade.detectMultiScale(gray)
if(len(eyes)>=1):
draw_text(np_frame, 'Mask',0,20)
cv2.imshow("capture", np_frame)
elif stage == 2:
ret, np_frame = videoIn.read()
if (len(database['name'])!=0):
np_frame = predict(np_frame)
else:
draw_text(np_frame, 'Database empty',0,20)
cv2.imshow("capture", np_frame)
elif stage == 3:
choice = 0
print('<------------------------------- 数据库修改 ------------------------------->')
choice = int(input('功能选项(返回--0 录入数据--1 删除数据--2 添加数据--3):'))
#功能实现
#直接返回
if choice == 0:
pass
#选择录入数据
elif choice == 1:
name = input('请输入姓名:')
id = input('请输入id号:')
number = int(input('请输入录入图片数量:'))
num = 1
database['name'].append(name)
database['id'].append(id)
while True:
ret, np_frame = videoIn.read()
cv2.imshow("capture", np_frame)
np_frame, num = imgSave(np_frame, num , id ,name)
draw_text(np_frame, 'num:'+str(num),0,20)
if num > number:
faces, labels = prepare_training_data(datapath)
face_recognizer.train(faces, np.array(labels))
break
if cv2.waitKey(1) & 0xFF == ord('q'):#强制退出
break
print('数据录入成功!')
#选择删除数据
elif choice == 2:
Dis_Database()
index = int(input('请输入想删除的数据的序号:'))
while(index > len(database['id'])-1 or index < 0):
index = int(input('输入序号不存在,请重新输入:'))
del_path = datapath + '/' + database['id'][index] + ' ' + database['name'][index]
if not os.path.exists(del_path):
print('文件夹不存在')
else:
shutil.rmtree(del_path)
del database['name'][index]
del database['id'][index]
if len(database['name']) != 0:#数据库非空则训练
faces, labels = prepare_training_data(datapath)
face_recognizer.train(faces, np.array(labels))
print('数据删除成功!')
#添加数据
elif choice == 3:
Dis_Database()
index = int(input('请输入想添加的数据的序号:'))
number = int(input('请输入想添加的图片的数量:'))
while(index > len(database['id'])-1 or index < 0):
index = int(input('输入序号不存在,请重新输入:'))
add_path = datapath + '/' + database['id'][index] + ' ' + database['name'][index]
old_len = len(os.listdir(add_path))
num = old_len + 1
id = database['id'][index]
name = database['name'][index]
while True:
ret, np_frame = videoIn.read()
cv2.imshow("capture", np_frame)
np_frame, num = imgSave(np_frame, num , id ,name)
draw_text(np_frame, 'num:'+str(num),0,20)
if num > old_len + number:
faces, labels = prepare_training_data(datapath)
face_recognizer.train(faces, np.array(labels))
break
if cv2.waitKey(1) & 0xFF == ord('q'):#强制退出
break
stage = 0
print('<------------------------------- 返回界面 ------------------------------->')
print('\n')
运行后,在非终端处:
- 按1开始口罩识别
- 按2开始人脸检测
- 按3开始人脸检测数据库修改:修改时要在终端输入,可实现数据录入、删除、添加
- 按0回到正常显示界面。
人脸检测效果
pynq实现如下
import cv2
import os
import numpy as np
from time import sleep
import time
import shutil
from pynq.overlays.base import BaseOverlay
from pynq.lib.video import *
# 显示器
base = BaseOverlay("base.bit")
# monitor configuration: 640*480 @ 60Hz
Mode = VideoMode(640,480,24)
hdmi_out = base.video.hdmi_out
hdmi_out.configure(Mode,PIXEL_BGR)
hdmi_out.start()
screen_w=320
screen_h=240
frame_out_w = 320
frame_out_h = 240
# 摄像头
frame_in_w = 320
frame_in_h = 240
videoIn = cv2.VideoCapture(0)
videoIn.set(cv2.CAP_PROP_FRAME_WIDTH, frame_in_w);
videoIn.set(cv2.CAP_PROP_FRAME_HEIGHT, frame_in_h);
print("Capture device is open: " + str(videoIn.isOpened()))
#利用cv2内置的函数,初步识别人脸
face_cascade = cv2.CascadeClassifier('xml/haarcascade_frontalface_default.xml')
nose_cascade = cv2.CascadeClassifier('xml/haarcascade_nose.xml')
mouth_cascade = cv2.CascadeClassifier('xml/haarcascade_mouth.xml')
eye_cascade = cv2.CascadeClassifier('xml/haarcascade_eye.xml')
#利用cv2内置的人脸识别模块,可进行预训练和预测
face_recognizer = cv2.face.createLBPHFaceRecognizer()
#全局使用
datapath = 'database'
database = {'name':[],'id':[]}
#读取数据库
def Read_Database():
dirs = os.listdir(datapath)
for dir in dirs:
if not dir == '.ipynb_checkpoints':
id,name = dir.split()#文件夹命名格式为id name
database['name'].append(name)
database['id'].append(id)
#用cv2内置的人脸检测模块,判断图片中是否有人。如果有,则将人脸部分和人脸坐标返回
def detect_face(img):
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
faces = face_cascade.detectMultiScale(gray, scaleFactor=1.2, minNeighbors=5)
if (len(faces) == 0):
return None, None
(x, y, w, h) = faces[0]
return gray[y:y + w, x:x + h], faces[0]
#进度条打印
def progress(per,width=50):
text = ('\r[%%-%ds]'%width)%('#'*int(per*width))
text += '%3s%%'
text = text%(round(per*100))
print('\t'+text,end='')
#准备训练数据集:faces放置人脸图像,labels放置下表标签0、1、2、3...
def prepare_training_data(data_folder_path):
dirs = os.listdir(data_folder_path)
faces = []
labels = []
print('<------------------- Train Mode ------------------->')
for dir_name in dirs:
if dir_name == '.ipynb_checkpoints':
continue
id,name = dir_name.split()
for i in range(len(database['id'])):
if database['id'][i] == id:
label = int(i)#训练的标签是id name在database中的下标
break
subject_dir_path = data_folder_path + "/" + dir_name
subject_images_names = os.listdir(subject_dir_path)
print('\n[Name] {0} [Id] {1}'.format(name.ljust(5,' '),id.ljust(3,' ')))
now_num = 0#本数据已经训练好的图片数量
for image_name in subject_images_names:
if not image_name == '.ipynb_checkpoints':
image_path = subject_dir_path + "/" + image_name
image = cv2.imread(image_path)
face, rect = detect_face(image)
if face is not None:
faces.append(face)
labels.append(label)
now_num += 1
#打印训练进度:
per = now_num / len(subject_images_names)
progress(per)
progress(1)
print('\n\n<------------------- Complete ------------------->\n')
return faces, labels
#画方框
def draw_rectangle(img, rect):
(x, y, w, h) = rect
cv2.rectangle(img, (x, y), (x + w, y + h), (128, 128, 0), 2)
#写文字
def draw_text(img, text, x, y):
cv2.putText(img, text, (x, y), cv2.FONT_HERSHEY_COMPLEX, 1, (128, 128, 0), 2)
#图片预测
def predict(test_img):
img = test_img.copy()
face, rect = detect_face(img)#得到人脸部分
#如果图像中检测到人脸
if not face is None:
label = face_recognizer.predict(face)
if label[1]<100:#loss越小,可信度越高
label_text = database['name'][label[0]]
base.rgbleds[5].write(2)
else:
label_text = 'Unknown'
base.rgbleds[5].write(4)
draw_rectangle(img, rect)
draw_text(img, label_text, rect[0], rect[1] - 5)
else:
draw_text(img, 'nofound',0,20)
base.rgbleds[5].write(0)
return img
#新图像保存
def imgSave(img,num,id,name):
newdir = datapath + '/' + id + ' ' +name#文件夹格式id name
if not os.path.exists(newdir):
os.mkdir(newdir)
face, rect = detect_face(img)
if not face is None:#识别到脸部才进行录入
cv2.imwrite(newdir+"/"+str(num)+".jpg", img[:,:,:])
num=num+1
return img, num
#显示当前数据库
def Dis_Database():
print('序号'.ljust(4,' ')+'id'.ljust(6,' ')+'name'.ljust(6,' '))
for i in range(len(database['name'])):
id = database['id'][i]
name = database['name'][i]
print(str(i).ljust(6,' ')+id.ljust(6,' ')+name.ljust(6,' '))
####################################################################################################
stage = 0
Read_Database()
faces, labels = prepare_training_data(datapath)#初始数据集
if (len(database['name']) != 0):
face_recognizer.train(faces, np.array(labels))#训练
while True:
# 检测按键,停止训练,保存当前模型和最佳模型
# 0正常显示;1口罩识别;2人脸检测;3数据录入
ret, np_frame = videoIn.read()
if base.buttons[0].read()==1:
stage=0
if base.buttons[1].read()==1:
stage=1
if base.buttons[2].read()==1:
stage=2
if base.buttons[3].read()==1:
stage=3
if stage == 0:
base.rgbleds[4].write(0)
base.rgbleds[5].write(0)
elif stage == 1:
gray = cv2.cvtColor(np_frame, cv2.COLOR_BGR2GRAY)
noses = nose_cascade.detectMultiScale(gray)
if(len(noses)>=1):
draw_text(np_frame, 'Nomask',0,20)
base.rgbleds[4].write(4)
else:
eyes = eye_cascade.detectMultiScale(gray)
if(len(eyes)>=1):
draw_text(np_frame, 'Mask',0,20)
base.rgbleds[4].write(2)
else:
base.rgbleds[4].write(0)
elif stage == 2:
if (len(database['name'])!=0):
np_frame = predict(np_frame)
else:
draw_text(np_frame, 'Database empty',0,20)
elif stage == 3:
choice = 0
print('<------------------------------- 数据库修改 ------------------------------->')
choice = int(input('功能选项(返回--0 录入数据--1 删除数据--2 添加数据--3):'))
#功能实现
#直接返回
if choice == 0:
pass
#选择录入数据
elif choice == 1:
Dis_Database()
name = input('请输入姓名:')
id = input('请输入id号:')
number = int(input('请输入录入图片数量:'))
num = 1
database['name'].append(name)
database['id'].append(id)
while True:
ret, np_frame = videoIn.read()
outframe = hdmi_out.newframe()
outframe[0:screen_h,0:screen_w,:] = np_frame[0:screen_h,0:screen_w,:]
hdmi_out.writeframe(outframe)
np_frame, num = imgSave(np_frame, num , id ,name)
draw_text(np_frame, 'num:'+str(num),0,20)
progress((num-1)/number)
if num > number:
progress(1)
print('\n')
faces, labels = prepare_training_data(datapath)
face_recognizer.train(faces, np.array(labels))
break
# if base.buttons[0].read()==0:#强制退出
# break
print('数据录入成功!')
#选择删除数据
elif choice == 2:
Dis_Database()
index = int(input('请输入想删除的数据的序号:'))
while(index > len(database['id'])-1 or index < 0):
index = int(input('输入序号不存在,请重新输入:'))
del_path = datapath + '/' + database['id'][index] + ' ' + database['name'][index]
if not os.path.exists(del_path):
print('文件夹不存在')
else:
shutil.rmtree(del_path)
del database['name'][index]
del database['id'][index]
if len(database['name']) != 0:#数据库非空则训练
faces, labels = prepare_training_data(datapath)
face_recognizer.train(faces, np.array(labels))
print('数据删除成功!')
#添加数据
elif choice == 3:
Dis_Database()
index = int(input('请输入想添加的数据的序号:'))
number = int(input('请输入想添加的图片的数量:'))
while(index > len(database['id'])-1 or index < 0):
index = int(input('输入序号不存在,请重新输入:'))
add_path = datapath + '/' + database['id'][index] + ' ' + database['name'][index]
old_len = len(os.listdir(add_path))
num = old_len + 1
id = database['id'][index]
name = database['name'][index]
print(add_path,old_len,id,name)
while True:
ret, np_frame = videoIn.read()
outframe = hdmi_out.newframe()
outframe[0:screen_h,0:screen_w,:] = np_frame[0:screen_h,0:screen_w,:]
hdmi_out.writeframe(outframe)
np_frame, num = imgSave(np_frame, num , id ,name)
draw_text(np_frame, 'num:'+str(num),0,20)
progress((num-old_len-1)/number)
if num > old_len + number:
progress(1)
print('\n')
faces, labels = prepare_training_data(datapath)
face_recognizer.train(faces, np.array(labels))
break
# if base.buttons[0].read()==0:#强制退出
# break
stage = 0
print('<------------------------------- 返回界面 ------------------------------->')
print('\n')
outframe = hdmi_out.newframe()
outframe[0:screen_h,0:screen_w,:] = np_frame[0:screen_h,0:screen_w,:]
hdmi_out.writeframe(outframe)