From 2f9518f3cca0b5910f29255eaadce2828f65f236 Mon Sep 17 00:00:00 2001 From: wailovet Date: Wed, 3 Jul 2024 02:49:02 +0800 Subject: [PATCH] 'update' --- mz_train_tools_core_HYDiT.py | 2 +- tools/face_resize_copy_dir2.py | 245 +++++++++++++++++++++++++++++++++ 2 files changed, 246 insertions(+), 1 deletion(-) create mode 100644 tools/face_resize_copy_dir2.py diff --git a/mz_train_tools_core_HYDiT.py b/mz_train_tools_core_HYDiT.py index 3680387..d8cbf7f 100644 --- a/mz_train_tools_core_HYDiT.py +++ b/mz_train_tools_core_HYDiT.py @@ -169,7 +169,7 @@ def MZ_HYDiTDatasetConfig_call(args={}): HYDiT_MODEL = { "HunyuanDiT/t2i/model/pytorch_model_ema.pt": { - "url": "https://www.modelscope.cn/models/wailovet/hy1.1/resolve/master/t2i/model/pytorch_model_module.pt", + "url": "https://www.modelscope.cn/models/wailovet/hy1.1/resolve/master/t2i/model/pytorch_model_ema.pt", }, "HunyuanDiT/t2i/model/pytorch_model_module.pt": { "url": "https://www.modelscope.cn/models/wailovet/hy1.1/resolve/master/t2i/model/pytorch_model_module.pt", diff --git a/tools/face_resize_copy_dir2.py b/tools/face_resize_copy_dir2.py new file mode 100644 index 0000000..7a78823 --- /dev/null +++ b/tools/face_resize_copy_dir2.py @@ -0,0 +1,245 @@ +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()