0, 1, 4, 5, 6, 7, 8, 10, 13, 14, 17, 21, 33, 37, 39, 40, 46, 52, 53, 54,

55, 58, 61, 63, 65, 66, 67, 70, 78, 80, 81, 82, 84, 87, 88, 91, 93, 95,

103, 105, 107, 109, 127, 132, 133, 136, 144, 145, 146, 148, 149, 150, 152,

153, 154, 155, 157, 158, 159, 160, 161, 162, 163, 168, 172, 173, 176, 178,

181, 185, 191, 195, 197, 234, 246, 249, 251, 263, 267, 269, 270, 276, 282,

283, 284, 285, 288, 291, 293, 295, 296, 297, 300, 308, 310, 311, 312, 314,

317, 318, 321, 323, 324, 332, 334, 336, 338, 356, 361, 362, 365, 373, 374,

375, 377, 378, 379, 380, 381, 382, 384, 385, 386, 387, 388, 389, 390, 397,

398, 400, 402, 405, 409, 415, 454, 466, 468, 469, 470, 471, 472, 473, 474,

475, 476, 477,

"_neutral", "browDownLeft", "browDownRight", "browInnerUp", "browOuterUpLeft",

"browOuterUpRight", "cheekPuff", "cheekSquintLeft", "cheekSquintRight",

"eyeBlinkLeft", "eyeBlinkRight", "eyeLookDownLeft", "eyeLookDownRight",

"eyeLookInLeft", "eyeLookInRight", "eyeLookOutLeft", "eyeLookOutRight",

"eyeLookUpLeft", "eyeLookUpRight", "eyeSquintLeft", "eyeSquintRight",

"eyeWideLeft", "eyeWideRight", "jawForward", "jawLeft", "jawOpen",

"jawRight", "mouthClose", "mouthDimpleLeft", "mouthDimpleRight",

"mouthFrownLeft", "mouthFrownRight", "mouthFunnel", "mouthLeft",

"mouthLowerDownLeft", "mouthLowerDownRight", "mouthPressLeft",

"mouthPressRight", "mouthPucker", "mouthRight", "mouthRollLower",

"mouthRollUpper", "mouthShrugLower", "mouthShrugUpper", "mouthSmileLeft",

"mouthSmileRight", "mouthStretchLeft", "mouthStretchRight",

"mouthUpperUpLeft", "mouthUpperUpRight", "noseSneerLeft", "noseSneerRight",

"""TF SAME pad (asymmetric on stride-2; PyTorch's symmetric pad undershoots by 1 px)."""

H, W = x.shape[-2], x.shape[-1]

