diff --git a/__init__.py b/__init__.py index 8a53abd..930b314 100644 --- a/__init__.py +++ b/__init__.py @@ -367,8 +367,8 @@ class MZ_LoadImagesFromDirectoryPath: if not os.path.exists(image_dir): return (images,) images = os.listdir(image_dir) - images = [x for x in images if x.endswith( - ".png") or x.endswith(".jpg")] + images = [x for x in images if x.lower().endswith( + ".png") or x.lower().endswith(".jpg")] images = [os.path.join(image_dir, x) for x in images] pil_images = [] @@ -430,6 +430,10 @@ class MZ_HYDiTDatasetConfig: "required": { "workspace_config": ("MZ_TT_HYDiT_WorkspaceConfig",), "resolution": ("INT", {"default": 1024}), + "force_clear": (["enable", "disable"], {"default": "disable"}), + "force_clear_only_images": (["enable", "disable"], {"default": "disable"}), + "same_caption_generate": (["enable", "disable"], {"default": "disable"}), + "same_caption": ("STRING", {"default": "", "dynamicPrompts": True, "multiline": True}), }, "optional": { "images": ("IMAGE",), diff --git a/mz_train_tools_core.py b/mz_train_tools_core.py index 96b3cb5..65460a2 100644 --- a/mz_train_tools_core.py +++ b/mz_train_tools_core.py @@ -120,7 +120,7 @@ def MZ_ImageSelecter_call(args={}): if force_clear: if force_clear_only_images: for file in os.listdir(train_images_dir): - if file.endswith(".png") or file.endswith(".jpg"): + if file.lower().endswith(".png") or file.lower().endswith(".jpg"): os.remove(os.path.join(train_images_dir, file)) else: shutil.rmtree(train_images_dir) @@ -359,7 +359,7 @@ def get_sample_images(train_config): if os.path.exists(sample_images_dir): image_files = os.listdir(sample_images_dir) image_files = list( - filter(lambda x: x.endswith(".png"), image_files)) + filter(lambda x: x.lower().endswith(".png"), image_files)) # 筛选 output_name 前缀 image_files = list( filter(lambda x: x.startswith(output_name), image_files)) diff --git a/mz_train_tools_core_HYDiT.py b/mz_train_tools_core_HYDiT.py index 07b64d1..63a0ebb 100644 --- a/mz_train_tools_core_HYDiT.py +++ b/mz_train_tools_core_HYDiT.py @@ -115,6 +115,17 @@ def MZ_HYDiTDatasetConfig_call(args={}): if workspace_name is None or workspace_name == "": raise Exception("训练名称不能为空(workspace_name is required)") + force_clear = args.get("force_clear") == "enable" + force_clear_only_images = args.get("force_clear_only_images") == "enable" + if force_clear: + if force_clear_only_images: + for file in os.listdir(train_images_dir): + if file.lower().endswith(".png") or file.lower().endswith(".jpg"): + os.remove(os.path.join(train_images_dir, file)) + else: + shutil.rmtree(train_images_dir) + os.makedirs(train_images_dir, exist_ok=True) + saved_images_path = [] for i, pil_image in enumerate(pil_images): @@ -124,6 +135,17 @@ def MZ_HYDiTDatasetConfig_call(args={}): pil_image.save(os.path.join(train_images_dir, filename)) saved_images_path.append(filename) + same_caption_generate = args.get("same_caption_generate") == "enable" + if same_caption_generate: + same_caption = args.get("same_caption").strip() + if same_caption != "": + # 循环已经保存的图片 + for i, filename in enumerate(saved_images_path): + base_filename = os.path.splitext(filename)[0] + caption_filename = base_filename + ".caption" + with open(os.path.join(train_images_dir, caption_filename), "w", encoding="utf-8") as f: + f.write(same_caption) + return ( train_images_dir, ) @@ -346,7 +368,7 @@ def MZ_HYDiTTrain_call(args={}): writer.writerow(["image_path", "text_zh"]) for filename in full_filenames: - if filename.endswith(".png"): + if filename.lower().endswith(".png"): image_path = os.path.join(workspace_images_dir, filename) pil_image = Image.open(image_path) @@ -596,7 +618,7 @@ def get_sample_images(train_config): if os.path.exists(sample_images_dir): image_files = os.listdir(sample_images_dir) image_files = list( - filter(lambda x: x.endswith(".png"), image_files)) + filter(lambda x: x.lower().endswith(".png"), image_files)) # 筛选 output_name 前缀 image_files = list( filter(lambda x: x.startswith(output_name), image_files)) diff --git a/mz_train_tools_utils.py b/mz_train_tools_utils.py index b67a9d2..d6f2136 100644 --- a/mz_train_tools_utils.py +++ b/mz_train_tools_utils.py @@ -70,7 +70,7 @@ class Utils: files = os.listdir(path) new_files = [] for file in files: - if file.endswith(".png"): + if file.lower().endswith(".png"): new_files.append(file) files = new_files files.sort(key=lambda x: int(os.path.basename(x).split(".")[0])) diff --git a/tools/face_resize_copy_dir.py b/tools/face_resize_copy_dir.py new file mode 100644 index 0000000..ac043ff --- /dev/null +++ b/tools/face_resize_copy_dir.py @@ -0,0 +1,191 @@ +import argparse +import hashlib +import io +import os +from PIL import Image +import numpy as np +import onnxruntime as ort +import cv2 + + +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 = newHeight // 8 * 8 + newWidth = 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("-s", "--score", type=float, default=0.8) + args = parser.parse_args() + limitscore = args.score + sess = ort.InferenceSession(args.face_model_path) + sess.set_providers(['CUDAExecutionProvider']) + input_dir = args.input_dir + output_dir = args.output_dir + + total = 0 + for root, dirs, files in os.walk(input_dir): + for file in files: + if not file.endswith(".png") and not file.endswith(".jpg"): + continue + total += 1 + + import tqdm + + with tqdm.tqdm(total=total) as pbar: + 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) + pbar.set_description(f"Processing {file_path}") + + 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) + image = Utils.resize_min(image, 1280, 1280) + + cv2_img = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR) + + scores, bboxes, keypoints, aligned_imgs, landmarks, affine_matrices = sess.run( + None, {"input": cv2_img}) + if len(bboxes) != 1: + continue + + bboxes = bboxes[0] + scores = scores[0] + if limitscore > scores: + continue + + 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}") + + + center_point = (bboxes[0] + bboxes[2]) / \ + 2, (bboxes[1] + bboxes[3]) / 2 + image = Utils.resize_crop(image, 960, 1280, center_point) + cv2_img = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR) + scores, bboxes, keypoints, aligned_imgs, landmarks, affine_matrices = sess.run( + None, {"input": cv2_img}) + if len(bboxes) != 1: + continue + # cv2.imshow("image", cv2.cvtColor( + # np.array(image), cv2.COLOR_RGB2BGR)) + # cv2.waitKey(0) + + output_path = os.path.join(output_dir, md5_str + ".png") + # print(f"Save {file_path} to {output_path}") + + image.save(output_path) + + +if __name__ == "__main__": + main()