From 7ab7b6684cb7fcdd6be3fd9a56624e1bc47d1047 Mon Sep 17 00:00:00 2001 From: matt3o Date: Thu, 22 Feb 2024 15:39:39 +0100 Subject: [PATCH] supports batch of reference images --- faceanalysis.py | 16 +++++++++++++--- 1 file changed, 13 insertions(+), 3 deletions(-) diff --git a/faceanalysis.py b/faceanalysis.py index 854766a..e825b38 100644 --- a/faceanalysis.py +++ b/faceanalysis.py @@ -69,11 +69,21 @@ class FaceEmbedDistance: self.analysis_models = analysis_models - ref = np.array(T.ToPILImage()(reference[0].permute(2, 0, 1)).convert('RGB')) - ref = self.get_descriptor(ref) - if ref is None: + #if reference.shape[0] > 1: + # reference = torch.mean(reference, dim=0).unsqueeze(0) + + ref = [] + for i in reference: + ref_emb = self.get_descriptor(np.array(T.ToPILImage()(i.permute(2, 0, 1)).convert('RGB'))) + if ref_emb is not None: + ref.append(torch.from_numpy(ref_emb)) + + if ref == []: raise Exception('No face detected in reference image') + ref = torch.stack(ref) + ref = np.array(torch.mean(ref, dim=0)) + out = [] out_eucl = [] out_cos = []