Spaces:
Runtime error
Runtime error
File size: 3,063 Bytes
d8265e9 98eeefd d8265e9 98eeefd d8265e9 98eeefd | 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 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 | import os
import sys
import numpy as np
import torch
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
from config import MODELS_DIR
INSIGHTFACE_DETECT_SIZE = 512
def cuda_to_int(cuda_str: str) -> int:
"""
Convert the string with format "cuda:X" to integer X.
"""
if cuda_str == "cuda":
return 0
device = torch.device(cuda_str)
if device.type != "cuda":
raise ValueError(f"Device type must be 'cuda', got: {device.type}")
return device.index
LMK_ADAPT_ORIGIN_ORDER = [
1,
10,
12,
14,
16,
3,
5,
7,
0,
23,
21,
19,
32,
30,
28,
26,
17,
43,
48,
49,
51,
50,
102,
103,
104,
105,
101,
73,
74,
86,
]
class FaceDetector:
def __init__(self, device="cuda"):
from insightface.app import FaceAnalysis
self.app = FaceAnalysis(
allowed_modules=["detection", "landmark_2d_106"],
root=f"{MODELS_DIR}/auxiliary",
providers=["CUDAExecutionProvider"],
)
self.app.prepare(
ctx_id=cuda_to_int(device),
det_size=(INSIGHTFACE_DETECT_SIZE, INSIGHTFACE_DETECT_SIZE),
)
def __call__(self, frame, threshold=0.5):
f_h, f_w, _ = frame.shape
faces = self.app.get(frame)
get_face_store = None
max_size = 0
if len(faces) == 0:
return None, None
else:
for face in faces:
bbox = face.bbox.astype(np.int_).tolist()
w, h = bbox[2] - bbox[0], bbox[3] - bbox[1]
if w < 50 or h < 80:
continue
if w / h > 1.5 or w / h < 0.2:
continue
if face.det_score < threshold:
continue
size_now = w * h
if size_now > max_size:
max_size = size_now
get_face_store = face
if get_face_store is None:
return None, None
else:
face = get_face_store
lmk = np.round(face.landmark_2d_106).astype(np.int_)
halk_face_coord = np.mean([lmk[74], lmk[73]], axis=0) # lmk[73]
sub_lmk = lmk[LMK_ADAPT_ORIGIN_ORDER]
halk_face_dist = np.max(sub_lmk[:, 1]) - halk_face_coord[1]
upper_bond = halk_face_coord[1] - halk_face_dist # *0.94
x1, y1, x2, y2 = (
np.min(sub_lmk[:, 0]),
int(upper_bond),
np.max(sub_lmk[:, 0]),
np.max(sub_lmk[:, 1]),
)
if y2 - y1 <= 0 or x2 - x1 <= 0 or x1 < 0:
x1, y1, x2, y2 = face.bbox.astype(np.int_).tolist()
y2 += int((x2 - x1) * 0.1)
x1 -= int((x2 - x1) * 0.05)
x2 += int((x2 - x1) * 0.05)
x1 = max(0, x1)
y1 = max(0, y1)
x2 = min(f_w, x2)
y2 = min(f_h, y2)
return (x1, y1, x2, y2), lmk
|