Spaces:
Build error
Build error
| import os | |
| import sys | |
| from importlib import import_module, invalidate_caches | |
| from importlib.util import module_from_spec, spec_from_file_location | |
| from tempfile import TemporaryDirectory | |
| import cv2 | |
| import mediapipe as mp | |
| import numpy as np | |
| import plotly.express as px | |
| import requests | |
| import torch | |
| from git import Repo | |
| from huggingface_hub import hf_hub_download | |
| class FECNetModel: | |
| def __init__(self, hf_token: str) -> None: | |
| self.hf_token = hf_token | |
| repo_dir = TemporaryDirectory() | |
| Repo.clone_from( | |
| "https://github.com/AmirSh15/FECNet.git", | |
| repo_dir.name, | |
| ) | |
| invalidate_caches() | |
| sys.path.append(repo_dir.name) | |
| fecnet_module_path = os.path.join(repo_dir.name, "models", "FECNet.py") | |
| with open(fecnet_module_path, "r") as f: | |
| content = f.read() | |
| content = content.replace( | |
| "cuda", | |
| "cpu", | |
| ) | |
| with open(fecnet_module_path, "w") as f: | |
| f.write(content) | |
| spec = spec_from_file_location("FECNet", fecnet_module_path) | |
| fecnet_module = module_from_spec(spec) # type: ignore | |
| spec.loader.exec_module(fecnet_module) # type: ignore | |
| self.model = self.__load_model( | |
| self.__download_weights(repo_dir.name), fecnet_module.FECNet | |
| ) | |
| self.face_detector = mp.solutions.face_detection.FaceDetection( | |
| min_detection_confidence=0.5 | |
| ) | |
| def __download_weights(self, model_dir: str) -> str: | |
| model_path = hf_hub_download( | |
| "natexcvi/pretrained-fecnet", | |
| "fecnet.pt", | |
| token=self.hf_token, | |
| ) | |
| return model_path | |
| def __load_model(self, model_path: str, model_class): | |
| model = model_class(pretrained=False) | |
| model_weights = torch.load(model_path, map_location=torch.device("cpu")) | |
| model.load_state_dict(model_weights) | |
| model.eval() | |
| return model | |
| def predict(self, image: np.ndarray): | |
| pred = self.model(image) | |
| return pred | |
| def distance(a, b): | |
| return np.linalg.norm(a - b) | |
| def embed_image(self, image, crop_face: bool = False) -> np.ndarray: | |
| image = cv2.imdecode(image, cv2.IMREAD_COLOR) | |
| if crop_face: | |
| image = self.extract_face(image) | |
| # image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) | |
| image = cv2.resize(image, (224, 224)) | |
| image = np.transpose(image, (2, 0, 1)) | |
| image = np.expand_dims(image, axis=0) | |
| image = torch.from_numpy(image.astype(np.float32)) | |
| pred = self.predict(image) | |
| return pred.detach().numpy() | |
| def extract_face(self, image): | |
| mp_face_detection = mp.solutions.face_detection | |
| # Convert the image to RGB | |
| image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) | |
| # Initialize the face detection model | |
| # Run the face detection model on the image | |
| results = self.face_detection.process(image) | |
| # If a face is detected, crop the image to the face box | |
| if results.detections: | |
| for detection in results.detections: | |
| x, y, w, h = ( | |
| int( | |
| detection.location_data.relative_bounding_box.xmin | |
| * image.shape[1] | |
| ), | |
| int( | |
| detection.location_data.relative_bounding_box.ymin | |
| * image.shape[0] | |
| ), | |
| int( | |
| detection.location_data.relative_bounding_box.width | |
| * image.shape[1] | |
| ), | |
| int( | |
| detection.location_data.relative_bounding_box.height | |
| * image.shape[0] | |
| ), | |
| ) | |
| cropped_image = image[y : y + h, x : x + w] | |
| return cv2.cvtColor(cropped_image, cv2.COLOR_RGB2BGR) | |