'update'

__init__.py

@@ -102,7 +102,8 @@ class MZ_KohyaSSUseConfig:
 
     @classmethod
     def INPUT_TYPES(s):
-        train_config_templates = os.listdir(s.train_config_template_dir)
+        train_config_templates = Utils.listdir(s.train_config_template_dir) 
+
         # 去掉json后缀
         train_config_templates = [os.path.splitext(x)[0]
                                   for x in train_config_templates]
@@ -240,6 +241,9 @@ class MZ_KohyaSSLoraTrain:
         # 使用walk查询所有的workspace中的所有lora模型,lora存放在每个workspace的output目录下
         workspaces_loras = []
         for root, dirs, files in os.walk(workspaces_dir):
+            
+            # 排除隐藏文件夹    
+            dirs[:] = [d for d in dirs if not d.startswith(".")]
             if root.endswith("output"):
                 for file in files:
                     if file.endswith(".safetensors"):
@@ -366,7 +370,8 @@ class MZ_LoadImagesFromDirectoryPath:
         image_dir = kwargs["directory"]
         if not os.path.exists(image_dir):
             return (images,)
-        images = os.listdir(image_dir)
+        images = Utils.listdir(image_dir) 
+
         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]

mz_train_tools_core.py

@@ -119,7 +119,8 @@ def MZ_ImageSelecter_call(args={}):
     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):
+            images_files = Utils.listdir(train_images_dir)
+            for file in images_files:
                 if file.lower().endswith(".png") or file.lower().endswith(".jpg"):
                     os.remove(os.path.join(train_images_dir, file))
         else:
@@ -357,7 +358,7 @@ def get_sample_images(train_config):
     pil_images = []
     pre_render_texts_x = []
     if os.path.exists(sample_images_dir):
-        image_files = os.listdir(sample_images_dir)
+        image_files = Utils.listdir(sample_images_dir)
         image_files = list(
             filter(lambda x: x.lower().endswith(".png"), image_files))
         # 筛选 output_name 前缀
@@ -483,7 +484,7 @@ def MZ_KohyaSSTrain_call(args={}):
     elif base_lora == "latest":
         workspace_lora_dir = os.path.join(workspace_dir, "output")
         if os.path.exists(workspace_lora_dir):
-            workspace_lora_files = os.listdir(workspace_lora_dir)
+            workspace_lora_files = Utils.listdir(workspace_lora_dir)
             workspace_lora_files = list(
                 filter(lambda x: x.endswith(".safetensors"), workspace_lora_files))
             workspace_lora_files = list(
@@ -520,7 +521,7 @@ def MZ_KohyaSSTrain_call(args={}):
     elif base_controlnet == "latest":
         workspace_controlnet_dir = os.path.join(workspace_dir, "output")
         if os.path.exists(workspace_controlnet_dir):
-            workspace_controlnet_files = os.listdir(workspace_controlnet_dir)
+            workspace_controlnet_files = Utils.listdir(workspace_controlnet_dir)
             workspace_controlnet_files = list(
                 filter(lambda x: x.endswith(".safetensors"), workspace_controlnet_files))
             workspace_controlnet_files = list(
@@ -564,7 +565,7 @@ def MZ_KohyaSSTrain_call(args={}):
             workspace_dir, "conditioning_images")
         conditioning_images_onec = ""
         if os.path.exists(conditioning_images_dir):
-            conditioning_images_onec = os.listdir(conditioning_images_dir)[0]
+            conditioning_images_onec = Utils.listdir(conditioning_images_dir)[0]
             other_config["controlnet_image"] = os.path.join(
                 conditioning_images_dir, conditioning_images_onec)
         

mz_train_tools_core_HYDiT.py

@@ -129,7 +129,8 @@ def MZ_HYDiTDatasetConfig_call(args={}):
     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):
+            images_files = Utils.listdir(train_images_dir)
+            for file in images_files:
                 if file.lower().endswith(".png") or file.lower().endswith(".jpg"):
                     os.remove(os.path.join(train_images_dir, file))
         else:
@@ -376,7 +377,7 @@ def MZ_HYDiTTrain_call(args={}):
     if workspace_images_dir is None:
         workspace_images_dir = os.path.join(workspace_dir, "train_images")
 
-    full_filenames = os.listdir(workspace_images_dir)
+    full_filenames = Utils.listdir(workspace_images_dir) 
 
     # 创建image_text.csv文件
     csv_filename = os.path.join(workspace_dir, "image_text.csv")
