'update'
This commit is contained in:
wailovet
+2
-1
@@ -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 = []
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user