pad_h = max(((H + stride - 1) // stride - 1) * stride + kernel - H, 0)

pad_w = max(((W + stride - 1) // stride - 1) * stride + kernel - W, 0)

if pad_h == 0 and pad_w == 0:

return x

(24, 24, 1), (24, 28, 1), (28, 32, 2), (32, 36, 1),

(36, 42, 1), (42, 48, 2), (48, 56, 1), (56, 64, 1),

(64, 72, 1), (72, 80, 1), (80, 88, 1), (88, 96, 2),

(96, 96, 1), (96, 96, 1), (96, 96, 1), (96, 96, 1),

"""DW 3x3 + PW + residual. Residual max-pools on stride>1, channel-pads on out_ch>in_ch."""

def __init__(self, in_ch: int, out_ch: int, stride: int, device=None, dtype=None, operations=None):

super().__init__()

ops = operations if operations is not None else nn

self.in_ch, self.out_ch, self.stride = in_ch, out_ch, stride

self.depthwise = ops.Conv2d(in_ch, in_ch, 3, stride=stride, padding=0, groups=in_ch, bias=True, device=device, dtype=dtype)

self.pointwise = ops.Conv2d(in_ch, out_ch, 1, padding=0, bias=True, device=device, dtype=dtype)

def forward(self, x: Tensor) -> Tensor:

residual = F.max_pool2d(x, 2, 2) if self.stride > 1 else x

if self.out_ch > self.in_ch:

residual = F.pad(residual, (0, 0, 0, 0, 0, self.out_ch - self.in_ch))

x = _tf_same_pad(x, 3, self.stride) if self.stride > 1 else F.pad(x, (1, 1, 1, 1))

"""Short-range BlazeFace: (B, 3, 128, 128) in [-1, 1] → 896 anchors x 17."""

def __init__(self, device=None, dtype=None, operations=None):

super().__init__()

ops = operations if operations is not None else nn

kw = dict(device=device, dtype=dtype)

self.stem = ops.Conv2d(3, 24, 5, stride=2, padding=0, bias=True, **kw)

self.blocks = nn.ModuleList(BlazeFaceBlock(i, o, s, device=device, dtype=dtype, operations=operations)

for (i, o, s) in _BLAZEFACE_BLOCKS)

# 16²x2 + 8²x6 = 512 + 384 = 896 anchors.

self.cls_16 = ops.Conv2d(88, 2, 1, padding=0, bias=True, **kw)

self.cls_8 = ops.Conv2d(96, 6, 1, padding=0, bias=True, **kw)

self.reg_16 = ops.Conv2d(88, 32, 1, padding=0, bias=True, **kw)

self.reg_8 = ops.Conv2d(96, 96, 1, padding=0, bias=True, **kw)

def forward(self, image_chw_normalized: Tensor) -> tuple[Tensor, Tensor]:

x = F.relu(self.stem(_tf_same_pad(image_chw_normalized, 5, 2)))

# 16x16 tap is block-10 output (before the 88→96 stride-2 in block 11).

for i in range(11):

x = self.blocks[i](x)

feat_16 = x

for i in range(11, 16):

x = self.blocks[i](x)

feat_8 = x

def flat(t, a, k): # NHWC flatten → (B, H*W*A, K)

B, _, H, W = t.shape

return t.permute(0, 2, 3, 1).reshape(B, H * W * a, k)

cls = torch.cat([flat(self.cls_16(feat_16), 2, 1), flat(self.cls_8(feat_8), 6, 1)], dim=1)

reg = torch.cat([flat(self.reg_16(feat_16), 2, 16), flat(self.reg_8(feat_8), 6, 16)], dim=1)

"""Double inverted residual: DW → PW(mid) → DW → PW(out) [+ residual].

def __init__(self, in_ch: int, mid_ch: int, out_ch: int, stride: int, device=None, dtype=None, operations=None):

super().__init__()

ops = operations if operations is not None else nn

kw = dict(device=device, dtype=dtype)

self.has_residual = (in_ch == out_ch and stride == 1)

self.dw1 = ops.Conv2d(in_ch, in_ch, 3, stride=stride, padding=0, groups=in_ch, bias=True, **kw)

self.pw1 = ops.Conv2d(in_ch, mid_ch, 1, padding=0, bias=True, **kw)

self.dw2 = ops.Conv2d(mid_ch, mid_ch, 3, stride=1, padding=0, groups=mid_ch, bias=True, **kw)

self.pw2 = ops.Conv2d(mid_ch, out_ch, 1, padding=0, bias=True, **kw)

def forward(self, x: Tensor) -> Tensor:

residual = x if self.has_residual else None

x = F.relu(self.pw1(self.dw1(F.pad(x, (1, 1, 1, 1)))))

x = self.pw2(self.dw2(F.pad(x, (1, 1, 1, 1))))

(32, 8, 32, 1), (32, 8, 32, 1), # 96²x32

(32, 16, 64, 2), (64, 16, 64, 1), (64, 16, 64, 1), # 48²x64 — tap[0]

(64, 32, 128, 2), (128, 32, 128, 1), (128, 32, 128, 1), # 24²x128 — tap[1]

(128, 48, 192, 2), (192, 48, 192, 1), (192, 48, 192, 1), # 12²x192 — tap[2]

(192, 96, 384, 2), (384, 96, 384, 1), (384, 96, 384, 1), (384, 96, 384, 1), # 6²x384

[(96, 48, 96, 1), (96, 48, 96, 1)], # 12²x96

[(64, 32, 64, 1), (64, 32, 64, 1)], # 24²x64

[(48, 24, 48, 1)], # 48²x48 — feeds the heads

"""TF DEPTH_TO_SPACE in DCR layout (input channels = (i, j, c_out)).

B_, _, H_, W_ = t.shape

t = t.reshape(B_, r, r, c_out, H_, W_)

t = t.permute(0, 3, 4, 1, 5, 2).contiguous()

"""Full-range face detector: (B, 3, 192, 192) in [-1, 1] → 2304 anchors x 17 values."""

def __init__(self, device=None, dtype=None, operations=None):

super().__init__()

ops = operations if operations is not None else nn

kw = dict(device=device, dtype=dtype)

mk_block = lambda i, m, o, s: FRBlock(i, m, o, s, device=device, dtype=dtype, operations=operations)

self.stem = ops.Conv2d(3, 32, 3, stride=2, padding=0, bias=True, **kw)

self.backbone = nn.ModuleList(mk_block(i, m, o, s) for (i, m, o, s) in _FR_BACKBONE_BLOCKS)

self.lateral_convs = nn.ModuleList(ops.Conv2d(i, o, 1, padding=0, bias=True, **kw) for (i, o) in _FR_LATERAL_CHANNELS)

self.top_conv = ops.Conv2d(384, 96, 1, padding=0, bias=True, **kw)

self.decoder_levels = nn.ModuleList(

nn.ModuleList(mk_block(i, m, o, s) for (i, m, o, s) in lvl) for lvl in _FR_DECODER_BLOCKS

)

# 96→64 before 12→24, 64→48 before 24→48.

self.decoder_reduce_convs = nn.ModuleList([

ops.Conv2d(96, 64, 1, padding=0, bias=True, **kw),

ops.Conv2d(64, 48, 1, padding=0, bias=True, **kw),

])

# Heads mix 2x2-cell info via DW-stride-2 + depth_to_space block_size=2.

self.cls_conv = ops.Conv2d(48, 4, 1, padding=0, bias=True, **kw)

self.cls_dw = ops.Conv2d(4, 4, 3, stride=2, padding=0, groups=4, bias=True, **kw)

self.reg_conv = ops.Conv2d(48, 64, 1, padding=0, bias=True, **kw)

self.reg_dw = ops.Conv2d(64, 64, 3, stride=2, padding=0, groups=64, bias=True, **kw)

def forward(self, image_chw_normalized: Tensor) -> tuple[Tensor, Tensor]:

# Symmetric pad-1 throughout (full-range tflite uses explicit TF PAD, not SAME).

x = F.relu(self.stem(F.pad(image_chw_normalized, (1, 1, 1, 1))))

tap_set = set(_FR_LATERAL_TAP_INDICES)

laterals: list[Tensor] = []

for i, blk in enumerate(self.backbone):

x = blk(x)

if i in tap_set:

laterals.append(x)

# top_conv / lateral_convs / decoder_reduce_convs all have fused ReLU in the tflite.

p = F.relu(self.top_conv(x))

laterals_rev = list(reversed(laterals))

lateral_convs_rev = list(reversed(self.lateral_convs))

for level in range(len(self.decoder_levels)):

lateral = laterals_rev[level]

p = F.interpolate(p, size=lateral.shape[-2:], mode="bilinear", align_corners=False)

p = p + F.relu(lateral_convs_rev[level](lateral))

for blk in self.decoder_levels[level]:

p = blk(p)

if level < len(self.decoder_reduce_convs):

p = F.relu(self.decoder_reduce_convs[level](p))

c = self.cls_dw(F.pad(self.cls_conv(p), (1, 1, 1, 1)))

c = _dcr_depth_to_space(c, r=2, c_out=1)

r = self.reg_dw(F.pad(self.reg_conv(p), (1, 1, 1, 1)))

r = _dcr_depth_to_space(r, r=2, c_out=16)

B = c.shape[0]

cls_out = c.permute(0, 2, 3, 1).reshape(B, _BF_FR_NUM_ANCHORS, 1)

reg_out = r.permute(0, 2, 3, 1).reshape(B, _BF_FR_NUM_ANCHORS, 16)

"""2304 anchors over 48x48; anchor_w=anchor_h=1 (fixed_anchor_size)."""

feat = _BF_FR_GRID

yy, xx = np.meshgrid(np.arange(feat, dtype=np.float32), np.arange(feat, dtype=np.float32), indexing="ij")

cx, cy, ones = (xx + 0.5) / feat, (yy + 0.5) / feat, np.ones_like(xx)

score_thresh: float = _BF_MIN_SCORE) -> np.ndarray:

"""Same decode as short-range with 2304-anchor grid and box_scale=192."""

scores = expit(np.clip(classificators[:, 0], -_BF_FR_SCORE_CLIP, _BF_FR_SCORE_CLIP))

keep = scores >= score_thresh

if not keep.any():

return np.empty((0, 17), dtype=np.float32)

r = regressors[keep] / _BF_FR_BOX_SCALE

a = _blazeface_full_range_anchors()[keep]

cxs, cys, aws, ahs = a[:, 0:1], a[:, 1:2], a[:, 2:3], a[:, 3:4]

xc, yc = r[:, 0:1] * aws + cxs, r[:, 1:2] * ahs + cys

w, h = r[:, 2:3] * aws, r[:, 3:4] * ahs

out = np.empty((r.shape[0], 17), dtype=np.float32)

out[:, 0:1], out[:, 1:2], out[:, 2:3], out[:, 3:4] = xc - w / 2, yc - h / 2, xc + w / 2, yc + h / 2

out[:, 4:16:2] = r[:, _BF_KP_OFFSET::2] * aws + cxs

out[:, 5:16:2] = r[:, _BF_KP_OFFSET + 1::2] * ahs + cys

out[:, 16] = scores[keep]

(16, 16, 1), (16, 16, 1), (16, 32, 2), (32, 32, 1), (32, 32, 1), (32, 64, 2),

(64, 64, 1), (64, 64, 1), (64, 128, 2), (128, 128, 1), (128, 128, 1), (128, 128, 2),

(128, 128, 1), (128, 128, 1), (128, 128, 2), (128, 128, 1), (128, 128, 1),

"""PReLU BlazeBlock: PReLU between DW and PW, and after the residual add."""

def __init__(self, in_ch: int, out_ch: int, stride: int, device=None, dtype=None, operations=None):

super().__init__()

ops = operations if operations is not None else nn

kw = dict(device=device, dtype=dtype)

self.in_ch, self.out_ch, self.stride = in_ch, out_ch, stride

self.depthwise = ops.Conv2d(in_ch, in_ch, 3, stride=stride, padding=0, groups=in_ch, bias=True, **kw)

self.prelu_dwise = nn.PReLU(num_parameters=in_ch, **kw)

self.pointwise = ops.Conv2d(in_ch, out_ch, 1, padding=0, bias=True, **kw)

self.prelu_out = nn.PReLU(num_parameters=out_ch, **kw)

def forward(self, x: Tensor) -> Tensor:

residual = F.max_pool2d(x, 2, 2) if self.stride > 1 else x

if self.out_ch > self.in_ch:

residual = F.pad(residual, (0, 0, 0, 0, 0, self.out_ch - self.in_ch))

x = _tf_same_pad(x, 3, self.stride) if self.stride > 1 else F.pad(x, (1, 1, 1, 1))

NUM_LANDMARKS = 478

def __init__(self, device=None, dtype=None, operations=None):

super().__init__()

ops = operations if operations is not None else nn

kw = dict(device=device, dtype=dtype)

self.stem = ops.Conv2d(3, 16, 3, stride=2, padding=0, bias=True, **kw)

self.prelu_stem = nn.PReLU(num_parameters=16, **kw)

self.blocks = nn.ModuleList(FaceMeshBlock(i, o, s, device=device, dtype=dtype, operations=operations)

for (i, o, s) in _FACEMESH_BLOCKS)

self.head_reduce = ops.Conv2d(128, 8, 1, padding=0, bias=True, **kw)

self.prelu_head_reduce = nn.PReLU(num_parameters=8, **kw)

self.head_block = FaceMeshBlock(8, 8, 1, device=device, dtype=dtype, operations=operations)

self.head_presence = ops.Conv2d(8, 1, 3, padding=0, bias=True, **kw)

self.head_landmarks = ops.Conv2d(8, self.NUM_LANDMARKS * 3, 3, padding=0, bias=True, **kw)

def forward(self, face_chw_normalized: Tensor) -> tuple[Tensor, Tensor]:

"""(B, 3, 192, 192) in [0, 1] → ((B, 478, 3) landmarks in 192-canonical, (B,) presence)."""

x = self.prelu_stem(self.stem(_tf_same_pad(face_chw_normalized, 3, 2)))

for blk in self.blocks:

x = blk(x)

x = self.prelu_head_reduce(self.head_reduce(x))

x = self.head_block(x)

B = x.shape[0]

presence = self.head_presence(x).reshape(B)

lmks = self.head_landmarks(x).reshape(B, self.NUM_LANDMARKS, 3)

"""MLP-Mixer block: token-mixing MLP (over tokens) → channel-mixing MLP (over dim).

def __init__(self, num_tokens: int, token_dim: int, token_hidden: int, channel_hidden: int,

device=None, dtype=None, operations=None):

super().__init__()

ops = operations if operations is not None else nn

kw = dict(device=device, dtype=dtype)

# bias=False → no LN beta (matches MP).

self.ln1 = ops.LayerNorm(token_dim, eps=_BS_LN_EPS, bias=False, **kw)

self.ln2 = ops.LayerNorm(token_dim, eps=_BS_LN_EPS, bias=False, **kw)

self.token_mlp1 = ops.Linear(num_tokens, token_hidden, bias=True, **kw)

self.token_mlp2 = ops.Linear(token_hidden, num_tokens, bias=True, **kw)

self.channel_mlp1 = ops.Linear(token_dim, channel_hidden, bias=True, **kw)

self.channel_mlp2 = ops.Linear(channel_hidden, token_dim, bias=True, **kw)

def forward(self, x: Tensor) -> Tensor:

y = self.ln1(x).transpose(1, 2)

x = x + self.token_mlp2(F.relu(self.token_mlp1(y))).transpose(1, 2)

def __init__(self, device=None, dtype=None, operations=None):

super().__init__()

ops = operations if operations is not None else nn

kw = dict(device=device, dtype=dtype)

self.token_reduce = ops.Linear(_BS_NUM_INPUT_LANDMARKS, _BS_NUM_TOKENS_REDUCED, bias=True, **kw)

self.token_embed = ops.Linear(2, _BS_TOKEN_DIM, bias=True, **kw)

self.cls_token = nn.Parameter(torch.zeros(1, 1, _BS_TOKEN_DIM, **kw))

self.blocks = nn.ModuleList(

MlpMixerBlock(_BS_NUM_TOKENS, _BS_TOKEN_DIM, _BS_TOKEN_MIX_HIDDEN, _BS_CHANNEL_MIX_HIDDEN,

device=device, dtype=dtype, operations=operations) for _ in range(4)

)

self.head = ops.Linear(_BS_TOKEN_DIM, _BS_NUM_BLENDSHAPES, bias=True, **kw)

@staticmethod

def _input_normalize(landmarks_2d: Tensor) -> Tensor:

# Centroid-subtract → L2 scale → x0.5. The 0.5 is baked into training.

centroid = landmarks_2d.mean(dim=1, keepdim=True)

x = landmarks_2d - centroid

mag = torch.sqrt((x * x).sum(dim=-1, keepdim=True))

scale = mag.mean(dim=1, keepdim=True)

return (x / scale.clamp(min=1e-12)) * 0.5

def forward(self, landmarks_2d: Tensor) -> Tensor:

"""(B, 146, 2) → (B, 52) in [0, 1]. Input units don't matter (centroid + L2 normalize)."""

x = self._input_normalize(landmarks_2d)

x = self.token_reduce(x.transpose(1, 2)).transpose(1, 2)

x = self.token_embed(x)

cls = self.cls_token.expand(x.shape[0], -1, -1)

x = torch.cat([cls, x], dim=1)

for blk in self.blocks:

x = blk(x)

"""896 anchors per SsdAnchorsCalculator (fixed_anchor_size → anchor_w=anchor_h=1)."""

per_ar = len(_BF_ASPECT_RATIOS) + (1 if _BF_INTERP_SCALE_AR > 0 else 0)

layer_anchors: List[np.ndarray] = []

layer = 0

while layer < _BF_NUM_LAYERS:

stride = _BF_STRIDES[layer]

last = layer

while last < _BF_NUM_LAYERS and _BF_STRIDES[last] == stride:

last += 1

per_cell = per_ar * (last - layer)

feat = (_BF_INPUT_SIZE + stride - 1) // stride

yy, xx = np.meshgrid(np.arange(feat, dtype=np.float32), np.arange(feat, dtype=np.float32), indexing="ij")

cx, cy, ones = (xx + _BF_ANCHOR_OFFSET_X) / feat, (yy + _BF_ANCHOR_OFFSET_Y) / feat, np.ones_like(xx)

cell = np.stack([cx, cy, ones, ones], axis=-1).reshape(-1, 4)

layer_anchors.append(np.repeat(cell, per_cell, axis=0))

layer = last

out = np.concatenate(layer_anchors, axis=0)

assert out.shape == (896, 4), out.shape

score_thresh: float = _BF_MIN_SCORE) -> np.ndarray:

"""Decode (regs (896,16), cls (896,1)) → (N, 17) = [xyxy, kp0x..kp5y, score] in [0, 1]."""

scores = expit(np.clip(classificators[:, 0], -_BF_SCORE_CLIP, _BF_SCORE_CLIP))

keep = scores >= score_thresh

if not keep.any():

return np.empty((0, 17), dtype=np.float32)

r = regressors[keep] / _BF_BOX_SCALE

a = _blazeface_anchors()[keep] # (N, 4) cx, cy, 1, 1

cxs, cys, aws, ahs = a[:, 0:1], a[:, 1:2], a[:, 2:3], a[:, 3:4]

xc, yc = r[:, 0:1] * aws + cxs, r[:, 1:2] * ahs + cys

w, h = r[:, 2:3] * aws, r[:, 3:4] * ahs

out = np.empty((r.shape[0], 17), dtype=np.float32)

out[:, 0:1], out[:, 1:2], out[:, 2:3], out[:, 3:4] = xc - w / 2, yc - h / 2, xc + w / 2, yc + h / 2

out[:, 4:16:2] = r[:, _BF_KP_OFFSET::2] * aws + cxs

out[:, 5:16:2] = r[:, _BF_KP_OFFSET + 1::2] * ahs + cys

out[:, 16] = scores[keep]

"""MP weighted NMS — kept boxes are score-weighted averages of overlapping detections."""

if detections.shape[0] == 0:

return detections

dets = detections[np.argsort(-detections[:, 16])]

N = dets.shape[0]

areas = np.clip(dets[:, 2] - dets[:, 0], 0, None) * np.clip(dets[:, 3] - dets[:, 1], 0, None)

kept: List[np.ndarray] = []

used = np.zeros(N, dtype=bool)

for i in range(N):

if used[i]:

continue

ax1, ay1, ax2, ay2 = dets[i, 0:4]

merge_idx = [i]

for j in range(i + 1, N):

if used[j]:

continue

bx1, by1, bx2, by2 = dets[j, 0:4]

iw = max(0.0, min(ax2, bx2) - max(ax1, bx1))

ih = max(0.0, min(ay2, by2) - max(ay1, by1))

inter = iw * ih

union = areas[i] + areas[j] - inter

if union > 0 and inter / union > iou_thresh: # strict > matches MP

merge_idx.append(j)

used[j] = True

used[i] = True

cluster = dets[merge_idx]

ws = cluster[:, 16:17]

ws_sum = ws.sum()

merged = np.copy(cluster[0])

if ws_sum > 0:

merged[:16] = (cluster[:, :16] * ws).sum(axis=0) / ws_sum

kept.append(merged)

"""Detection (normalized) → rotated 1.5xbbox ROI in image pixels (anisotropic)."""

xmin, ymin, xmax, ymax = detection[0:4]

lx = detection[4 + _FACE_LEFT_EYE_KP * 2 + 0] * image_w

ly = detection[4 + _FACE_LEFT_EYE_KP * 2 + 1] * image_h

rx = detection[4 + _FACE_RIGHT_EYE_KP * 2 + 0] * image_w

ry = detection[4 + _FACE_RIGHT_EYE_KP * 2 + 1] * image_h

# Image-y-down convention: angle = target - atan2(-dy, dx).

angle = _FACE_ROI_TARGET_ANGLE - math.atan2(ly - ry, rx - lx)

return (float((xmin + xmax) * 0.5 * image_w),

float((ymin + ymax) * 0.5 * image_h),

float((xmax - xmin) * image_w * _FACE_ROI_SCALE_X),

float((ymax - ymin) * image_h * _FACE_ROI_SCALE_Y),

"""Bilinear-sample image_chw at corner-aligned (src_x, src_y)."""

H, W = int(image_chw.shape[-2]), int(image_chw.shape[-1])

grid = torch.stack([(2.0 * src_x + 1.0) / W - 1.0,

(2.0 * src_y + 1.0) / H - 1.0], dim=-1).unsqueeze(0)

return F.grid_sample(image_chw.unsqueeze(0), grid, mode="bilinear",

angle: float, output_size: int = _FM_INPUT_SIZE) -> Tensor:

"""Rotated rect → output_size² with BORDER_REPLICATE. image_chw must be in [0, 1]."""

s_x, s_y = width / output_size, height / output_size

cos_a, sin_a = math.cos(angle), math.sin(angle)

arange = torch.arange(output_size, dtype=image_chw.dtype, device=image_chw.device) - output_size * 0.5

v_grid, u_grid = torch.meshgrid(arange, arange, indexing="ij")

src_x = cx + u_grid * s_x * cos_a - v_grid * s_y * sin_a

src_y = cy + u_grid * s_x * sin_a + v_grid * s_y * cos_a

"""Centered max(W,H) square → target² with BORDER_ZERO + [-1, 1] norm.

H, W = int(image_chw_raw.shape[1]), int(image_chw_raw.shape[2])

sub_rect_size = float(max(W, H))

sub_rect_cx, sub_rect_cy = W * 0.5, H * 0.5

s = sub_rect_size / target

arange = torch.arange(target, dtype=image_chw_raw.dtype, device=image_chw_raw.device) - target * 0.5

v_grid, u_grid = torch.meshgrid(arange, arange, indexing="ij")

out = _sample_warp(image_chw_raw, sub_rect_cx + u_grid * s, sub_rect_cy + v_grid * s, "zeros")

"""BlazeFace → FaceMesh v2 → blendshapes."""

def __init__(self, device=None, dtype=None, operations=None, detector_variant: str = "short"):

super().__init__()

self.detector_variant = detector_variant

if detector_variant == "both":

self.detector = BlazeFace(device=device, dtype=dtype, operations=operations)

self.detector_full = BlazeFaceFullRange(device=device, dtype=dtype, operations=operations)

else:

det_cls = {"short": BlazeFace, "full": BlazeFaceFullRange}[detector_variant]

self.detector = det_cls(device=device, dtype=dtype, operations=operations)

self.mesh = FaceMesh(device=device, dtype=dtype, operations=operations)

self.blendshapes = FaceBlendshapes(device=device, dtype=dtype, operations=operations)

self.register_buffer("_bs_idx", torch.tensor(_BS_INPUT_INDICES, dtype=torch.long), persistent=False)

def _detector(self, variant: str) -> nn.Module:

if variant not in ("short", "full"):

raise ValueError(f"Unknown face detector variant: {variant!r}")

if self.detector_variant == "both":

return self.detector_full if variant == "full" else self.detector

if variant != self.detector_variant:

raise ValueError(

f"FaceLandmarker was initialized with the {self.detector_variant!r} detector, not {variant!r}"

)

return self.detector

def run_detector_batch(self, images_rgb_uint8: List[np.ndarray],

score_thresh: float = _BF_MIN_SCORE,

iou_thresh: float = 0.5,

variant: Optional[str] = None):

"""Batched detector pass. Returns (img_raws, sub_rects, sizes, per_frame_decoded)

if not images_rgb_uint8:

return [], [], [], []

if variant is None:

variant = "short" if self.detector_variant == "both" else self.detector_variant

detector = self._detector(variant)

device, dtype = detector.stem.weight.device, detector.stem.weight.dtype

det_input_size, decode_fn = ((_BF_FR_INPUT_SIZE, _decode_blazeface_full_range)

if variant == "full"

else (_BF_INPUT_SIZE, _decode_blazeface))

# Same-size frames: stack once and transfer once. Variable size falls back

# to per-image (only triggers for SAM3DBody's head crops).

sizes = [tuple(img.shape[:2]) for img in images_rgb_uint8]

if len(set(sizes)) == 1:

batch_chw = torch.from_numpy(np.stack(images_rgb_uint8, axis=0)).to(device, dtype).movedim(-1, -3).contiguous()

img_raws = [batch_chw[bi] for bi in range(batch_chw.shape[0])]

else:

img_raws = [torch.from_numpy(img).to(device, dtype).movedim(-1, -3).contiguous() for img in images_rgb_uint8]

warps = [_blazeface_input_warp(img_raw, det_input_size) for img_raw in img_raws]

det_crops = [w[0] for w in warps]

sub_rects = [(w[1], w[2], w[3]) for w in warps]

regs_b, cls_b = detector(torch.stack(det_crops, dim=0))

regs_np, cls_np = regs_b.float().cpu().numpy(), cls_b.float().cpu().numpy()

per_frame = []

for b in range(len(images_rgb_uint8)):

decoded = decode_fn(regs_np[b], cls_np[b], score_thresh=score_thresh)

per_frame.append(_weighted_nms(decoded, iou_thresh=iou_thresh) if decoded.shape[0] > 0 else decoded)

return img_raws, sub_rects, sizes, per_frame

def detect_batch(self, images_rgb_uint8: List[np.ndarray], num_faces: int = 1,

score_thresh: float = _BF_MIN_SCORE,

variant: Optional[str] = None) -> List[List[dict]]:

"""Full pipeline batched across `images_rgb_uint8`. Returns one face-dict

def normalize_detections(sub_rects, sizes, detections):

# tensor-normalized → image-normalized [0,1] for _detection_to_face_rect.

for b, decoded in enumerate(detections):

if decoded.shape[0] == 0:

continue

cx, cy, size = sub_rects[b]

H, W = sizes[b]

sx0, sy0 = cx - size * 0.5, cy - size * 0.5

decoded[:, 0:16:2] = (sx0 + size * decoded[:, 0:16:2]) / W

decoded[:, 1:16:2] = (sy0 + size * decoded[:, 1:16:2]) / H

if num_faces > 0:

detections[b] = decoded[: int(num_faces)]

if variant is None and self.detector_variant == "both":

img_raws, short_rects, sizes, short_dets = self.run_detector_batch(

images_rgb_uint8, score_thresh=score_thresh, variant="short",

)

_, full_rects, full_sizes, full_dets = self.run_detector_batch(

images_rgb_uint8, score_thresh=score_thresh, variant="full",

)

normalize_detections(short_rects, sizes, short_dets)

normalize_detections(full_rects, full_sizes, full_dets)

per_frame_dets = [

short if len(short) >= len(full) else full

for short, full in zip(short_dets, full_dets)

]

else:

img_raws, sub_rects, sizes, per_frame_dets = self.run_detector_batch(

images_rgb_uint8, score_thresh=score_thresh, variant=variant,

)

normalize_detections(sub_rects, sizes, per_frame_dets)

# Collect every detected face across all frames into one mesh input.

face_params: List[Tuple[int, float, float, float, float, float, float]] = []

mesh_crops: List[Tensor] = []

for b, dets in enumerate(per_frame_dets):

if dets.shape[0] == 0:

continue

H, W = sizes[b]

img_for_mesh = img_raws[b] / 255.0

for det in dets:

cx, cy, w, h, angle = _detection_to_face_rect(det, W, H)

mesh_crops.append(_warp_face_crop(img_for_mesh, cx, cy, w, h, angle, _FM_INPUT_SIZE))

face_params.append((b, float(det[16]), cx, cy, w, h, angle))

results: List[List[dict]] = [[] for _ in range(len(images_rgb_uint8))]

if not mesh_crops:

return results

lmks_canon_b, presence_b = self.mesh(torch.stack(mesh_crops, dim=0))

bs_out_b = self.blendshapes(lmks_canon_b[:, self._bs_idx, :2])

# Batched canonical→image affine

params_t = torch.tensor(

[(cx, cy, w, h, math.cos(a), math.sin(a)) for (_b, _s, cx, cy, w, h, a) in face_params],

device=lmks_canon_b.device, dtype=lmks_canon_b.dtype,

)

cxs, cys, ws, hs, cos_a, sin_a = params_t.unbind(dim=1)

inv = 1.0 / _FM_INPUT_SIZE

u = lmks_canon_b[..., 0] - _FM_INPUT_SIZE * 0.5

v = lmks_canon_b[..., 1] - _FM_INPUT_SIZE * 0.5

lmks_xy_t = torch.stack([

cxs[:, None] + u * (ws * inv * cos_a)[:, None] - v * (hs * inv * sin_a)[:, None],

cys[:, None] + u * (ws * inv * sin_a)[:, None] + v * (hs * inv * cos_a)[:, None],

], dim=-1)

lmks_xy_np = lmks_xy_t.float().cpu().numpy()

lmks_canon_np = lmks_canon_b.float().cpu().numpy()

presence_np = presence_b.float().cpu().numpy()

bs_np = bs_out_b.float().cpu().numpy()

for i, (b, score, *_) in enumerate(face_params):

lmks_xy = lmks_xy_np[i]

mn, mx = lmks_xy.min(0), lmks_xy.max(0)

results[b].append({

"bbox_xyxy": np.array([mn[0], mn[1], mx[0], mx[1]], dtype=np.float32),