@@ -446,7 +447,7 @@ def MZ_HYDiTTrain_call(args={}):
     except Exception as e:
         raise Exception(f"生成arrow文件时出现异常,详细信息请查看控制台...")
 
-    arrows_files = os.listdir(arrows_dir)
+    arrows_files = Utils.listdir(arrows_dir) 
 
     dataset_yaml_path = os.path.join(workspace_dir, "dataset.yaml")
 
@@ -671,9 +672,11 @@ def get_sample_images(train_config):
     pil_images = []
     pre_render_texts_x = []
     if os.path.exists(sample_images_dir):
-        image_files = os.listdir(sample_images_dir)
+        image_files = Utils.listdir(sample_images_dir) 
+        
         image_files = list(
             filter(lambda x: x.lower().endswith(".png"), image_files))
+        
         # 筛选 output_name 前缀
         image_files = list(
             filter(lambda x: x.startswith(output_name), image_files))
@@ -1075,6 +1078,9 @@ def search_loras(wdirs):
     loras = []
     for wdir in wdirs:
         for root, dirs, files in os.walk(wdir):
+            # 排除隐藏文件夹    
+            dirs[:] = [d for d in dirs if not d.startswith('.')]
+
             for file in files:
                 if file.lower().endswith(".safetensors"):
                     # 并且在同一目录下有相同名字的json文件

mz_train_tools_utils.py

@@ -67,7 +67,7 @@ class Utils:
 
     def listdir_png(path):
         try:
-            files = os.listdir(path)
+            files = Utils.listdir(path)
             new_files = []
             for file in files:
                 if file.lower().endswith(".png"):
@@ -78,6 +78,15 @@ class Utils:
         except Exception as e:
             return []
 
+    def listdir(path):
+        try:
+            files = os.listdir(path)
+            # 排除.开头的文件
+            files = [file for file in files if not file.startswith(".")]
+            return files
+        except Exception as e:
+            return []
+
     def tensor2pil(image):
         return Image.fromarray(np.clip(255.0 * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8))
 
@@ -676,6 +685,8 @@ class Utils:
 
             if os.path.exists(find_fullpath):
                 for root, dirs, files in os.walk(find_fullpath):
+                    # 排除隐藏文件夹
+                    dirs[:] = [d for d in dirs if not d.startswith('.')]
                     for file in files:
                         if target_sha256 == Utils.file_sha256(os.path.join(root, file)):
                             return os.path.join(root, file)
@@ -919,14 +930,19 @@ class Utils:
     def get_models_by_folder(dir_path):
         models = []
         for root, dirs, files in os.walk(dir_path):
+            # 排除隐藏文件夹    
+            dirs[:] = [d for d in dirs if not d.startswith('.')]
             for file in files:
                 if file.endswith(".pth") or file.endswith(".pt") or file.endswith(".pkl") or file.endswith(".onnx") or file.endswith(".safetensors"):
                     models.append(os.path.join(root, file))
         return models
-    
+
     def get_folders_by_folder(dir_path):
         folders = []
         for root, dirs, files in os.walk(dir_path):
+            
+            # 排除隐藏文件夹    
+            dirs[:] = [d for d in dirs if not d.startswith('.')]
             for dir in dirs:
                 folders.append(os.path.join(root, dir))
-        return folders
+        return folders

tools/aesthetic.py

