Compare commits
2 Commits
70bb7fa5ba
...
bf22cbf918
| Author | SHA1 | Date |
|---|---|---|
|
|
bf22cbf918 | |
|
|
a488de2a93 |
|
|
@ -0,0 +1,37 @@
|
|||
import cupy
|
||||
import numpy
|
||||
import os
|
||||
import shutil
|
||||
|
||||
TYPE = "facenet"
|
||||
|
||||
DEST_ROOT = r"D:\helloproject-ai-data\similar_face"
|
||||
SAMPLE_FILE = r"D:\helloproject-ai-data\face_cropped\井上春華\井上春華=morningmusume16ki=12815441476-0-0.jpg"
|
||||
PICK_N = 500
|
||||
PROBE = 0.6
|
||||
embeddings = cupy.load(rf"C:\Users\tomokazu\PycharmProjects\helloproject-ai\embeddings_{TYPE}.npy")
|
||||
embeddings_label = numpy.load(rf"C:\Users\tomokazu\PycharmProjects\helloproject-ai\embeddings_{TYPE}_label.npy")
|
||||
|
||||
sample_pos = numpy.argwhere(embeddings_label == os.path.basename(SAMPLE_FILE))
|
||||
sample_emb = embeddings[*sample_pos, :].T
|
||||
print(sample_pos)
|
||||
print(embeddings_label.shape)
|
||||
print(embeddings.shape)
|
||||
print(sample_emb.shape)
|
||||
similarity: cupy.ndarray = cupy.dot(embeddings, sample_emb).T / (
|
||||
cupy.linalg.norm(embeddings, axis=1) * cupy.linalg.norm(sample_emb))
|
||||
# sample_norm = cupy.linalg.norm(sample_emb)
|
||||
# similarity = cupy.apply_along_axis(lambda x: cupy.dot(x, sample_emb) / (cupy.linalg.norm(x) * sample_norm), 1,
|
||||
# embeddings)
|
||||
print(similarity.shape)
|
||||
similar_pos: cupy.ndarray = similarity.reshape((-1,)).argsort()[::-1]
|
||||
print(similar_pos)
|
||||
p = int(cupy.count_nonzero(cupy.argwhere(similarity > PROBE)))
|
||||
# p = 300
|
||||
sim_list = embeddings_label[cupy.asnumpy(similar_pos)[:p]]
|
||||
print(numpy.stack([sim_list, cupy.asnumpy(similarity.reshape((-1,))[similar_pos])[:p]], axis=1))
|
||||
os.makedirs(os.path.join(DEST_ROOT, os.path.splitext(os.path.basename(SAMPLE_FILE))[0] + f"_{TYPE}"), exist_ok=True)
|
||||
for file in sim_list:
|
||||
shutil.copyfile(os.path.join(r"D:\helloproject-ai-data\face_cropped", file.split("=")[0], file),
|
||||
os.path.join(DEST_ROOT, os.path.splitext(os.path.basename(SAMPLE_FILE))[0] + f"_{TYPE}", file))
|
||||
print(f"Copied {p} file(s).")
|
||||
|
|
@ -0,0 +1,18 @@
|
|||
import gradio
|
||||
import pathlib
|
||||
import cupy
|
||||
import numpy
|
||||
import os
|
||||
|
||||
TYPE = "facenet"
|
||||
|
||||
embeddings = cupy.load(rf"C:\Users\tomokazu\PycharmProjects\helloproject-ai\embeddings_{TYPE}.npy")
|
||||
embeddings_label = numpy.load(rf"C:\Users\tomokazu\PycharmProjects\helloproject-ai\embeddings_{TYPE}_label.npy")
|
||||
|
||||
|
||||
def get_similar_pic(_input: pathlib.Path):
|
||||
search_key = os.path.basename(_input)
|
||||
key_pos = numpy.argwhere(embeddings_label == search_key)
|
||||
search_emb = embeddings[*key_pos, :].T
|
||||
similarity_array: cupy.ndarray = (cupy.dot(embeddings, search_emb).T /
|
||||
(cupy.linalg.norm(embeddings, axis=1) * cupy.linalg.norm(search_emb)))
|
||||
|
|
@ -0,0 +1,83 @@
|
|||
import os.path
|
||||
from io import BytesIO
|
||||
from itertools import chain
|
||||
from PIL import Image
|
||||
from more_itertools import chunked
|
||||
from os import listdir
|
||||
from torchinfo import summary
|
||||
from torchvision import transforms
|
||||
from torchvision.io import decode_jpeg
|
||||
from tqdm import tqdm
|
||||
from facenet_pytorch import InceptionResnetV1
|
||||
import torch
|
||||
import numpy
|
||||
from insightface.app import FaceAnalysis
|
||||
|
||||
CROPPED_DIR = r"D:\helloproject-ai-data\face_cropped"
|
||||
CHUNK_SIZE = 64
|
||||
DEVICE = torch.device("cuda")
|
||||
INPUT_SIZE = 256
|
||||
TYPE = "facenet"
|
||||
face_analysis = FaceAnalysis(providers=[
|
||||
('TensorrtExecutionProvider', {
|
||||
'trt_engine_cache_enable': True,
|
||||
'trt_engine_cache_path': 'trt_cache',
|
||||
'trt_fp16_enable': True,
|
||||
}),
|
||||
'CUDAExecutionProvider',
|
||||
'CPUExecutionProvider',
|
||||
])
|
||||
face_analysis.prepare(ctx_id=0, det_size=(INPUT_SIZE, INPUT_SIZE))
|
||||
|
||||
transform = transforms.Compose([
|
||||
# transforms.ToTensor(),
|
||||
transforms.Resize(size=int(INPUT_SIZE * 1.2), interpolation=transforms.InterpolationMode.BILINEAR),
|
||||
transforms.CenterCrop(size=INPUT_SIZE)
|
||||
# transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
|
||||
])
|
||||
|
||||
# model: torch.nn.Module = get_model(name=MODEL_NAME)
|
||||
# model.load_state_dict(
|
||||
# torch.load(os.path.join(os.path.dirname(__file__), "edgeface", f"{MODEL_NAME}.pt"), weights_only=False))
|
||||
# model = torch.load(r"\\192.168.250.1\share\helloproject-ai-data\artifact\vggface2_facenet.pth", weights_only=False)
|
||||
# model = model.eval().cuda(device=DEVICE)
|
||||
with torch.no_grad():
|
||||
# print(model.eval())
|
||||
# summary(model, input_size=[CHUNK_SIZE, 3, INPUT_SIZE, INPUT_SIZE])
|
||||
# trt_model = torch.compile(model)
|
||||
|
||||
embeddings: numpy.ndarray | None = None
|
||||
labels = []
|
||||
if os.path.exists(f"embeddings_{TYPE}_label.npy"):
|
||||
embeddings: numpy.ndarray = numpy.load(f"embeddings_{TYPE}.npy")
|
||||
labels: list[str] = numpy.load(f"embeddings_{TYPE}_label.npy").tolist()
|
||||
|
||||
all_cropped_list = list(
|
||||
chain.from_iterable([listdir(os.path.join(CROPPED_DIR, name)) for name in listdir(CROPPED_DIR)]))
|
||||
# all_cropped_list = all_cropped_list[:1000]
|
||||
|
||||
labels_set = set(labels)
|
||||
pbar = tqdm(total=all_cropped_list.__len__())
|
||||
|
||||
for file_name in all_cropped_list:
|
||||
sub_dir_name = file_name.split("=")[0]
|
||||
pbar.update(1)
|
||||
if pbar.desc != sub_dir_name:
|
||||
pbar.set_description(sub_dir_name)
|
||||
if file_name in labels_set:
|
||||
continue
|
||||
|
||||
image = numpy.array(Image.open(os.path.join(CROPPED_DIR, sub_dir_name, file_name)))[:, :, [2, 1, 0]]
|
||||
emb = face_analysis.get(image)
|
||||
if not emb:
|
||||
continue
|
||||
if embeddings is not None:
|
||||
embeddings = numpy.concatenate([embeddings, numpy.expand_dims(emb[0].embedding, 0)], axis=0)
|
||||
# print(embeddings.shape)
|
||||
else:
|
||||
embeddings = numpy.expand_dims(emb[0].embedding, 0)
|
||||
labels.append(file_name)
|
||||
|
||||
# print(embeddings.shape, labels.__len__())
|
||||
numpy.save(f"embeddings_{TYPE}.npy", embeddings)
|
||||
numpy.save(f"embeddings_{TYPE}_label.npy", numpy.array(labels))
|
||||
|
|
@ -0,0 +1,96 @@
|
|||
import concurrent.futures
|
||||
import os.path
|
||||
from concurrent.futures.process import ProcessPoolExecutor
|
||||
from itertools import chain
|
||||
from PIL import Image
|
||||
from more_itertools import chunked
|
||||
from os import listdir
|
||||
from torchvision import transforms
|
||||
from torchvision.io import decode_jpeg
|
||||
from tqdm import tqdm
|
||||
from edgeface.backbones import get_model
|
||||
import torch
|
||||
import numpy
|
||||
from edgeface.face_alignment import align
|
||||
|
||||
CROPPED_DIR = r"D:\helloproject-ai-data\face_cropped"
|
||||
MODEL_NAME = "edgeface_s_gamma_05"
|
||||
CHUNK_SIZE = 64
|
||||
DEVICE = torch.device("cuda")
|
||||
INPUT_SIZE = 112
|
||||
TYPE = "edgeface"
|
||||
|
||||
transform = transforms.Compose([
|
||||
transforms.ToTensor(),
|
||||
# lambda x: x.to(torch.float32) / 255.,
|
||||
transforms.Resize(size=INPUT_SIZE, interpolation=transforms.InterpolationMode.BILINEAR),
|
||||
# transforms.Resize(size=int(INPUT_SIZE * 1.2), interpolation=transforms.InterpolationMode.BILINEAR),
|
||||
# transforms.CenterCrop(size=INPUT_SIZE)
|
||||
transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
|
||||
])
|
||||
|
||||
|
||||
def align_edgeface(p: str):
|
||||
sub_dir_name = p.split("=")[0]
|
||||
aligned = align.get_aligned_face(os.path.join(CROPPED_DIR, sub_dir_name, p))
|
||||
if aligned is None:
|
||||
return None
|
||||
return transform(aligned).to(DEVICE)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
model: torch.nn.Module = get_model(name=MODEL_NAME)
|
||||
model.load_state_dict(
|
||||
torch.load(os.path.join(os.path.dirname(__file__), "edgeface", "checkpoints", f"{MODEL_NAME}.pt"),
|
||||
weights_only=False))
|
||||
model = model.eval().cuda(device=DEVICE)
|
||||
with torch.no_grad(), torch.autocast(device_type='cuda', dtype=torch.float16):
|
||||
# print(model.eval())
|
||||
# summary(model, input_size=[CHUNK_SIZE, 3, INPUT_SIZE, INPUT_SIZE])
|
||||
# trt_model = torch.compile(model)
|
||||
|
||||
embeddings: numpy.ndarray | None = None
|
||||
labels = []
|
||||
if os.path.exists(f"embeddings_{TYPE}_label.npy"):
|
||||
embeddings: numpy.ndarray = numpy.load(f"embeddings_{TYPE}.npy")
|
||||
labels: list[str] = numpy.load(f"embeddings_{TYPE}_label.npy").tolist()
|
||||
|
||||
all_cropped_list = list(
|
||||
chain.from_iterable([listdir(os.path.join(CROPPED_DIR, name)) for name in listdir(CROPPED_DIR)]))
|
||||
# all_cropped_list = all_cropped_list[:1000]
|
||||
|
||||
labels_set = set(labels)
|
||||
pbar = tqdm(total=all_cropped_list.__len__())
|
||||
|
||||
for chk in chunked(all_cropped_list, n=CHUNK_SIZE):
|
||||
decoded_images = []
|
||||
pool_res_list = []
|
||||
for file_name in chk:
|
||||
pbar.update(1)
|
||||
if pbar.desc != file_name.split("=")[0]:
|
||||
pbar.set_description(file_name.split("=")[0])
|
||||
if file_name in labels_set:
|
||||
continue
|
||||
pool_res = align_edgeface(file_name)
|
||||
pool_res_list.append(pool_res)
|
||||
for result, name in zip(pool_res_list, chk):
|
||||
# result = result.result()
|
||||
if result is not None:
|
||||
decoded_images.append(result)
|
||||
labels.append(name)
|
||||
|
||||
if not decoded_images:
|
||||
continue
|
||||
stacked = torch.stack(decoded_images)
|
||||
# print(stacked.shape)
|
||||
res = model(stacked)
|
||||
# print(res.shape)
|
||||
if embeddings is not None:
|
||||
embeddings = numpy.concatenate([embeddings, res.cpu().numpy()], axis=0)
|
||||
# print(embeddings.shape)
|
||||
else:
|
||||
embeddings = res.cpu().numpy()
|
||||
|
||||
# print(embeddings.shape, labels.__len__())
|
||||
numpy.save(f"embeddings_{TYPE}.npy", embeddings)
|
||||
numpy.save(f"embeddings_{TYPE}_label.npy", numpy.array(labels))
|
||||
|
|
@ -17,13 +17,13 @@ if not isfile(argv[1]):
|
|||
exit(1)
|
||||
|
||||
face_analysis = FaceAnalysis(providers=[
|
||||
# 'CUDAExecutionProvider',
|
||||
# 'CPUExecutionProvider',
|
||||
('TensorrtExecutionProvider', {
|
||||
'trt_engine_cache_enable': True,
|
||||
'trt_engine_cache_path': join(getcwd(), 'onnx_cache'),
|
||||
'trt_fp16_enable': True,
|
||||
})
|
||||
}),
|
||||
'CUDAExecutionProvider',
|
||||
'CPUExecutionProvider',
|
||||
], allowed_modules=['recognition', 'detection'])
|
||||
face_analysis.prepare(ctx_id=0, det_size=(160, 160))
|
||||
|
||||
|
|
@ -55,7 +55,7 @@ for file in image_files:
|
|||
cosine = dot(emb[0].embedding, collect_image_emb[0].embedding) / \
|
||||
(norm(emb[0].embedding) * norm(collect_image_emb[0].embedding))
|
||||
print(file, cosine)
|
||||
if cosine > 0.3:
|
||||
if cosine > 0.6:
|
||||
copyfile(join(getcwd(), argv[2], file), join(getcwd(), dir_name, "true", file))
|
||||
else:
|
||||
|
||||
|
|
|
|||
|
|
@ -0,0 +1,64 @@
|
|||
import time
|
||||
from itertools import chain
|
||||
from shutil import copyfile
|
||||
from insightface.app import FaceAnalysis
|
||||
from os import getcwd, listdir, makedirs, supports_bytes_environ
|
||||
from os.path import join, isdir, isfile, basename, dirname
|
||||
from numpy import dot, array
|
||||
from numpy.linalg import norm
|
||||
from PIL import Image
|
||||
from sys import argv
|
||||
|
||||
if argv.__len__() != 2:
|
||||
exit(1)
|
||||
|
||||
if not isfile(argv[1]):
|
||||
exit(1)
|
||||
|
||||
face_analysis = FaceAnalysis(providers=[
|
||||
# 'CUDAExecutionProvider',
|
||||
# 'CPUExecutionProvider',
|
||||
('TensorrtExecutionProvider', {
|
||||
'trt_engine_cache_enable': True,
|
||||
'trt_engine_cache_path': join(getcwd(), 'onnx_cache'),
|
||||
'trt_fp16_enable': True,
|
||||
})
|
||||
], allowed_modules=['recognition', 'detection'])
|
||||
face_analysis.prepare(ctx_id=0, det_size=(160, 160))
|
||||
|
||||
print(argv)
|
||||
ROOT_DIR = r"D:\helloproject-ai-data\face_cropped"
|
||||
collect_image = array(Image.open(join(getcwd(), argv[1])))[:, :, [2, 1, 0]]
|
||||
image_files: list[str] = chain.from_iterable([[join() for file in listdir(join(ROOT_DIR, subdir))] for subdir in listdir(ROOT_DIR)])
|
||||
|
||||
collect_image_emb = face_analysis.get(collect_image)
|
||||
if collect_image_emb.__len__() == 0:
|
||||
print("Not found face: ", argv[1])
|
||||
exit(1)
|
||||
|
||||
# collect_image_emb = collect_image_emb[0].embedding
|
||||
|
||||
dir_name = basename(dirname(argv[2]))
|
||||
print(dir_name)
|
||||
makedirs(join(getcwd(), dir_name, "true"), exist_ok=True)
|
||||
makedirs(join(getcwd(), dir_name, "false"), exist_ok=True)
|
||||
|
||||
images = []
|
||||
begin = time.time()
|
||||
for file in image_files:
|
||||
if isfile(join(getcwd(), argv[2], file)):
|
||||
# print(join(getcwd(), argv[2], file))
|
||||
image = array(Image.open(join(getcwd(), argv[2], file)))[:, :, [2, 1, 0]]
|
||||
emb = face_analysis.get(image)
|
||||
if not emb:
|
||||
continue
|
||||
cosine = dot(emb[0].embedding, collect_image_emb[0].embedding) / \
|
||||
(norm(emb[0].embedding) * norm(collect_image_emb[0].embedding))
|
||||
print(file, cosine)
|
||||
if cosine > 0.3:
|
||||
copyfile(join(getcwd(), argv[2], file), join(getcwd(), dir_name, "true", file))
|
||||
else:
|
||||
pass
|
||||
# copyfile(join(getcwd(), argv[2], file), join(getcwd(), dir_name, "false", file))
|
||||
|
||||
print(f"{time.time() - begin}sec")
|
||||
|
|
@ -0,0 +1,115 @@
|
|||
# import cv2
|
||||
import msgspec
|
||||
from torch import tensor
|
||||
import torch
|
||||
from torchvision.transforms import functional, InterpolationMode
|
||||
from torchvision.io import decode_jpeg
|
||||
import os
|
||||
# import shutil
|
||||
import numpy
|
||||
from PIL import Image
|
||||
from io import BytesIO
|
||||
from more_itertools import chunked
|
||||
from tqdm import tqdm
|
||||
import math
|
||||
|
||||
ROOT_DIR = r"D:\helloproject-ai-data\blog_images"
|
||||
CROPPED_DIR = r"D:\helloproject-ai-data\face_cropped"
|
||||
CROP_THRESHOLD = 0.8
|
||||
|
||||
inference_size = 640
|
||||
device = torch.device("cuda")
|
||||
|
||||
|
||||
def calc_rotate(landmark: list[list[float]]) -> tuple[tuple[int, int], float]:
|
||||
left_eye, right_eye, nose, left_mouth, right_mouth = landmark
|
||||
center_x = sum((left_eye[0], right_eye[0], left_mouth[0], right_mouth[0])) / 4
|
||||
center_y = sum((left_eye[1], right_eye[1], left_mouth[1], right_mouth[1])) / 4
|
||||
eye_center = (right_eye[0] + left_eye[0]) / 2, (right_eye[1] + left_eye[1]) / 2
|
||||
mouth_center = (right_mouth[0] + left_mouth[0]) / 2, (right_mouth[1] + left_mouth[1]) / 2
|
||||
return (int(center_x), int(center_y)), numpy.arctan2(eye_center[0] - mouth_center[0],
|
||||
mouth_center[1] - eye_center[1])
|
||||
|
||||
|
||||
def cropper(pos_list: list[tuple[str, list[list] | None]], tqdm_pbar: tqdm, exist_set: set[str]):
|
||||
for name, pos_s in pos_list:
|
||||
file_name = name
|
||||
sub_dir_name = file_name.split("=", maxsplit=1)[0]
|
||||
if pos_s:
|
||||
decoded_image = None
|
||||
# host_image = decoded_image.cpu().numpy().transpose([1, 2, 0])
|
||||
for order, pos in enumerate(pos_s):
|
||||
dest_name = file_name.split(".")[0] + f"-{order}.jpg"
|
||||
if dest_name in exist_set:
|
||||
continue
|
||||
# print(dest_name)
|
||||
if decoded_image is None:
|
||||
dat = numpy.fromfile(os.path.join(ROOT_DIR, sub_dir_name, file_name),
|
||||
dtype=numpy.uint8)
|
||||
try:
|
||||
decoded_image = decode_jpeg(tensor(dat), device=device)
|
||||
except:
|
||||
decoded_image = tensor(
|
||||
numpy.array(Image.open(BytesIO(dat.tobytes()))).transpose([2, 0, 1])).to(device)
|
||||
bbox, acc, landmark = pos
|
||||
if acc > CROP_THRESHOLD:
|
||||
# scale = 1.0
|
||||
# if max(decoded_image.shape[1:]) > inference_size:
|
||||
scale = max(decoded_image.shape[1:]) / inference_size
|
||||
|
||||
bbox = list(map(lambda x: x * scale, bbox))
|
||||
landmark = list(map(lambda x: x * scale, landmark))
|
||||
|
||||
# print(file_name, decoded_image.shape[1:], scale, acc)
|
||||
center, rotate_angle = calc_rotate(list(chunked(landmark, n=2)))
|
||||
# print(bbox, landmark, center)
|
||||
rotated = functional.rotate(decoded_image, angle=(360 / (2 * math.pi)) * rotate_angle,
|
||||
center=list(center), interpolation=InterpolationMode.BILINEAR)
|
||||
crop_size = max([int(bbox[3] - bbox[1]), int(bbox[2] - bbox[0])])
|
||||
if crop_size < 100:
|
||||
continue
|
||||
cropped = functional.crop(rotated,
|
||||
top=int(center[1] - crop_size / 2),
|
||||
left=int(center[0] - crop_size / 2),
|
||||
height=crop_size, width=crop_size)
|
||||
with open(os.path.join(CROPPED_DIR, sub_dir_name, dest_name), mode="wb") as fp:
|
||||
functional.to_pil_image(cropped, mode="RGB").save(fp, format="jpeg", quality=85)
|
||||
if tqdm_pbar.desc != sub_dir_name:
|
||||
tqdm_pbar.set_description(sub_dir_name)
|
||||
tqdm_pbar.update(n=1)
|
||||
|
||||
|
||||
# cv2.rectangle(host_image, (int(bbox[0]), int(bbox[1])),
|
||||
# (int(bbox[2]), int(bbox[3])),
|
||||
# (255, 0, 0), 2, cv2.LINE_AA)
|
||||
# with open(
|
||||
# os.path.join(CROPPED_DIR, sub_dir_name, file_name),
|
||||
# mode="wb") as fp:
|
||||
# cv2.imencode(".jpg", cv2.cvtColor(host_image, cv2.COLOR_BGR2RGB))[1].tofile(fp)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
with open(file="faces.jsonl", mode="r", encoding="utf-8") as fp:
|
||||
face_pos_list: list[dict[str, list | None]] = [msgspec.json.decode(line) for line in
|
||||
fp.read().removesuffix("\n").split("\n")]
|
||||
|
||||
face_pos_dict = {}
|
||||
for face_pos in face_pos_list:
|
||||
if next(iter(face_pos.values())) is None:
|
||||
face_pos_dict[next(iter(face_pos.keys()))] = None
|
||||
else:
|
||||
if not next(iter(face_pos.keys())) in face_pos_dict.keys():
|
||||
face_pos_dict[next(iter(face_pos.keys()))] = []
|
||||
face_pos_dict[next(iter(face_pos.keys()))].append(next(iter(face_pos.values())))
|
||||
|
||||
# shutil.rmtree(CROPPED_DIR)
|
||||
os.makedirs(CROPPED_DIR, exist_ok=True)
|
||||
# print(face_pos)
|
||||
names_set = {next(iter(_dict.keys())).split("=", maxsplit=1)[0] for _dict in face_pos_list}
|
||||
print(names_set)
|
||||
[os.makedirs(os.path.join(CROPPED_DIR, name), exist_ok=True) for name in names_set]
|
||||
pbar = tqdm(total=face_pos_dict.items().__len__())
|
||||
|
||||
exist_set = set().union(
|
||||
*[set(os.listdir(os.path.join(CROPPED_DIR, sub_dir))) for sub_dir in os.listdir(CROPPED_DIR)])
|
||||
cropper(list(face_pos_dict.items()), pbar, exist_set)
|
||||
Binary file not shown.
|
|
@ -5,7 +5,7 @@ import tkinter
|
|||
from nvjpeg_decoder import decode
|
||||
from os import listdir, getcwd
|
||||
from os.path import join
|
||||
from numpy import array
|
||||
from numpy import array, fromfile, uint8
|
||||
# from matplotlib import pyplot, figure
|
||||
# from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg
|
||||
# import matplotlib_fontja
|
||||
|
|
@ -24,8 +24,9 @@ session = InferenceSession(
|
|||
)
|
||||
for member in listdir(datadir):
|
||||
for file in listdir(join(datadir, member)):
|
||||
with open(join(datadir, member, file), mode="rb") as f:
|
||||
(data, (width, height)) = decode((f.read()), "imagenet")
|
||||
print(width, height)
|
||||
image_arr = array(data).reshape((1, 3, height, width)) # .transpose([1, 2, 0])
|
||||
session.run(input_feed={'input': image_arr}, output_names=['bbox', 'confidence', 'landmark'])
|
||||
# with open(join(datadir, member, file), mode="rb") as f:
|
||||
(data, (scale, (width, height))) = decode(fromfile(join(datadir, member, file), dtype=uint8), "imagenet",
|
||||
(1080, 1080))
|
||||
print(width, height)
|
||||
image_arr = array(data).reshape((1, 3, height, width)) # .transpose([1, 2, 0])
|
||||
session.run(input_feed={'input': image_arr}, output_names=['bbox', 'confidence', 'landmark'])
|
||||
|
|
|
|||
|
|
@ -1,14 +1,17 @@
|
|||
import json
|
||||
import os
|
||||
from concurrent.futures.process import ProcessPoolExecutor
|
||||
from itertools import chain
|
||||
from multiprocessing import shared_memory
|
||||
from io import BytesIO
|
||||
from os import listdir, path, pathsep, makedirs
|
||||
from pprint import pprint
|
||||
import more_itertools
|
||||
import numpy as np
|
||||
import msgspec
|
||||
import pandas.io.json
|
||||
import tqdm
|
||||
from PIL import Image
|
||||
from numpy import ndarray
|
||||
from uuid import uuid4
|
||||
from onnxruntime import InferenceSession
|
||||
from torch import tensor
|
||||
import aiofiles
|
||||
|
|
@ -42,21 +45,18 @@ async def gather_runner(l: list, fn):
|
|||
return await gather(*[fn(p, sem) for p in l])
|
||||
|
||||
|
||||
def post_processor(outputs, batch_size, image_size):
|
||||
# print("aaa", flush=True)
|
||||
outputs = [numpy.ascontiguousarray(output.astype(numpy.float32)) for output in outputs]
|
||||
res = resnet_post_process([output.__array_interface__["data"][0] for output in outputs], batch_size, image_size)
|
||||
return res
|
||||
|
||||
|
||||
def post_processor_memmap(tmp_filename, sizes, batch_size, image_size):
|
||||
# print("aaa", flush=True)
|
||||
outputs = [
|
||||
numpy.memmap(filename=path.join("memmap", tmp_filename + str(order)), dtype=numpy.float16, mode="r", shape=size)
|
||||
for order, size in enumerate(sizes)]
|
||||
outputs = [numpy.ascontiguousarray(output.astype(numpy.float32)) for output in outputs]
|
||||
res = resnet_post_process([output.__array_interface__["data"][0] for output in outputs], batch_size, image_size)
|
||||
return res
|
||||
# def post_processor(outputs, batch_size, image_size):
|
||||
# # print("aaa", flush=True)
|
||||
# outputs = [numpy.ascontiguousarray(output.astype(numpy.float32)) for output in outputs]
|
||||
# res = resnet_post_process([output.__array_interface__["data"][0] for output in outputs], batch_size, image_size)
|
||||
# return res
|
||||
#
|
||||
#
|
||||
# def post_processor_memmap(tmp_filename, sizes, batch_size, image_size): # print("aaa", flush=True) outputs = [
|
||||
# numpy.memmap(filename=path.join("memmap", tmp_filename + str(order)), dtype=numpy.float16, mode="r", shape=size)
|
||||
# for order, size in enumerate(sizes)] outputs = [numpy.ascontiguousarray(output.astype(numpy.float32)) for output in
|
||||
# outputs] res = resnet_post_process([output.__array_interface__["data"][0] for output in outputs], batch_size,
|
||||
# image_size) return res
|
||||
|
||||
|
||||
def post_processor_shm(shm_name, sizes, batch_size, image_size):
|
||||
|
|
@ -65,7 +65,7 @@ def post_processor_shm(shm_name, sizes, batch_size, image_size):
|
|||
[numpy.ascontiguousarray(numpy.ndarray(shape=size, dtype=numpy.float16, buffer=shm.buf).astype(numpy.float32))
|
||||
for size, shm in zip(sizes, shms)]
|
||||
res = resnet_post_process([output.__array_interface__["data"][0] for output in outputs], batch_size, image_size)
|
||||
[shm.close() for shm in shms]
|
||||
# print(res)
|
||||
return res
|
||||
|
||||
|
||||
|
|
@ -73,7 +73,7 @@ if __name__ == '__main__':
|
|||
from kornia.augmentation import LongestMaxSize, PadTo, Normalize
|
||||
from kornia.constants import Resample
|
||||
|
||||
longest_max_size = LongestMaxSize(max_size=640,resample=Resample.NEAREST)
|
||||
longest_max_size = LongestMaxSize(max_size=640, resample=Resample.NEAREST)
|
||||
pad_to = PadTo(size=(640, 640), pad_value=1.)
|
||||
normalize = Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225))
|
||||
|
||||
|
|
@ -92,35 +92,51 @@ if __name__ == '__main__':
|
|||
'CPUExecutionProvider'
|
||||
]
|
||||
)
|
||||
with open(file="faces.jsonl", mode="r", encoding="utf-8") as fp:
|
||||
already = {list(msgspec.json.decode(line).keys())[0] for line in fp.read().removesuffix("\n").split("\n")}
|
||||
pbar = tqdm.tqdm(
|
||||
total=(set().union(*[listdir(path.join(root_dir, name)) for name in listdir(root_dir)]) - already).__len__())
|
||||
|
||||
# print(len(already))
|
||||
# exit(0)
|
||||
|
||||
for name in listdir(root_dir):
|
||||
with (ProcessPoolExecutor(max_workers=4) as executor):
|
||||
if name != "下井谷幸穂":
|
||||
pbar.set_description_str(desc=name, refresh=True)
|
||||
if name != "ブログ":
|
||||
# continue
|
||||
pass
|
||||
file_names = listdir(path.join(root_dir, name))
|
||||
file_names_set = set(file_names) - already
|
||||
file_names = list(file_names_set)
|
||||
name_files = [path.join(root_dir, name, file_name) for file_name in file_names]
|
||||
files_data = {file_name: numpy.frombuffer(dat, dtype=numpy.uint8) for file_name, dat in
|
||||
zip(file_names, run(gather_runner(name_files, async_read)))}
|
||||
|
||||
if files_data.__len__() == 0:
|
||||
continue
|
||||
futures = []
|
||||
shms = []
|
||||
namess = []
|
||||
# print(k_1)
|
||||
for cnk in more_itertools.chunked(tqdm.tqdm(files_data.items(), desc=name), n=chunk_size):
|
||||
for cnk in more_itertools.chunked(files_data.items(), n=chunk_size):
|
||||
stack = []
|
||||
tmp_file_name = ""
|
||||
names = []
|
||||
for file, dat in cnk:
|
||||
tmp_file_name = file
|
||||
try:
|
||||
decoded_image = decode_jpeg(tensor(dat), device=device)
|
||||
except:
|
||||
decoded_image = tensor(numpy.array(Image.open(BytesIO(dat.tobytes()))).transpose([2, 0, 1])).to(
|
||||
decoded_image = tensor(
|
||||
numpy.array(Image.open(BytesIO(dat.tobytes()))).transpose([2, 0, 1])).to(
|
||||
device)
|
||||
decoded_image = decoded_image.to(torch.float16) / 255
|
||||
decoded_image = normalize(decoded_image)
|
||||
decoded_image_resized = longest_max_size(decoded_image)
|
||||
decoded_image_padded = pad_to(decoded_image_resized)
|
||||
stack.append(decoded_image_padded.squeeze())
|
||||
names.append(file)
|
||||
namess.append(names)
|
||||
[stack.append(torch.zeros(size=[3, 640, 640], dtype=torch.float16, device=device)) for _ in
|
||||
range(32 - stack.__len__())]
|
||||
range(chunk_size - stack.__len__())]
|
||||
stacked = torch.stack(stack).contiguous()
|
||||
# print(stacked.shape)
|
||||
io_binding = session.io_binding()
|
||||
|
|
@ -139,15 +155,46 @@ if __name__ == '__main__':
|
|||
outputs: list[numpy.ndarray] = io_binding.copy_outputs_to_cpu()
|
||||
# [numpy.memmap(filename=path.join("memmap", tmp_file_name + str(order)), dtype=numpy.float16,
|
||||
# mode="w+", shape=output.shape) for order, output in enumerate(outputs)]
|
||||
uuid = uuid4().__str__()
|
||||
shared_array: list[shared_memory.SharedMemory] = \
|
||||
[shared_memory.SharedMemory(name=tmp_file_name + "_" + str(order), create=True, size=output.nbytes)
|
||||
[shared_memory.SharedMemory(name=uuid + "_" + str(order), create=True, size=output.nbytes)
|
||||
for order, output in enumerate(outputs)]
|
||||
shared_ndarray = [numpy.ndarray(shape=output.shape, dtype=numpy.float16, buffer=shm.buf)
|
||||
for shm, output in zip(shared_array, outputs, strict=True)]
|
||||
for shm, output in zip(shared_ndarray, outputs, strict=True):
|
||||
shm[:] = output[:]
|
||||
future = executor.submit(post_processor_shm, tmp_file_name, [output.shape for output in outputs],
|
||||
future = executor.submit(post_processor_shm, uuid, [output.shape for output in outputs],
|
||||
chunk_size, [image_size, image_size])
|
||||
# print(future.result())
|
||||
# future.add_done_callback(pprint)
|
||||
futures.append(future)
|
||||
shms.extend(shared_array)
|
||||
# exit(0)
|
||||
pbar.update(n=cnk.__len__())
|
||||
# result_dict = dict()
|
||||
with open("faces.jsonl", mode="a", encoding="utf-8") as fp:
|
||||
futures_results = [future.result() for future in futures]
|
||||
# pprint(futures_results)
|
||||
for names, futures_result in zip(namess, futures_results):
|
||||
|
||||
for name, results in zip(names, futures_result):
|
||||
results_list = []
|
||||
if results:
|
||||
# print(name)
|
||||
for result in results:
|
||||
# [print(int(a), end=" ") for a in result[0]]
|
||||
# print(*result[1], end=" ")
|
||||
# [print(int(a), end=" ") for a in result[2]]
|
||||
# print()
|
||||
# results_list.append(list(chain.from_iterable([result])))
|
||||
fp.write(
|
||||
pandas.io.json.ujson_dumps({name: [result[0], result[1][0], result[2]]},
|
||||
ensure_ascii=False, double_precision=5) + "\n")
|
||||
pass
|
||||
else:
|
||||
fp.write(
|
||||
pandas.io.json.ujson_dumps({name: None}, ensure_ascii=False) + "\n")
|
||||
|
||||
# print(name, [])
|
||||
pass
|
||||
# result_dict[name] = results_list
|
||||
# pprint(result_dict)
|
||||
[shm.close() for shm in shms]
|
||||
|
|
|
|||
Loading…
Reference in New Issue