Upload 8 files

app.py

@@ -0,0 +1,20 @@
+import gradio as gr
+from src.skeletal import FaceMesh
+
+facemesh = FaceMesh("config.yaml")
+
+with gr.Blocks() as demo:
+    gr.Markdown("Estimate smilar person using this demo.")
+    with gr.Row():
+        with gr.Column(scale=1):
+            input = gr.Image(type="filepath", label="Input image")
+            dropdown = gr.Dropdown([5, 10, 20, 30, 40, 50], value="20", label="Top K")
+            button = gr.Button("Estimate")
+        with gr.Column(scale=2):
+            output = gr.Dataframe()
+
+    button.click(
+        facemesh.estimate_similar_person, inputs=[input, dropdown], outputs=output
+    )
+
+demo.launch()

config.yaml

@@ -0,0 +1,4 @@
+path_data: data
+path_garbage: garbage
+path_csv: result.csv
+path_skeletal: skeletal.npy

requirements.txt

@@ -0,0 +1,8 @@
+notebook==6.5.3
+requests==2.30.0
+bs4==0.0.1
+pandas==2.0.1
+omegaconf==2.3.0
+opencv-python==4.7.0
+matplotlib==3.7.1
+gradio==3.28.3

skeletal.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:18d7a3a30feb38ed39755e90b1cb6b2f1a578d7ad80c8f4be7a1bae009958085
+size 13835360

src/scrape.py

@@ -0,0 +1,126 @@
+import argparse
+import glob
+import os
+import shutil
+import time
+import urllib
+
+import cv2
+import mediapipe as mp
+import pandas as pd
+import requests
+from bs4 import BeautifulSoup
+from omegaconf import OmegaConf
+
+
+class Scraper:
+    def __init__(self, config):
+        self.config = OmegaConf.load(config)
+        self.base_url = "https://hominis.media/person/"
+        if os.path.exists(self.config.path_csv):
+            self.df = pd.read_csv(self.config.path_csv)
+            self.idx = len(self.df)
+        else:
+            self.df = pd.DataFrame([], columns=["filepath", "name", "url"])
+            self.idx = 0
+        os.makedirs(self.config.path_data, exist_ok=True)
+        os.makedirs(self.config.path_garbage, exist_ok=True)
+
+    def run(self):
+        html = requests.get(self.base_url, timeout=5)
+        soup = BeautifulSoup(html.content, "html.parser")
+        pages = soup.find_all("input", class_="selectButton")
+        before = 0
+
+        for page in pages:
+            url = self.base_url + page.get("onclick").split("'")[1].replace(
+                "/person/", ""
+            )
+            html = requests.get(url, timeout=5)
+            soup = BeautifulSoup(html.content, "html.parser")
+            people = soup.find_all("li", class_="card people")
+            for person in people:
+                name = person.find("p", class_="name").text
+                img_url = (
+                    person.find("p", class_="thumbnail")
+                    .get("style")
+                    .replace("background-image:url('", "")
+                    .replace("');", "")
+                )
+                img_path = os.path.join(self.config.path_data, name + ".png")
+                if os.path.exists(img_path):
+                    continue
+                try:
+                    urllib.request.urlretrieve(img_url, img_path)
+                    self.df.loc[self.idx] = {
+                        "filepath": img_path,
+                        "name": name,
+                        "url": img_url,
+                    }
+                    self.idx += 1
+                    time.sleep(1)
+                except Exception:
+                    continue
+
+            imgs = glob.glob(os.path.join(self.config.path_data, "*.png"))
+            assert len(imgs) == len(self.df)
+            print(f"Get {len(imgs) - before} images")
+            before = len(imgs)
+
+        self.df.to_csv(self.config.path_csv, index=False)
+
+    def post_processing(self):
+        mp_face_mesh = mp.solutions.face_mesh
+        with mp_face_mesh.FaceMesh(
+            static_image_mode=True,
+            max_num_faces=10,
+            refine_landmarks=True,
+            min_detection_confidence=0.5,
+        ) as face_mesh:
+            for file in glob.glob(os.path.join(self.config.path_data, "*.png")):
+                image = cv2.imread(file)
+                results = face_mesh.process(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
+                if not results.multi_face_landmarks:
+                    shutil.move(
+                        file,
+                        os.path.join(self.config.path_garbage, os.path.split(file)[-1]),
+                    )
+                if len(results.multi_face_landmarks) > 1:
+                    shutil.move(
+                        file,
+                        os.path.join(self.config.path_garbage, os.path.split(file)[-1]),
+                    )
+
+        idx = []
+        for path in glob.glob(os.path.join(self.config.path_garbage, "*.png")):
+            idx.append(
+                self.df[
+                    self.df["filepath"]
+                    == os.path.join(self.config.path_data, os.path.split(path)[-1])
+                ].index.values[0]
+            )
+        self.df = self.df.drop(idx)
+        assert len(glob.glob(os.path.join(self.config.path_data, "*.png"))) == len(
+            self.df
+        )
+        self.df.to_csv(self.config.path_csv, index=False)
+
+
+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 = Scraper(args.config)
+    scraper.run()
+    scraper.post_processing()

src/skeletal.py

@@ -0,0 +1,79 @@
+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()