@@ -0,0 +1,210 @@
+import json
+import os
+import shutil
+from PIL import Image
+import matplotlib.pyplot as plt
+
+from warnings import filterwarnings
+
+import numpy as np
+import torch
+import pytorch_lightning as pl
+import torch.nn as nn
+from PIL import Image
+
+
+from transformers import CLIPProcessor, CLIPModel
+import torch.nn.functional as F
+
+
+class MLP(pl.LightningModule):
+    def __init__(self, input_size, xcol='emb', ycol='avg_rating'):
+        super().__init__()
+        self.input_size = input_size
+        self.xcol = xcol
+        self.ycol = ycol
+        self.layers = nn.Sequential(
+            nn.Linear(self.input_size, 1024),
+            # nn.ReLU(),
+            nn.Dropout(0.2),
+            nn.Linear(1024, 128),
+            # nn.ReLU(),
+            nn.Dropout(0.2),
+            nn.Linear(128, 64),
+            # nn.ReLU(),
+            nn.Dropout(0.1),
+
+            nn.Linear(64, 16),
+            # nn.ReLU(),
+
+            nn.Linear(16, 1)
+        )
+
+    def forward(self, x):
+        return self.layers(x)
+
+    def training_step(self, batch, batch_idx):
+        x = batch[self.xcol]
+        y = batch[self.ycol].reshape(-1, 1)
+        x_hat = self.layers(x)
+        loss = F.mse_loss(x_hat, y)
+        return loss
+
+    def validation_step(self, batch, batch_idx):
+        x = batch[self.xcol]
+        y = batch[self.ycol].reshape(-1, 1)
+        x_hat = self.layers(x)
+        loss = F.mse_loss(x_hat, y)
+        return loss
+
+    def configure_optimizers(self):
+        optimizer = torch.optim.Adam(self.parameters(), lr=1e-3)
+        return optimizer
+
+
+def normalized(a, axis=-1, order=2):
+    import numpy as np
+
+    l2 = np.atleast_1d(np.linalg.norm(a, order, axis))
+    l2[l2 == 0] = 1
+    return a / np.expand_dims(l2, axis)
+
+
+class Aestheic:
+    def __init__(self, clip_path, aes_model):
+
+        model = MLP(768)  # CLIP embedding dim is 768 for CLIP ViT L 14
+
+        # load the model you trained previously or the model available in this repo
+        s = torch.load(aes_model)
+
+        model.load_state_dict(s)
+
+        model.to("cuda")
+        model.eval()
+
+        self.model = model
+
+        self.device = "cuda"
+
+        model2 = CLIPModel.from_pretrained(
+            clip_path, local_files_only=True).to("cuda")
+        preprocess = CLIPProcessor.from_pretrained(
+            clip_path, local_files_only=True)
+
+        self.model2 = model2
+        self.preprocess = preprocess
+
+    def prediction(self, img_paths):
+        if not isinstance(img_paths, list):
+            img_paths = [img_paths]
+
+        pil_images = []
+        for img_path in img_paths:
+            pil_image = Image.open(img_path)
+            pil_images.append(pil_image)
+
+        inputs = self.preprocess(
+            images=pil_images, return_tensors="pt",
+        ).to("cuda")
+        with torch.no_grad():
+            image_features = self.model2.get_image_features(**inputs)
+            print(image_features.shape)
+
+        im_emb_arr = normalized(image_features.cpu().detach().numpy())
+
+        prediction = self.model(torch.from_numpy(
+            im_emb_arr).to(self.device).type(torch.cuda.FloatTensor))
+
+        prediction_score = prediction.cpu().detach().numpy()
+        prediction_score = prediction_score.flatten()
+
+        images_score = {
+            img_path: float(score) for img_path, score in zip(img_paths, prediction_score)
+        }
+
+        # 按value排序
+        images_score = dict(
+            sorted(images_score.items(), key=lambda x: x[1], reverse=True))
+
+        return images_score
+
+
+import argparse
+if __name__ == "__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("-c", "--clip_path", type=str,
+                        default="D:\\数据集\\clip-vit-large-patch14")
+    parser.add_argument("-m", "--aes_model", type=str,
+                        default="D:\\数据集\\sac+logos+ava1-l14-linearMSE.pth")
+    parser.add_argument("-s", "--score", type=float, default=5.9)
+
+    parser.add_argument("-b", "--batch_size", type=int, default=1)
+
+    args = parser.parse_args()
+    limitscore = args.score
+
+    aes = Aestheic(args.clip_path, args.aes_model)
+
+    input_dir = args.input_dir
+    output_dir = args.output_dir
+    batch_size = args.batch_size
+    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}")
+
+    images = []
+
+    with tqdm.tqdm(total=total) as pbar:
+        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
+                file_path = os.path.join(root, file)
+                pbar.update(1)
+
+                pbar.set_description(
+                    f"Processing {os.path.basename(file_path)}")
+
+                if len(images) < batch_size:
+                    images.append(file_path)
+                else:
+                    image_scroe_map = aes.prediction(images)
+                    images = []
+                    # print(scroe)
+                    # print(json.dumps(scroe, indent=4, ensure_ascii=False))
+                    # exit(0)
+
+                    for img_path, score in image_scroe_map.items():
+                        print(f"{img_path}: {score}")
+
+                        if score > limitscore:
+                            if not os.path.exists(output_dir):
+                                os.makedirs(output_dir)
+                            shutil.copy(img_path, output_dir)
+
+    if len(images) > 0:
+        aes.prediction(images)
+        images = []
+
+        for img_path, score in image_scroe_map.items():
+            print(f"{img_path}: {score}")
+
+            if score > limitscore:
+                if not os.path.exists(output_dir):
+                    os.makedirs(output_dir)
+                shutil.copy(img_path, output_dir)
+
+                

