From 19fe6d1435663ae4e195effcd7f21f002d92e843 Mon Sep 17 00:00:00 2001 From: wailovet Date: Mon, 1 Jul 2024 23:49:58 +0800 Subject: [PATCH] 'update' --- __init__.py | 9 +- mz_train_tools_core.py | 11 +- mz_train_tools_core_HYDiT.py | 14 +- mz_train_tools_utils.py | 22 ++- tools/aesthetic.py | 210 +++++++++++++++++++++++++++ tools/aesthetic2_5.py | 260 ++++++++++++++++++++++++++++++++++ tools/face_resize_copy_dir.py | 47 +++--- 7 files changed, 540 insertions(+), 33 deletions(-) create mode 100644 tools/aesthetic.py create mode 100644 tools/aesthetic2_5.py diff --git a/__init__.py b/__init__.py index 7397c2a..2d9847c 100644 --- a/__init__.py +++ b/__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] diff --git a/mz_train_tools_core.py b/mz_train_tools_core.py index 65460a2..db1a3e0 100644 --- a/mz_train_tools_core.py +++ b/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) diff --git a/mz_train_tools_core_HYDiT.py b/mz_train_tools_core_HYDiT.py index 2a3f4d4..e0e0d50 100644 --- a/mz_train_tools_core_HYDiT.py +++ b/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文件 diff --git a/mz_train_tools_utils.py b/mz_train_tools_utils.py index b21b147..04c099d 100644 --- a/mz_train_tools_utils.py +++ b/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 \ No newline at end of file + return folders diff --git a/tools/aesthetic.py b/tools/aesthetic.py new file mode 100644 index 0000000..ee3891d --- /dev/null +++ b/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) + + diff --git a/tools/aesthetic2_5.py b/tools/aesthetic2_5.py new file mode 100644 index 0000000..d51f377 --- /dev/null +++ b/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) + + diff --git a/tools/face_resize_copy_dir.py b/tools/face_resize_copy_dir.py index ac043ff..03eacf5 100644 --- a/tools/face_resize_copy_dir.py +++ b/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__":