基于MTCNN+arcface的人脸检测识别(pytorch)

import pickle import shutil import numpy as np import random import os import cv2 from tqdm import tqdm class BBox: # 人脸的box def __init__(self, box): self.left = box[0] self.top = box[1] self.right = box[2] self.bottom = box[3] self.x = box[0] self.y = box[1] self.w = box[2] - box[0] self.h = box[3] - box[1] def project(self, point): """将关键点的绝对值转换为相对于左上角坐标偏移并归一化 参数： point：某一关键点坐标(x,y) 返回值： 处理后偏移 """ x = (point[0] - self.x) / self.w y = (point[1] - self.y) / self.h return np.asarray([x, y]) def reproject(self, point): """将关键点的相对值转换为绝对值，与project相反 参数： point:某一关键点的相对归一化坐标 返回值： 处理后的绝对坐标 """ x = self.x + self.w * point[0] y = self.y + self.h * point[1] return np.asarray([x, y]) def reprojectLandmark(self, landmark): """对所有关键点进行reproject操作""" p = np.zeros((len(landmark), 2)) for i in range(len(landmark)): p[i] = self.reproject(landmark[i]) return p def projectLandmark(self, landmark): """对所有关键点进行project操作""" p = np.zeros((len(landmark), 2)) for i in range(len(landmark)): p[i] = self.project(landmark[i]) return p # 预处理数据，转化图像尺度并对像素归一 def processed_image(img, scale): height, width, channels = img.shape new_height = int(height * scale) new_width = int(width * scale) new_dim = (new_width, new_height) img_resized = cv2.resize(img, new_dim, interpolation=cv2.INTER_LINEAR) # 把图片转换成numpy值 image = np.array(img_resized).astype(np.float32) # 转换成CHW image = image.transpose((2, 0, 1)) # 归一化 image = (image - 127.5) / 128 return image 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 get_landmark_from_lfw_neg(txt, data_path, with_landmark=True): """获取txt中的图像路径，人脸box，人脸关键点 参数： txt：数据txt文件 data_path:数据存储目录 with_landmark:是否留有关键点 返回值： result包含(图像路径，人脸box，关键点) """ with open(txt, 'r') as f: lines = f.readlines() result = [] for line in lines: line = line.strip() components = line.split(' ') # 获取图像路径 img_path = os.path.join(data_path, components[0]).replace('\\', '/') # 人脸box box = (components[1], components[3], components[2], components[4]) box = [float(_) for _ in box] box = list(map(int, box)) if not with_landmark: result.append((img_path, BBox(box))) continue # 五个关键点(x,y) landmark = np.zeros((5, 2)) for index in range(5): rv = (float(components[5 + 2 * index]), float(components[5 + 2 * index + 1])) landmark[index] = rv result.append((img_path, BBox(box), landmark)) return result def get_landmark_from_celeba(data_path, with_landmark=True): """获取celeba的脸box，人脸关键点 参数： bbox_txt：数据bbox文件 landmarks_txt：数据landmarks文件 data_path:数据存储目录 with_landmark:是否留有关键点 返回值： result包含(图像路径，人脸box，关键点) """ bbox_txt = os.path.join(data_path, 'list_bbox_celeba.txt') landmarks_txt = os.path.join(data_path, 'list_landmarks_celeba.txt') # 获取图像路径，box，关键点 if not os.path.exists(bbox_txt): return [] with open(bbox_txt, 'r') as f: bbox_lines = f.readlines() with open(landmarks_txt, 'r') as f: landmarks_lines = f.readlines() result = [] for i in range(2, len(bbox_lines)): bbox_line = bbox_lines[i] landmarks_line = landmarks_lines[i] bbox_components = bbox_line.strip().split() landmarks_components = landmarks_line.strip().split() # 获取图像路径 img_path = os.path.join(data_path, 'img_celeba', bbox_components[0]).replace('\\', '/') # 人脸box box = (bbox_components[1], bbox_components[2], bbox_components[3], bbox_components[4]) box = [float(_) for _ in box] box = list(map(int, box)) box = [box[0], box[1], box[2] + box[0], box[3] + box[1]] if not with_landmark: result.append((img_path, BBox(box))) continue # 五个关键点(x,y) landmark = np.zeros((5, 2)) for index in range(5): rv = (float(landmarks_components[1 + 2 * index]), float(landmarks_components[1 + 2 * index + 1])) landmark[index] = rv result.append((img_path, BBox(box), landmark)) return result def combine_data_list(data_dir): """把每个数据列表放在同一个文件上 参数： data_dir：已经裁剪后的文件夹 """ npr = np.random with open(os.path.join(data_dir, 'positive.txt'), 'r') as f: pos = f.readlines() with open(os.path.join(data_dir, 'negative.txt'), 'r') as f: neg = f.readlines() with open(os.path.join(data_dir, 'part.txt'), 'r') as f: part = f.readlines() with open(os.path.join(data_dir, 'landmark.txt'), 'r') as f: landmark = f.readlines() with open(os.path.join(data_dir, 'all_data_list.txt'), 'w') as f: base_num = len(pos) // 1000 * 1000 s1 = '整理前的数据：neg数量：{} pos数量：{} part数量:{} landmark: {} 基数:{}'.format(len(neg), len(pos), len(part), len(landmark), base_num) print(s1) # 打乱写入的数据顺序，并这里这里设置比例，设置size参数的比例就能得到数据集比例, 论文比例为：3:1:1:2 neg_keep = npr.choice(len(neg), size=base_num * 3, replace=base_num * 3 > len(neg)) part_keep = npr.choice(len(part), size=base_num, replace=base_num > len(part)) pos_keep = npr.choice(len(pos), size=base_num, replace=base_num > len(pos)) landmark_keep = npr.choice(len(landmark), size=base_num * 2, replace=base_num * 2 > len(landmark)) s2 = '整理后的数据：neg数量：{} pos数量：{} part数量:{} landmark数量：{}'.format(len(neg_keep), len(pos_keep), len(part_keep), len(landmark_keep)) print(s2) with open(os.path.join(data_dir, 'temp.txt'), 'a', encoding='utf-8') as f_temp: f_temp.write('%s\n' % s1) f_temp.