tools/aesthetic2_5.py

@@ -0,0 +1,260 @@
+import json
+import os
+from collections import OrderedDict
+from os import PathLike
+import shutil
+from typing import Final
+
+import torch
+import torch.nn as nn
+from transformers import (
+    SiglipImageProcessor,
+    SiglipVisionConfig,
+    SiglipVisionModel,
+    logging,
+)
+from transformers.image_processing_utils import BatchFeature
+from transformers.modeling_outputs import ImageClassifierOutputWithNoAttention
+
+logging.set_verbosity_error()
+
+URL: Final[str] = (
+    "https://github.com/discus0434/aesthetic-predictor-v2-5/raw/main/models/aesthetic_predictor_v2_5.pth"
+)
+
+
+class AestheticPredictorV2_5Head(nn.Module):
+    def __init__(self, config: SiglipVisionConfig) -> None:
+        super().__init__()
+        self.scoring_head = nn.Sequential(
+            nn.Linear(config.hidden_size, 1024),
+            nn.Dropout(0.5),
+            nn.Linear(1024, 128),
+            nn.Dropout(0.5),
+            nn.Linear(128, 64),
+            nn.Dropout(0.5),
+            nn.Linear(64, 16),
+            nn.Dropout(0.2),
+            nn.Linear(16, 1),
+        )
+
+    def forward(self, image_embeds: torch.Tensor) -> torch.Tensor:
+        return self.scoring_head(image_embeds)
+
+
+class AestheticPredictorV2_5Model(SiglipVisionModel):
+    PATCH_SIZE = 14
+
+    def __init__(self, config: SiglipVisionConfig, *args, **kwargs) -> None:
+        super().__init__(config, *args, **kwargs)
+        self.layers = AestheticPredictorV2_5Head(config)
+        self.post_init()
+
+    def forward(
+        self,
+        pixel_values: torch.FloatTensor | None = None,
+        labels: torch.Tensor | None = None,
+        return_dict: bool | None = None,
+    ) -> tuple | ImageClassifierOutputWithNoAttention:
+        return_dict = (
+            return_dict if return_dict is not None else self.config.use_return_dict
+        )
+
+        outputs = super().forward(
+            pixel_values=pixel_values,
+            return_dict=return_dict,
+        )
+        image_embeds = outputs.pooler_output
+        image_embeds_norm = image_embeds / \
+            image_embeds.norm(dim=-1, keepdim=True)
+        prediction = self.layers(image_embeds_norm)
+
+        loss = None
+        if labels is not None:
+            loss_fct = nn.MSELoss()
+            loss = loss_fct()
+
+        if not return_dict:
+            return (loss, prediction, image_embeds)
+
+        return ImageClassifierOutputWithNoAttention(
+            loss=loss,
+            logits=prediction,
+            hidden_states=image_embeds,
+        )
+
+
+class AestheticPredictorV2_5Processor(SiglipImageProcessor):
+    def __init__(self, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+
+    def __call__(self, *args, **kwargs) -> BatchFeature:
+        return super().__call__(*args, **kwargs)
+
+    @classmethod
+    def from_pretrained(
+        self,
+        pretrained_model_name_or_path: str
+        | PathLike = "google/siglip-so400m-patch14-384",
+        *args,
+        **kwargs,
+    ) -> "AestheticPredictorV2_5Processor":
+        return super().from_pretrained(pretrained_model_name_or_path, *args, **kwargs)
+
+
+def convert_v2_5_from_siglip(
+    predictor_name_or_path: str | PathLike | None = None,
+    encoder_model_name: str = "google/siglip-so400m-patch14-384",
+    *args,
+    **kwargs,
+) -> tuple[AestheticPredictorV2_5Model, AestheticPredictorV2_5Processor]:
+    model = AestheticPredictorV2_5Model.from_pretrained(
+        encoder_model_name, *args, **kwargs
+    )
+
+    processor = AestheticPredictorV2_5Processor.from_pretrained(
+        encoder_model_name, *args, **kwargs
+    )
+
+    if predictor_name_or_path is None or not os.path.exists(predictor_name_or_path):
+        state_dict = torch.hub.load_state_dict_from_url(
+            URL, map_location="cpu")
+    else:
+        state_dict = torch.load(predictor_name_or_path, map_location="cpu")
+
+    assert isinstance(state_dict, OrderedDict)
+
+    model.layers.load_state_dict(state_dict)
+    model.eval()
+
+    return model, processor
+
+
+from PIL import Image
+
+
+class AestheticPredictor:
+    def __init__(self, clip_path, aesthetic_predictor_path):
+        # load model and preprocessor
+        self.model, self.preprocessor = convert_v2_5_from_siglip(
+            low_cpu_mem_usage=True,
+            trust_remote_code=True,
+            encoder_model_name=clip_path,
+            predictor_name_or_path=aesthetic_predictor_path,
+        )
+        if torch.cuda.is_available():
+            self.model = self.model.to(torch.bfloat16).cuda()
+
+    def inference(self, img_paths) -> float:
+        if not isinstance(img_paths, list):
+            img_paths = [img_paths]
+
+        pil_images = []
+        for img_path in img_paths:
+            pil_image = Image.open(img_path)
+            pil_images.append(pil_image)
+
+        # preprocess image
+        pixel_values = self.preprocessor(
+            images=pil_images, return_tensors="pt"
+        ).pixel_values
+
+        if torch.cuda.is_available():
+            pixel_values = pixel_values.to(torch.bfloat16).cuda()
+
+        # predict aesthetic score
+        with torch.inference_mode():
+            score = self.model(
+                pixel_values).logits.squeeze().float().cpu().numpy()
+        prediction_score = score.flatten()
+
+        images_score = {
+            img_path: float(score) for img_path, score in zip(img_paths, prediction_score)
+        }
+
+        # 按value排序
+        images_score = dict(
+            sorted(images_score.items(), key=lambda x: x[1], reverse=True))
+
+        return images_score
+
+
+
+import argparse
+if __name__ == "__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("-c", "--clip_path", type=str,
+                        default="D:\\数据集\\siglip-so400m-patch14-384")
+    parser.add_argument("-m", "--aes_model", type=str,
+                        default="D:\\数据集\\aesthetic_predictor_v2_5.pth")
+    parser.add_argument("-s", "--score", type=float, default=5.9)
+
+    parser.add_argument("-b", "--batch_size", type=int, default=1)
+
+    args = parser.parse_args()
+    limitscore = args.score
+
+    aes = AestheticPredictor(args.clip_path, args.aes_model)
+
+    input_dir = args.input_dir
+    output_dir = args.output_dir
+    batch_size = args.batch_size
+    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}")
+
+    images = []
+
+    with tqdm.tqdm(total=total) as pbar:
+        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
+                file_path = os.path.join(root, file)
+                pbar.update(1)
+
+                pbar.set_description(
+                    f"Processing {os.path.basename(file_path)}")
+
+                if os.path.exists(os.path.join(output_dir, os.path.basename(file_path))):
+                    continue
+
+                if len(images) < batch_size:
+                    images.append(file_path)
+                else:
+                    image_scroe_map = aes.inference(images)
+                    images = []
+                    for img_path, score in image_scroe_map.items():
+                        print(f"{img_path}: {score}")
+
+                        if score > limitscore:
+                            if not os.path.exists(output_dir):
+                                os.makedirs(output_dir)
+                            shutil.copy(img_path, output_dir)
+
+    if len(images) > 0:
+        image_scroe_map = aes.inference(images)
+        images = []
+
+        for img_path, score in image_scroe_map.items():
+            print(f"{img_path}: {score}")
+
+            if score > limitscore:
+                if not os.path.exists(output_dir):
+                    os.makedirs(output_dir)
+                shutil.copy(img_path, output_dir)
+
+                

tools/face_resize_copy_dir.py

@@ -64,20 +64,18 @@ class Utils:
             newWidth = dst_h * src_w // src_h
             newHeight = dst_h
 
-        newHeight = newHeight // 8 * 8
-        newWidth = newWidth // 8 * 8
+        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 
+        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:
@@ -101,32 +99,41 @@ def main():
     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("-s", "--score", type=float, default=0.8)
+    parser.add_argument("-f", "--format", type=str, default="jpg")
+
     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
-
+    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
 
+    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)
-                pbar.set_description(f"Processing {file_path}")
 
                 with open(file_path, "rb") as f:
                     md5_str = Utils.Md5(f.read())
@@ -134,10 +141,11 @@ def main():
                     continue
 
                 image = Image.open(file_path)
-                image = Utils.resize_min(image, 1280, 1280)
+                image = Utils.resize_min(image, resize, resize)
 
                 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:
@@ -167,11 +175,11 @@ def main():
                     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)
+                image = Utils.resize_crop(
+                    image, resize * 3 / 4, resize, 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})
@@ -179,12 +187,13 @@ def main():
                     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)
+                # 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)
 
 
 if __name__ == "__main__":