import argparse import hashlib import io import os from PIL import Image import numpy as np import onnxruntime as ort import cv2 import piexif class Utils: def Md5(bytes): return hashlib.md5(bytes).hexdigest() def resize_by(image, percent): # 判断类型是否为PIL if not isinstance(image, Image.Image): image = Image.fromarray(image) width, height = image.size new_width = int(width * percent) new_height = int(height * percent) return image.resize((new_width, new_height), Image.LANCZOS) def resize_max(im, dst_w, dst_h): src_w, src_h = im.size if src_h > src_w: newWidth = dst_w newHeight = dst_w * src_h // src_w else: newWidth = dst_h * src_w // src_h newHeight = dst_h newHeight = newHeight // 8 * 8 newWidth = newWidth // 8 * 8 return im.resize((newWidth, newHeight), Image.Resampling.LANCZOS) def resize_min(im, dst_w, dst_h): src_w, src_h = im.size if src_h < src_w: newWidth = dst_w newHeight = dst_w * src_h // src_w else: newWidth = dst_h * src_w // src_h newHeight = dst_h newHeight = newHeight // 8 * 8 newWidth = newWidth // 8 * 8 return im.resize((newWidth, newHeight), Image.Resampling.LANCZOS) # 根据人脸进行剪裁 def resize_crop(im, dst_w, dst_h, focus_point): src_w, src_h = im.size if src_h > src_w: newWidth = dst_w newHeight = dst_w * src_h // src_w else: newWidth = dst_h * src_w // src_h newHeight = dst_h newHeight = int(newHeight // 8 * 8) newWidth = int(newWidth // 8 * 8) im = im.resize((newWidth, newHeight), Image.Resampling.LANCZOS) left = (newWidth - dst_w) // 2 top = (newHeight - dst_h) // 2 right = left + dst_w bottom = top + dst_h focus_point_x, focus_point_y = focus_point offset_w = dst_w // 8 offset_h = dst_h // 8 if focus_point_x < left + dst_w // 8: left = 0 right = dst_w elif focus_point_x > right - dst_w // 8: right = newWidth left = newWidth - dst_w if focus_point_y < top + dst_h // 8: top = 0 bottom = dst_h elif focus_point_y > bottom - dst_h // 8: bottom = newHeight top = newHeight - dst_h return im.crop((left, top, right, bottom)) def main(): parser = argparse.ArgumentParser() parser.add_argument("-i", "--input_dir", type=str, default="") parser.add_argument("-o", "--output_dir", type=str, default="") parser.add_argument("-m", "--face_model_path", type=str, default="") parser.add_argument("-r", "--resize", type=int, default=1024) parser.add_argument("-p", "--face_proportion", type=float, default=0.5) parser.add_argument("-f", "--format", type=str, default="jpg") args = parser.parse_args() sess = ort.InferenceSession(args.face_model_path) sess.set_providers(['CUDAExecutionProvider']) input_dir = args.input_dir output_dir = args.output_dir resize = args.resize img_format = args.format total = 0 for root, dirs, files in os.walk(input_dir): # 排除隐藏文件夹 dirs[:] = [d for d in dirs if not d.startswith('.')] for file in files: if not file.endswith(".png") and not file.endswith(".jpg"): continue total += 1 import tqdm print(f"total: {total}") os.makedirs(output_dir, exist_ok=True) with tqdm.tqdm(total=total) as pbar: # 排除隐藏文件夹 dirs[:] = [d for d in dirs if not d.startswith('.')] for root, dirs, files in os.walk(input_dir): for file in files: if not file.endswith(".png") and not file.endswith(".jpg"): continue file_path = os.path.join(root, file) pbar.update(1) with open(file_path, "rb") as f: md5_str = Utils.Md5(f.read()) if os.path.exists(os.path.join(output_dir, md5_str + ".png")): continue image = Image.open(file_path) max_image = image.copy() image = Utils.resize_min(image, 1920, 1920) zoom = image.size[0] / max_image.size[0] cv2_img = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR) print(f"start check {file_path}") scores, bboxes, keypoints, aligned_imgs, landmarks, affine_matrices = sess.run( None, {"input": cv2_img}) if len(bboxes) != 1: continue bboxes = bboxes[0] affine_matrices = affine_matrices[0] angle = np.degrees(np.arctan2( affine_matrices[1][0], affine_matrices[0][0])) real_angle = 0 if angle < -45 and angle > -135: # 向左旋转90度 real_angle = 90 elif angle > 45 and angle < 135: # 向右旋转90度 real_angle = -90 elif angle < -135 or angle > 135: # 向左旋转180度 real_angle = 180 if real_angle != 0: continue # print(f"{file_path} bboxes: \n{bboxes} {scores} {real_angle}") bboxes = bboxes / zoom image = max_image # 预期人脸高度占图片高度的一定比例 pr = args.face_proportion face_width = bboxes[2] - bboxes[0] pre_image_width = face_width / pr face_center_x = (bboxes[2] + bboxes[0]) / 2 crop_x1 = face_center_x - pre_image_width / 2 crop_x2 = face_center_x + pre_image_width / 2 # 判断是否超出图片边界 if crop_x1 < 0 or crop_x2 > image.width: print( f"crop_x1: {crop_x1}, crop_x2: {crop_x2}, image.width: {image.width}") continue face_center_y = (bboxes[3] + bboxes[1]) / 2 # 总高度和宽度的比例是4:3 pre_image_height = pre_image_width * 4 / 3 # 中点位于高度的0.382比例 crop_y1 = face_center_y - pre_image_height * 0.382 crop_y2 = crop_y1 + pre_image_height # 判断是否超出图片边界 if crop_y1 < 0 or crop_y2 > image.height: print( f"crop_y1: {crop_y1}, crop_y2: {crop_y2}, image.height: {image.height}") continue # 判断是否比resize小 if pre_image_width < resize or pre_image_height < resize: print( f"pre_image_width: {pre_image_width}, pre_image_height: {pre_image_height}") continue # 进行裁剪 image = image.crop((crop_x1, crop_y1, crop_x2, crop_y2)) # 进行缩放并居中剪裁 image = image.resize(( resize * 3 // 4, resize ), Image.LANCZOS) # cv2_img = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR) # cv2.imshow("image", cv2_img) # cv2.waitKey(0) if img_format == "png": output_path = os.path.join(output_dir, md5_str + ".png") image.save(output_path) else: output_path = os.path.join(output_dir, md5_str + ".jpg") image.save(output_path, "jpeg", quality=95) # 插入IPTC # piexif.insert(output_path, piexif.ImageIFD, { # "CameraOwnerName": "CameraOwnerName", # }) print(f"save to {output_path}") if __name__ == "__main__": main()