scripts/libraries: Add face detection post-processing.

This commit is contained in:
Kwabena W. Agyeman 2025-09-05 16:28:47 -07:00
parent 7124b5134d
commit 65829023f0

View File

@ -34,6 +34,14 @@ from ulab import numpy as np
_NO_DETECTION = const(()) _NO_DETECTION = const(())
def logit(x):
return np.log(x / (1.0 - x))
def sigmoid(x):
return 1.0 / (1.0 + np.exp(-x))
def mod(a, b): def mod(a, b):
return a - (b * (a // b)) return a - (b * (a // b))
@ -49,16 +57,16 @@ def threshold(scores, threshold, scale, find_max=False, find_max_axis=1):
return np.nonzero(scores < threshold)[0] return np.nonzero(scores < threshold)[0]
def quantize(model, value): def quantize(model, value, index=0):
if model.output_dtype[0] == 'f': if model.output_dtype[index] == 'f':
return value return value
return (value / model.output_scale[0]) + model.output_zero_point[0] return (value / model.output_scale[index]) + model.output_zero_point[index]
def dequantize(model, value): def dequantize(model, value, index=0):
if model.output_dtype[0] == 'f': if model.output_dtype[index] == 'f':
return value return value
return (value - model.output_zero_point[0]) * model.output_scale[0] return (value - float(model.output_zero_point[index])) * model.output_scale[index]
class fomo_postprocess: class fomo_postprocess:
@ -141,12 +149,6 @@ class yolo_v2_postprocess:
def __call__(self, model, inputs, outputs): def __call__(self, model, inputs, outputs):
def logit(x):
return np.log(x / (1.0 - x))
def sigmoid(x):
return 1.0 / (1.0 + np.exp(-x))
def softmax(x): def softmax(x):
e_x = np.exp(x - np.max(x, axis=1, keepdims=True)) e_x = np.exp(x - np.max(x, axis=1, keepdims=True))
return e_x / np.sum(e_x, axis=1, keepdims=True) return e_x / np.sum(e_x, axis=1, keepdims=True)
@ -330,3 +332,94 @@ class yolo_v8_postprocess:
nms.add_bounding_box(xmin[i], ymin[i], xmax[i], ymax[i], nms.add_bounding_box(xmin[i], ymin[i], xmax[i], ymax[i],
bb_scores[i], bb_classes[i]) bb_scores[i], bb_classes[i])
return nms.get_bounding_boxes(threshold=self.nms_threshold, sigma=self.nms_sigma) return nms.get_bounding_boxes(threshold=self.nms_threshold, sigma=self.nms_sigma)
class mediapipe_face_detection_postprocess:
_BLAZEFACE_CX = const(0)
_BLAZEFACE_CY = const(1)
_BLAZEFACE_CW = const(2)
_BLAZEFACE_CH = const(3)
_BLAZEFACE_KP = const(4)
def __init__(self, threshold=0.6, anchors=None, nms_threshold=0.1, nms_sigma=0.1):
self.threshold = threshold
self.anchors = anchors
if self.anchors is None:
self.anchors = np.empty((896, 2))
idx = 0
# Generate anchors for 16x16 grid with 2 duplicates and
# 8x8 grid with 6 duplicates to match the model output size.
for grid_size, scales in [(16, 2), (8, 6)]:
for gy in range(grid_size):
cy = (gy + 0.5) / grid_size
for gx in range(grid_size):
cx = (gx + 0.5) / grid_size
for _ in range(scales):
self.anchors[idx, 0] = cx
self.anchors[idx, 1] = cy
idx += 1
self.nms_threshold = nms_threshold
self.nms_sigma = nms_sigma
def blazeface_post_process(self, ih, iw, nms, model, inputs, outputs, score_idx, cords_idx, t, anchors):
s_oh, s_ow, s_oc = model.output_shape[score_idx]
scale = model.output_scale[score_idx]
t = quantize(model, logit(t), index=score_idx)
# Threshold all the scores
score_row_outputs = outputs[score_idx].reshape((s_oh * s_ow * s_oc))
score_indices = threshold(score_row_outputs, t, scale)
if not len(score_indices):
return _NO_DETECTION
# Get the score information
bb_scores = np.take(score_row_outputs, score_indices, axis=0)
bb_scores = sigmoid(dequantize(model, bb_scores, index=score_idx))
# Get the bounding boxes that have a valid score
c_oh, c_ow, c_oc = model.output_shape[cords_idx]
cords_row_outputs = outputs[cords_idx].reshape((c_oh * c_ow, c_oc))
bb = dequantize(model, np.take(cords_row_outputs, score_indices, axis=0), index=cords_idx)
# Get the anchor box information
bb_a_array = np.take(anchors, score_indices, axis=0)
# Compute the bounding box information
ax = bb_a_array[:, _BLAZEFACE_CX]
ay = bb_a_array[:, _BLAZEFACE_CY]
x_center = bb[:, _BLAZEFACE_CX] / iw + ax
y_center = bb[:, _BLAZEFACE_CY] / ih + ay
w_rel = bb[:, _BLAZEFACE_CW] / iw * 0.5
h_rel = bb[:, _BLAZEFACE_CH] / ih * 0.5
# Get the keypoint information
row_count = bb.shape[0]
keypoints = np.empty((row_count, (c_oc - _BLAZEFACE_KP) // 2, 2))
keypoints[:, :, 0] = (bb[:, _BLAZEFACE_KP::2] / iw + ax.reshape((row_count, 1))) * iw
keypoints[:, :, 1] = (bb[:, _BLAZEFACE_KP + 1::2] / ih + ay.reshape((row_count, 1))) * ih
# Scale the bounding boxes to have enough integer precision for NMS
xmin = (x_center - w_rel) * iw
ymin = (y_center - h_rel) * ih
xmax = (x_center + w_rel) * iw
ymax = (y_center + h_rel) * ih
for i in range(bb.shape[0]):
nms.add_bounding_box(xmin[i], ymin[i], xmax[i], ymax[i], bb_scores[i], 0, keypoints=keypoints[i])
def __call__(self, model, inputs, outputs):
ib, ih, iw, ic = model.input_shape[0]
nms = NMS(iw, ih, inputs[0].roi)
output_len = outputs[0].shape[1]
self.blazeface_post_process(ih, iw, nms, model, inputs, outputs, 1, 0,
self.threshold, self.anchors[:output_len])
if output_len < len(self.anchors):
self.blazeface_post_process(ih, iw, nms, model, inputs, outputs, 2, 3,
self.threshold, self.anchors[output_len:])
return nms.get_bounding_boxes(threshold=self.nms_threshold, sigma=self.nms_sigma)