| |
|
| |
|
| |
|
| | import torch
|
| | import torchvision.transforms as T
|
| |
|
| | import numpy as np
|
| |
|
| | from insightface.utils import face_align
|
| | from insightface.app import FaceAnalysis
|
| | from facexlib.recognition import init_recognition_model
|
| |
|
| |
|
| | __all__ = [
|
| | "FaceEncoderArcFace",
|
| | "get_landmarks_from_image",
|
| | ]
|
| |
|
| |
|
| | detector = None
|
| |
|
| | def get_landmarks_from_image(image):
|
| | """
|
| | Detect landmarks with insightface.
|
| |
|
| | Args:
|
| | image (np.ndarray or PIL.Image):
|
| | The input image in RGB format.
|
| |
|
| | Returns:
|
| | 5 2D keypoints, only one face will be returned.
|
| | """
|
| | global detector
|
| | if detector is None:
|
| | detector = FaceAnalysis()
|
| | detector.prepare(ctx_id=0, det_size=(640, 640))
|
| |
|
| | in_image = np.array(image).copy()
|
| |
|
| | faces = detector.get(in_image)
|
| | if len(faces) == 0:
|
| | raise ValueError("No face detected in the image")
|
| |
|
| |
|
| | face = max(faces, key=lambda x: (x.bbox[2] - x.bbox[0]) * (x.bbox[3] - x.bbox[1]))
|
| |
|
| |
|
| | keypoints = face.kps
|
| |
|
| | return keypoints
|
| |
|
| |
|
| | class FaceEncoderArcFace():
|
| | """ Official ArcFace, no_grad-only """
|
| |
|
| | def __repr__(self):
|
| | return "ArcFace"
|
| |
|
| |
|
| | def init_encoder_model(self, device, eval_mode=True):
|
| | self.device = device
|
| | self.encoder_model = init_recognition_model('arcface', device=device)
|
| |
|
| | if eval_mode:
|
| | self.encoder_model.eval()
|
| |
|
| |
|
| | @torch.no_grad()
|
| | def input_preprocessing(self, in_image, landmarks, image_size=112):
|
| | assert landmarks is not None, "landmarks are not provided!"
|
| |
|
| | in_image = np.array(in_image)
|
| | landmark = np.array(landmarks)
|
| |
|
| | face_aligned = face_align.norm_crop(in_image, landmark=landmark, image_size=image_size)
|
| |
|
| | image_transform = T.Compose([
|
| | T.ToTensor(),
|
| | T.Normalize([0.5], [0.5]),
|
| | ])
|
| | face_aligned = image_transform(face_aligned).unsqueeze(0).to(self.device)
|
| |
|
| | return face_aligned
|
| |
|
| |
|
| | @torch.no_grad()
|
| | def __call__(self, in_image, need_proc=False, landmarks=None, image_size=112):
|
| |
|
| | if need_proc:
|
| | in_image = self.input_preprocessing(in_image, landmarks, image_size)
|
| | else:
|
| | assert isinstance(in_image, torch.Tensor)
|
| |
|
| | in_image = in_image[:, [2, 1, 0], :, :].contiguous()
|
| | image_embeds = self.encoder_model(in_image)
|
| |
|
| | return image_embeds |
