Spaces:
Runtime error
Runtime error
File size: 2,336 Bytes
a1e77ee |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 |
import argparse
import os
import cv2
import mediapipe as mp
import numpy as np
import pandas as pd
from omegaconf import OmegaConf
class FaceMesh:
def __init__(self, config):
self.config = OmegaConf.load(config)
self.df = pd.read_csv(self.config.path_csv)
if os.path.exists(self.config.path_skeletal):
self.reference = np.load(self.config.path_skeletal)
def normalize(self, landmarks):
output = []
for landmark in landmarks:
landmark = np.array(landmark)
landmark = (landmark - landmark.min()) / (landmark.max() - landmark.min())
output.append(landmark)
return np.array(output)
def get_facemesh(self, path):
mp_face_mesh = mp.solutions.face_mesh
with mp_face_mesh.FaceMesh(
static_image_mode=True,
max_num_faces=1,
refine_landmarks=True,
min_detection_confidence=0.5,
) as face_mesh:
results = face_mesh.process(
cv2.cvtColor(cv2.imread(path), cv2.COLOR_BGR2RGB)
)
x, y, z = [], [], []
result = results.multi_face_landmarks[0]
for lands in result.landmark:
x.append(lands.x)
y.append(lands.y)
z.append(lands.z)
landmark = self.normalize([x, y, z])
return landmark
def create_dataset(self):
landmarks = []
for i in range(len(self.df)):
landmark = self.get_facemesh(self.df.iloc[i]["filepath"])
landmarks.append(landmark)
np.save(self.config.path_skeletal, np.array(landmarks))
def estimate_similar_person(self, path, topK):
print(path)
facemesh = self.get_facemesh(path)
diff = abs(self.reference - facemesh).mean((1, 2))
rank = np.argsort(diff)[0 : int(topK)]
top = self.df.iloc[rank]
return top.drop("filepath", axis=1)
def argparser():
parser = argparse.ArgumentParser()
parser.add_argument(
"-c",
"--config",
type=str,
default="config.yaml",
help="File path for config file.",
)
args = parser.parse_args()
return args
if __name__ == "__main__":
args = argparser()
scraper = FaceMesh(args.config)
scraper.create_dataset()
|