'update'

__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",),

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))

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))

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]))

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()