61 lines
1.7 KiB
Python
61 lines
1.7 KiB
Python
import numpy as np
|
|
import cv2
|
|
import threading
|
|
from typing import Any, Optional, List
|
|
import insightface
|
|
from insightface.app.common import Face
|
|
try:
|
|
import torch.cuda as cuda
|
|
except:
|
|
cuda = None
|
|
if cuda is not None:
|
|
if cuda.is_available():
|
|
providers = ["CUDAExecutionProvider"]
|
|
else:
|
|
providers = ["CPUExecutionProvider"]
|
|
else:
|
|
providers = ["CPUExecutionProvider"]
|
|
|
|
THREAD_LOCK = threading.Lock()
|
|
FACE_ANALYSER = None
|
|
FS_MODEL = None
|
|
CURRENT_FS_MODEL_PATH = None
|
|
|
|
def getFaceSwapModel(model_path: str):
|
|
global FS_MODEL
|
|
global CURRENT_FS_MODEL_PATH
|
|
if CURRENT_FS_MODEL_PATH is None or CURRENT_FS_MODEL_PATH != model_path:
|
|
CURRENT_FS_MODEL_PATH = model_path
|
|
FS_MODEL = insightface.model_zoo.get_model(model_path, providers=providers)
|
|
|
|
return FS_MODEL
|
|
|
|
def get_face_analyser() -> Any:
|
|
global FACE_ANALYSER
|
|
|
|
with THREAD_LOCK:
|
|
if FACE_ANALYSER is None:
|
|
FACE_ANALYSER = insightface.app.FaceAnalysis(name='buffalo_l', providers=providers)
|
|
return FACE_ANALYSER
|
|
|
|
def getFaceAnalyser():
|
|
face_analyser=insightface.app.FaceAnalysis(name="buffalo_l", providers=providers)
|
|
return face_analyser
|
|
|
|
def get_face_single(img_data: np.ndarray,face_index=0, det_size=(640, 640)):
|
|
face_analyser = get_face_analyser()
|
|
face_analyser.prepare(ctx_id=0, det_size=det_size)
|
|
face = face_analyser.get(img_data)
|
|
|
|
if len(face) == 0 and det_size[0] > 320 and det_size[1] > 320:
|
|
det_size_half = (det_size[0] // 2, det_size[1] // 2)
|
|
return get_face_single(img_data, face_index=face_index, det_size=det_size_half)
|
|
|
|
try:
|
|
return sorted(face, key=lambda x: x.bbox[0])[face_index]
|
|
except IndexError:
|
|
return None
|
|
|
|
|
|
|