diff --git a/embeddings_search.py b/embeddings_search.py new file mode 100644 index 0000000..b1d6f37 --- /dev/null +++ b/embeddings_search.py @@ -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).") diff --git a/face_annotator.py b/face_annotator.py new file mode 100644 index 0000000..8746980 --- /dev/null +++ b/face_annotator.py @@ -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))) diff --git a/generate_embeddings.py b/generate_embeddings.py index 42149ec..edbd448 100644 --- a/generate_embeddings.py +++ b/generate_embeddings.py @@ -9,17 +9,24 @@ from torchvision import transforms from torchvision.io import decode_jpeg from tqdm import tqdm from facenet_pytorch import InceptionResnetV1 -from edgeface import get_model import torch import numpy from insightface.app import FaceAnalysis CROPPED_DIR = r"D:\helloproject-ai-data\face_cropped" -MODEL_NAME = "edgeface_s_gamma_05" CHUNK_SIZE = 64 DEVICE = torch.device("cuda") INPUT_SIZE = 256 -face_analysis = FaceAnalysis() +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([ @@ -32,53 +39,45 @@ transform = transforms.Compose([ # 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) +# 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]) + # 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[:10000] + # all_cropped_list = all_cropped_list[:1000] + labels_set = set(labels) pbar = tqdm(total=all_cropped_list.__len__()) - for name_chunk in chunked(all_cropped_list, n=CHUNK_SIZE): - decoded_images = [] - for file_name in name_chunk: - sub_dir_name = file_name.split("=")[0] - dat = numpy.fromfile(os.path.join(CROPPED_DIR, sub_dir_name, file_name), dtype=numpy.uint8) - try: - decoded_image = decode_jpeg(torch.tensor(dat), device=DEVICE) - except BaseException as e: - decoded_image = (torch.tensor(numpy.array(Image.open(BytesIO(dat.tobytes()))) - .transpose([2, 0, 1])).to(DEVICE)) - decoded_images.append(decoded_image) # transform(decoded_image.to(torch.float32) / 255.)) - if pbar.desc != sub_dir_name: - pbar.set_description(sub_dir_name) - pbar.update(1) - # input_tensor = torch.stack(decoded_images) - # res: torch.Tensor = model(input_tensor) - # print(res.shape) - # print(face_analysis.get(decoded_images[0].cpu().numpy().transpose([1, 2, 0])[:, :, ::-1])) - _label = [] - res = [] - for decoded_image, name in zip(decoded_images, name_chunk): - if a := face_analysis.get(decoded_image.cpu().numpy().transpose([1, 2, 0])[:, :, ::-1]): - _label.append(name) - res.append(a[0].embedding) - # print(a[0].embedding.shape) - res = numpy.stack(res) + 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, res], axis=0) + embeddings = numpy.concatenate([embeddings, numpy.expand_dims(emb[0].embedding, 0)], axis=0) + # print(embeddings.shape) else: - embeddings = res - labels.extend(_label) + embeddings = numpy.expand_dims(emb[0].embedding, 0) + labels.append(file_name) + # print(embeddings.shape, labels.__len__()) - numpy.save("embeddings.npy", embeddings) - numpy.save("embeddings_label.npy", numpy.array(all_cropped_list)) + numpy.save(f"embeddings_{TYPE}.npy", embeddings) + numpy.save(f"embeddings_{TYPE}_label.npy", numpy.array(labels)) diff --git a/generate_embeddings_edgeface.py b/generate_embeddings_edgeface.py new file mode 100644 index 0000000..26b3ef5 --- /dev/null +++ b/generate_embeddings_edgeface.py @@ -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)) diff --git a/similar_face.py b/similar_face.py index 8a14477..343b61c 100755 --- a/similar_face.py +++ b/similar_face.py @@ -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: diff --git a/similar_face_all.py b/similar_face_all.py new file mode 100644 index 0000000..2606e42 --- /dev/null +++ b/similar_face_all.py @@ -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") diff --git a/test_script/face_cropper.py b/test_script/face_cropper.py new file mode 100644 index 0000000..58f9b5c --- /dev/null +++ b/test_script/face_cropper.py @@ -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) diff --git a/test_script/nvjpeg_decoder.pyd b/test_script/nvjpeg_decoder.pyd new file mode 100644 index 0000000..5724bc5 Binary files /dev/null and b/test_script/nvjpeg_decoder.pyd differ diff --git a/test_script/nvjpeg_infer.py b/test_script/nvjpeg_infer.py index 608f6a7..710bf10 100644 --- a/test_script/nvjpeg_infer.py +++ b/test_script/nvjpeg_infer.py @@ -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']) diff --git a/test_script/torchivision_jpeg_decode.py b/test_script/torchivision_jpeg_decode.py index d6661b9..b922361 100644 --- a/test_script/torchivision_jpeg_decode.py +++ b/test_script/torchivision_jpeg_decode.py @@ -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]