from PIL import Image from insightface.app import FaceAnalysis import numpy as np import os from pathlib import Path import time import argparse import cv2 class FaceInference: """人脸检测推理类，封装insightface的推理功能""" def __init__(self, det_thresh=0.5, det_size=(640, 640), ctx_id=0): """ 初始化人脸检测器 Args: det_thresh: 检测阈值 det_size: 检测图像尺寸 ctx_id: GPU设备ID，如果为-1则使用CPU，否则使用GPU """ # 如果 ctx_id 为 -1，使用 CPU；否则使用 GPU if ctx_id == -1: providers = ['CPUExecutionProvider'] provider_options = [{}] ctx_id = -1 # InsightFace 使用 -1 表示 CPU else: providers = ['CUDAExecutionProvider', 'CPUExecutionProvider'] provider_options = [{"device_id": str(ctx_id)}, {}] self.face_analysis = FaceAnalysis( allowed_modules=['detection'], providers=providers, provider_options=provider_options, ) self.face_analysis.prepare(ctx_id=ctx_id, det_thresh=det_thresh, det_size=det_size) def _make_square_bbox(self, x1, y1, x2, y2, image_width, image_height): """ 将矩形bbox转换为方形bbox，保持人脸比例不变 Args: x1, y1, x2, y2: 原始bbox坐标 image_width, image_height: 图像尺寸 Returns: tuple: (new_x1, new_y1, new_x2, new_y2) 方形bbox坐标 """ # 计算原始bbox的中心点和尺寸 center_x = (x1 + x2) / 2 center_y = (y1 + y2) / 2 width = x2 - x1 height = y2 - y1 # 取较大的边作为方形的边长 square_size = max(width, height) # 计算方形bbox的坐标 half_size = square_size / 2 new_x1 = center_x - half_size new_y1 = center_y - half_size new_x2 = center_x + half_size new_y2 = center_y + half_size # 处理边界情况，确保方形bbox在图像范围内 if new_x1 < 0: new_x1 = 0 new_x2 = square_size if new_y1 < 0: new_y1 = 0 new_y2 = square_size if new_x2 > image_width: new_x2 = image_width new_x1 = image_width - square_size if new_y2 > image_height: new_y2 = image_height new_y1 = image_height - square_size # 再次确保坐标在有效范围内 new_x1 = max(0, new_x1) new_y1 = max(0, new_y1) new_x2 = min(image_width, new_x2) new_y2 = min(image_height, new_y2) return new_x1, new_y1, new_x2, new_y2 def infer_from_array(self, image_array, n=None): """ 对输入numpy数组进行人脸检测推理 Args: image_array: numpy数组，形状为[H, W, 3]，值范围为0-255 n: 选择前n个最大的人脸，如果为None则选择所有人脸 Returns: dict: 包含检测结果的字典，格式为： { 'faces': 检测到的人脸列表, 'bboxes': bbox列表，每个元素为[x, y, width, height], 'masks': mask列表，每个元素为单通道mask图像, 'masked_images': masked图像列表，每个元素为应用mask后的图像, 'image_shape': 原始图像的形状 (height, width, channels) } 如果未检测到人脸，返回中心区域矩形作为默认bbox """ try: if image_array is None: print("错误：输入图像数组为空") return {} # 确保图像数组是正确的格式 if len(image_array.shape) != 3 or image_array.shape[2] != 3: print(f"错误：图像数组形状不正确，期望[H, W, 3]，实际{image_array.shape}") return {} # 确保数据类型和值范围正确 if image_array.dtype != np.uint8: image_array = image_array.astype(np.uint8) faces = self.face_analysis.get(image_array) height, width = image_array.shape[:2] if not faces: return { 'faces': [], 'bboxes': [], 'masks': [], 'masked_images': [], 'image_shape': image_array.shape } # 先按人脸面积大小排序，选择前n个最大的人脸 if n is not None and n > 0: # 计算每个人脸的面积并排序 faces_with_area = [(face, (face['bbox'][2] - face['bbox'][0]) * (face['bbox'][3] - face['bbox'][1])) for face in faces] faces_with_area.sort(key=lambda x: x[1], reverse=True) # 按面积从大到小排序 faces = [face for face, _ in faces_with_area[:n]] # 取前n个最大的人脸 # print(f"选择了前{n}个最大的人脸，总面积分别为: {[area for _, area in faces_with_area[:n]]}") # 再按x坐标从左到右排序 faces = sorted(faces, key=lambda x: x['bbox'][0]) # 生成bbox、mask和masked图像 bboxes = [] masks = [] masked_images = [] for i, face in enumerate(faces): bbox = face['bbox'] x1, y1, x2, y2 = bbox[0], bbox[1], bbox[2], bbox[3] # 将矩形bbox转换为方形bbox square_x1, square_y1, square_x2, square_y2 = self._make_square_bbox( x1, y1, x2, y2, width, height ) # 创建方形mask mask = np.zeros(image_array.shape[:2], dtype=np.uint8) mask[int(square_y1):int(square_y2), int(square_x1):int(square_x2)] = 1.0 # 创建mask与原图相乘的结果 masked_image = image_array.copy() masked_image = cv2.bitwise_and(masked_image, masked_image, mask=mask) bboxes.append([square_x1, square_y1, square_x2 - square_x1, square_y2 - square_y1]) masks.append(mask) masked_images.append(masked_image) # print(f" 人脸 {i+1}: 原始bbox=[{x1:.1f}, {y1:.1f}, {x2-x1:.1f}, {y2-y1:.1f}] -> 方形bbox=[{square_x1:.1f}, {square_y1:.1f}, {square_x2-square_x1:.1f}, {square_y2-square_y1:.1f}]") return { 'faces': faces, 'bboxes': bboxes, 'masks': masks, 'masked_images': masked_images, 'image_shape': image_array.shape } except Exception as e: print(f"处理图像数组时出错: {str(e)}") # 异常情况下也返回中心区域 if 'image_array' in locals() and image_array is not None: return { 'faces': [], 'bboxes': [], 'masks': [], 'masked_images': [], 'image_shape': image_array.shape } return {} def infer(self, image_path, n=None): """ 对输入图像进行人脸检测推理 Args: image_path: 图像文件路径或图片 n: 选择前n个最大的人脸，如果为None则选择所有人脸 Returns: dict: 包含检测结果的字典，格式为： { 'faces': 检测到的人脸列表, 'bboxes': bbox列表，每个元素为[x, y, width, height], 'masks': mask列表，每个元素为单通道mask图像, 'masked_images': masked图像列表，每个元素为应用mask后的图像, 'image_shape': 原始图像的形状 (height, width, channels) } 如果未检测到人脸，返回中心区域矩形作为默认bbox """ try: image = cv2.imread(image_path) if image is None: print(f"错误：无法读取图像 {image_path}") # 无法读取图像，返回空结果 return {} faces = self.face_analysis.get(image) height, width = image.shape[:2] if not faces: print(f"警告：图像 {os.path.basename(image_path)} 中未检测到人脸，使用中心区域作为默认方形bbox") # 计算中心区域方形（边长为原图较小边的50%） min_dim = min(width, height) square_size = min_dim // 2 center_x, center_y = width // 2, height // 2 x1 = center_x - square_size // 2 y1 = center_y - square_size // 2 x2 = x1 + square_size y2 = y1 + square_size # 确保bbox在图像范围内 x1 = max(0, x1) y1 = max(0, y1) x2 = min(width, x2) y2 = min(height, y2) # 创建中心区域的方形mask mask = np.zeros(image.shape[:2], dtype=np.uint8) mask[int(y1):int(y2), int(x1):int(x2)] = 1.0 # 创建masked图像 masked_image = image.copy() masked_image = cv2.bitwise_and(masked_image, masked_image, mask=mask) return { 'faces': [], 'bboxes': [[x1, y1, x2 - x1, y2 - y1]], 'masks': [mask], 'masked_images': [masked_image], 'image_shape': image.shape } # 先按人脸面积大小排序，选择前n个最大的人脸 if n is not None and n > 0: # 计算每个人脸的面积并排序 faces_with_area = [(face, (face['bbox'][2] - face['bbox'][0]) * (face['bbox'][3] - face['bbox'][1])) for face in faces] faces_with_area.sort(key=lambda x: x[1], reverse=True) # 按面积从大到小排序 faces = [face for face, _ in faces_with_area[:n]] # 取前n个最大的人脸 # 再按x坐标从左到右排序 faces = sorted(faces, key=lambda x: x['bbox'][0]) # 生成bbox、mask和masked图像 bboxes = [] masks = [] masked_images = [] for i, face in enumerate(faces): bbox = face['bbox'] x1, y1, x2, y2 = bbox[0], bbox[1], bbox[2], bbox[3] # 将矩形bbox转换为方形bbox square_x1, square_y1, square_x2, square_y2 = self._make_square_bbox( x1, y1, x2, y2, width, height ) # 创建方形mask mask = np.zeros(image.shape[:2], dtype=np.uint8) mask[int(square_y1):int(square_y2), int(square_x1):int(square_x2)] = 1.0 # 创建mask与原图相乘的结果 masked_image = image.copy() masked_image = cv2.bitwise_and(masked_image, masked_image, mask=mask) bboxes.append([square_x1, square_y1, square_x2 - square_x1, square_y2 - square_y1]) masks.append(mask) masked_images.append(masked_image) return { 'faces': faces, 'bboxes': bboxes, 'masks': masks, 'masked_images': masked_images, 'image_shape': image.shape } except Exception as e: print(f"处理图像 {image_path} 时出错: {str(e)}") # 异常情况下也返回中心区域方形 if 'image' in locals() and image is not None: height, width = image.shape[:2] # 计算中心区域方形（边长为原图较小边的50%） min_dim = min(width, height) square_size = min_dim // 2 center_x, center_y = width // 2, height // 2 x1 = center_x - square_size // 2 y1 = center_y - square_size // 2 x2 = x1 + square_size y2 = y1 + square_size # 确保bbox在图像范围内 x1 = max(0, x1) y1 = max(0, y1) x2 = min(width, x2) y2 = min(height, y2) # 创建中心区域的方形mask mask = np.zeros(image.shape[:2], dtype=np.uint8) mask[int(y1):int(y2), int(x1):int(x2)] = 1.0 # 创建masked图像 masked_image = image.copy() masked_image = cv2.bitwise_and(masked_image, masked_image, mask=mask) return { 'faces': [], 'bboxes': [[x1, y1, x2 - x1, y2 - y1]], 'masks': [mask], 'masked_images': [masked_image], 'image_shape': image.shape } return {} class FaceProcessor: def __init__(self, det_thresh=0.5, det_size=(640, 640)): self.face_analysis = FaceAnalysis(allowed_modules=['detection']) self.face_analysis.prepare(ctx_id=0, det_thresh=det_thresh, det_size=det_size) def _make_square_bbox(self, x1, y1, x2, y2, image_width, image_height): """ 将矩形bbox转换为方形bbox，保持人脸比例不变 Args: x1, y1, x2, y2: 原始bbox坐标 image_width, image_height: 图像尺寸 Returns: tuple: (new_x1, new_y1, new_x2, new_y2) 方形bbox坐标 """ # 计算原始bbox的中心点和尺寸 center_x = (x1 + x2) / 2 center_y = (y1 + y2) / 2 width = x2 - x1 height = y2 - y1 # 取较大的边作为方形的边长 square_size = max(width, height) # 计算方形bbox的坐标 half_size = square_size / 2 new_x1 = center_x - half_size new_y1 = center_y - half_size new_x2 = center_x + half_size new_y2 = center_y + half_size # 处理边界情况，确保方形bbox在图像范围内 if new_x1 < 0: new_x1 = 0 new_x2 = square_size if new_y1 < 0: new_y1 = 0 new_y2 = square_size if new_x2 > image_width: new_x2 = image_width new_x1 = image_width - square_size if new_y2 > image_height: new_y2 = image_height new_y1 = image_height - square_size # 再次确保坐标在有效范围内 new_x1 = max(0, new_x1) new_y1 = max(0, new_y1) new_x2 = min(image_width, new_x2) new_y2 = min(image_height, new_y2) return new_x1, new_y1, new_x2, new_y2 def get_face_bbox_and_mask(self, image): faces = self.face_analysis.get(image) if not faces: print("警告：图像中未检测到人脸。") return None, None, None # 按x坐标从左到右排序 faces = sorted(faces, key=lambda x: x['bbox'][0]) height, width = image.shape[:2] bboxes = [] masks = [] masked_images = [] for i, face in enumerate(faces): bbox = face['bbox'] x1, y1, x2, y2 = bbox[0], bbox[1], bbox[2], bbox[3] # 将矩形bbox转换为方形bbox square_x1, square_y1, square_x2, square_y2 = self._make_square_bbox( x1, y1, x2, y2, width, height ) # 创建方形mask mask = np.zeros(image.shape[:2], dtype=np.uint8) mask[int(square_y1):int(square_y2), int(square_x1):int(square_x2)] = 1.0 # 创建mask与原图相乘的结果 masked_image = image.copy() masked_image = cv2.bitwise_and(masked_image, masked_image, mask=mask) bboxes.append([square_x1, square_y1, square_x2 - square_x1, square_y2 - square_y1]) masks.append(mask) masked_images.append(masked_image) return bboxes, masks, masked_images def main(): parser = argparse.ArgumentParser(description='Process images to detect faces and save bbox, mask, and masked images.') parser.add_argument('--input_dir', type=str, default="./data/bbox_test_input", help='Directory containing input images.') parser.add_argument('--bbox_output_dir', type=str, default="./temp/bbox", help='Directory to save bbox npy files.') parser.add_argument('--mask_output_dir', type=str, default="./temp/mask", help='Directory to save mask images.') parser.add_argument('--masked_image_output_dir', type=str, default="./temp/masked_images", help='Directory to save masked images.') args = parser.parse_args() # 创建输出目录 os.makedirs(args.bbox_output_dir, exist_ok=True) os.makedirs(args.mask_output_dir, exist_ok=True) os.makedirs(args.masked_image_output_dir, exist_ok=True) # 初始化人脸检测器 face_processor = FaceProcessor() # 支持的图像格式 supported_formats = {'.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.tif'} # 获取所有图像文件 image_files = [] for file in os.listdir(args.input_dir): if Path(file).suffix.lower() in supported_formats: image_files.append(file) if not image_files: print(f"警告：在目录 {args.input_dir} 中未找到支持的图像文件") return # 处理每个图像 for image_file in image_files: image_path = os.path.join(args.input_dir, image_file) # 读取图像 image = cv2.imread(image_path) if image is None: print(f" 错误：无法读取图像 {image_path}") continue # 获取人脸检测结果 bboxes, masks, masked_images = face_processor.get_face_bbox_and_mask(image) if bboxes is None: print(f" 跳过：未检测到人脸") continue # 生成基础文件名（不含扩展名） base_name = Path(image_file).stem # 保存bbox为npy文件 bbox_file = os.path.join(args.bbox_output_dir, f"{base_name}_bbox.npy") np.save(bbox_file, np.array(bboxes)) # 保存mask和masked图像 for i, (mask, masked_image) in enumerate(zip(masks, masked_images)): # 保存mask mask_file = os.path.join(args.mask_output_dir, f"{base_name}_face{i+1}_mask.png") cv2.imwrite(mask_file, mask) # 保存masked图像 masked_image_file = os.path.join(args.masked_image_output_dir, f"{base_name}_face{i+1}_masked.png") cv2.imwrite(masked_image_file, masked_image) print(f" 已保存人脸{i+1}的mask: {mask_file}") print(f" 已保存人脸{i+1}的masked图像: {masked_image_file}") print(f"\n处理完成！") print(f"bbox文件保存在: {args.bbox_output_dir}") print(f"mask文件保存在: {args.mask_output_dir}") print(f"masked图像保存在: {args.masked_image_output_dir}") if __name__ == "__main__": main()