facenet-pytorch

facenet-pytorch はFaceNetをPyTorchで実装したものである。pip install facenet-pytorch でインストールする。

from facenet_pytorch import MTCNN, InceptionResnetV1
from PIL import Image

mtcnn = MTCNN() # オプションがいろいろあるがとりあえず
resnet = InceptionResnetV1(pretrained='vggface2').eval()

img = Image.open("face.jpg")
img_cropped = mtcnn(img)
img_embedding = resnet(img_cropped.unsqueeze(0)) # shape: (1,512)

ディレクトリに入っている画像の類似度の行列を表示する：

from torch.nn.functional import cosine_similarity
import pathlib

path = pathlib.Path("/path/to/folder")
names = [p for p in path.iterdir() if p.match("*.jp*g")]  

def encoding(name):
    img = Image.open(name)
    img_cropped = mtcnn(img)
    return resnet(img_cropped.unsqueeze(0)).detach()

for i in range(1, len(names)):
    for j in range(i):
        ei = encoding(names[i])
        ej = encoding(names[j])
        similarity = cosine_similarity(ei, ej).item()
        if similarity > 0.6:
            print(f" {similarity:4.2f}*", end="")
        else:
            print(f" {similarity:4.2f} ", end="")
    print()