Spaces:
Running
Running
File size: 2,807 Bytes
8d11d43 |
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 |
from insightface.app import FaceAnalysis
import numpy as np
import torch
INSIGHTFACE_DETECT_SIZE = 512
class FaceDetector:
def __init__(self, device="cuda"):
self.app = FaceAnalysis(
allowed_modules=["detection", "landmark_2d_106"],
root="checkpoints/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
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,
]
|