From a56d99d1dd2f44f924d0e261ace1accea25f14db Mon Sep 17 00:00:00 2001 From: wailovet Date: Sun, 7 Jul 2024 21:28:33 +0800 Subject: [PATCH] 'update' --- __init__.py | 3 +- mz_train_tools_core.py | 5 +- mz_train_tools_core_HYDiT.py | 6 +- tools/aesthetic.py | 210 -------------------------- tools/aesthetic2_5.py | 260 --------------------------------- tools/face_resize_copy_dir.py | 200 ------------------------- tools/face_resize_copy_dir2.py | 249 ------------------------------- 7 files changed, 9 insertions(+), 924 deletions(-) delete mode 100644 tools/aesthetic.py delete mode 100644 tools/aesthetic2_5.py delete mode 100644 tools/face_resize_copy_dir.py delete mode 100644 tools/face_resize_copy_dir2.py diff --git a/__init__.py b/__init__.py index 852d390..1b73304 100644 --- a/__init__.py +++ b/__init__.py @@ -64,6 +64,7 @@ class MZ_KohyaSSDatasetConfig: "force_clear_only_images": (["enable", "disable"], {"default": "disable"}), "same_caption_generate": (["enable", "disable"], {"default": "disable"}), "same_caption": ("STRING", {"default": "", "dynamicPrompts": True, "multiline": True}), + "image_format": (["png", "jpg", "webp"], {"default": "webp"}), }, "optional": { "conditioning_images": ("IMAGE",), @@ -632,7 +633,7 @@ class MZ_LoadImagesFromDirectoryPath: images = Utils.listdir(image_dir) images = [x for x in images if x.lower().endswith( - ".png") or x.lower().endswith(".jpg")] + ".png") or x.lower().endswith(".jpg") or file.lower().endswith(".webp")] images = [os.path.join(image_dir, x) for x in images] pil_images = [] diff --git a/mz_train_tools_core.py b/mz_train_tools_core.py index 06d3194..a14ff3e 100644 --- a/mz_train_tools_core.py +++ b/mz_train_tools_core.py @@ -135,11 +135,14 @@ def MZ_ImageSelecter_call(args={}): shutil.rmtree(train_images_dir) os.makedirs(train_images_dir, exist_ok=True) + image_format = args.get("image_format") + file_extension = "." + image_format saved_images_path = [] for i, pil_image in enumerate(pil_images): pil_image = Utils.resize_max(pil_image, resolution, resolution) width, height = pil_image.size - filename = hashlib.md5(pil_image.tobytes()).hexdigest() + ".png" + filename = hashlib.md5( + pil_image.tobytes()).hexdigest() + file_extension pil_image.save(os.path.join(train_images_dir, filename)) saved_images_path.append(filename) diff --git a/mz_train_tools_core_HYDiT.py b/mz_train_tools_core_HYDiT.py index f3427fe..4147e61 100644 --- a/mz_train_tools_core_HYDiT.py +++ b/mz_train_tools_core_HYDiT.py @@ -126,7 +126,7 @@ def MZ_HYDiTDatasetConfig_call(args={}): if same_caption != "": # 循环已经保存的图片 for i, filename in enumerate(saved_images_path): - if filename.lower().endswith(".png") or filename.lower().endswith(".jpg"): + if filename.lower().endswith(".png") or filename.lower().endswith(".jpg") or file.lower().endswith(".webp"): 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: @@ -150,7 +150,7 @@ def MZ_HYDiTDatasetConfig_call(args={}): if force_clear_only_images: images_files = Utils.listdir(train_images_dir) for file in images_files: - if file.lower().endswith(".png") or file.lower().endswith(".jpg"): + if file.lower().endswith(".png") or file.lower().endswith(".jpg") or file.lower().endswith(".webp"): os.remove(os.path.join(train_images_dir, file)) else: shutil.rmtree(train_images_dir) @@ -427,7 +427,7 @@ def MZ_HYDiTTrain_call(args={}): writer.writerow(["image_path", "text_zh"]) for filename in full_filenames: print(f"处理文件: {filename}") - if filename.lower().endswith(".png") or filename.lower().endswith(".jpg"): + if filename.lower().endswith(".png") or filename.lower().endswith(".jpg") or file.lower().endswith(".webp"): image_path = os.path.join(workspace_images_dir, filename) pil_image = Image.open(image_path) diff --git a/tools/aesthetic.py b/tools/aesthetic.py deleted file mode 100644 index ee3891d..0000000 --- a/tools/aesthetic.py +++ /dev/null @@ -1,210 +0,0 @@ -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 deleted file mode 100644 index d51f377..0000000 --- a/tools/aesthetic2_5.py +++ /dev/null @@ -1,260 +0,0 @@ -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 deleted file mode 100644 index 03eacf5..0000000 --- a/tools/face_resize_copy_dir.py +++ /dev/null @@ -1,200 +0,0 @@ -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 = 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("-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 - 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) - 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: - 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, 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}) - if len(bboxes) != 1: - continue - # cv2.imshow("image", cv2.cvtColor( - # np.array(image), cv2.COLOR_RGB2BGR)) - # 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__": - main() diff --git a/tools/face_resize_copy_dir2.py b/tools/face_resize_copy_dir2.py deleted file mode 100644 index 79c1253..0000000 --- a/tools/face_resize_copy_dir2.py +++ /dev/null @@ -1,249 +0,0 @@ -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="D:\\数据集\\model.onnx") - 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 - - pbar.update(1) - - file_path = os.path.join(root, file) - - 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()