53 lines
1.6 KiB
Python
Executable File
53 lines
1.6 KiB
Python
Executable File
from shutil import copyfile
|
|
from insightface.app import FaceAnalysis
|
|
from os import getcwd, listdir, makedirs
|
|
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__() != 3:
|
|
exit(1)
|
|
if not isdir(argv[2]):
|
|
exit(1)
|
|
|
|
if not isfile(argv[1]):
|
|
exit(1)
|
|
|
|
face_analysis = FaceAnalysis()
|
|
face_analysis.prepare(ctx_id=0, det_size=(160, 160))
|
|
|
|
print(argv)
|
|
collect_image = array(Image.open(join(getcwd(), argv[1])))[:, :, [2, 1, 0]]
|
|
image_files: list[str] = listdir(join(getcwd(), argv[2]))
|
|
|
|
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 = []
|
|
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:
|
|
|
|
copyfile(join(getcwd(), argv[2], file), join(getcwd(), dir_name, "false", file))
|