463 lines
13 KiB
Go
463 lines
13 KiB
Go
package service
|
|
|
|
import (
|
|
"bytes"
|
|
_ "embed"
|
|
"errors"
|
|
"fmt"
|
|
"image"
|
|
"image/color"
|
|
"image/draw"
|
|
"image/png"
|
|
"math"
|
|
"os"
|
|
"runtime"
|
|
"sort"
|
|
"sync"
|
|
|
|
ort "github.com/yalue/onnxruntime_go"
|
|
)
|
|
|
|
// YuNet 人脸检测模型（opencv_zoo face_detection_yunet_2023mar，输入尺寸已改为
|
|
// 动态），编译进二进制，运行时只额外依赖 onnxruntime 动态库。
|
|
//
|
|
//go:embed assets/yunet.onnx
|
|
var yunetModel []byte
|
|
|
|
const (
|
|
// 检出分数下限与 NMS 的 IoU 阈值
|
|
faceScoreThreshold = 0.35
|
|
faceNMSIoU = 0.3
|
|
// 推理输入的长边上限：更大的图先等比缩小，检出框再映射回原图，
|
|
// 避免超大参考图把内存和耗时拉爆。
|
|
faceMaxInferSide = 2560
|
|
)
|
|
|
|
// YuNet 的三个输出分支步长
|
|
var faceStrides = []int{8, 16, 32}
|
|
|
|
// ErrNoFaceDetected 表示图中没有检出人脸，调用方按原图处理。
|
|
var ErrNoFaceDetected = errors.New("no face detected")
|
|
|
|
var (
|
|
faceOnce sync.Once
|
|
faceSession *ort.DynamicAdvancedSession
|
|
faceInitErr error
|
|
)
|
|
|
|
// onnxruntimeLibPath 返回 onnxruntime 动态库路径：ONNXRUNTIME_LIB_PATH 优先，
|
|
// 否则用各平台的默认位置。
|
|
func onnxruntimeLibPath() string {
|
|
if p := os.Getenv("ONNXRUNTIME_LIB_PATH"); p != "" {
|
|
return p
|
|
}
|
|
if runtime.GOOS == "windows" {
|
|
return "onnxruntime.dll"
|
|
}
|
|
return "/usr/local/lib/libonnxruntime.so"
|
|
}
|
|
|
|
// faceOutputNames 是 YuNet 需要读取的输出名，顺序与 readOutputs 的下标约定一致：
|
|
// 先 cls_*、再 obj_*、最后 bbox_*（关键点分支用不到）。
|
|
func faceOutputNames() []string {
|
|
names := make([]string, 0, len(faceStrides)*3)
|
|
for _, prefix := range []string{"cls", "obj", "bbox"} {
|
|
for _, s := range faceStrides {
|
|
names = append(names, fmt.Sprintf("%s_%d", prefix, s))
|
|
}
|
|
}
|
|
return names
|
|
}
|
|
|
|
func faceDetector() (*ort.DynamicAdvancedSession, error) {
|
|
faceOnce.Do(func() {
|
|
ort.SetSharedLibraryPath(onnxruntimeLibPath())
|
|
if err := ort.InitializeEnvironment(); err != nil {
|
|
faceInitErr = fmt.Errorf("onnxruntime init: %w", err)
|
|
return
|
|
}
|
|
faceSession, faceInitErr = ort.NewDynamicAdvancedSessionWithONNXData(
|
|
yunetModel, []string{"input"}, faceOutputNames(), nil)
|
|
})
|
|
if faceInitErr != nil {
|
|
return nil, faceInitErr
|
|
}
|
|
return faceSession, nil
|
|
}
|
|
|
|
type faceDetection struct {
|
|
rect image.Rectangle
|
|
score float32
|
|
}
|
|
|
|
// detectFaces 返回图中的人脸矩形框（坐标基于 src 的原始尺寸）。
|
|
func detectFaces(src image.Image) ([]image.Rectangle, error) {
|
|
sess, err := faceDetector()
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
bounds := src.Bounds()
|
|
long := bounds.Dx()
|
|
if bounds.Dy() > long {
|
|
long = bounds.Dy()
|
|
}
|
|
scale := 1.0
|
|
if long > faceMaxInferSide {
|
|
scale = float64(faceMaxInferSide) / float64(long)
|
|
}
|
|
inW := int(float64(bounds.Dx()) * scale)
|
|
inH := int(float64(bounds.Dy()) * scale)
|
|
if inW < 1 || inH < 1 {
|
|
return nil, nil
|
|
}
|
|
// 输入补齐到 32 的整数倍，三个步长分支才有整数网格
|
|
padW := (inW + 31) / 32 * 32
|
|
padH := (inH + 31) / 32 * 32
|
|
|
|
// YuNet 吃 BGR、NCHW、未归一化的 0~255 像素
|
|
pixels := make([]float32, 3*padW*padH)
|
|
plane := padW * padH
|
|
for y := 0; y < inH; y++ {
|
|
srcY := bounds.Min.Y + int(float64(y)/scale)
|
|
for x := 0; x < inW; x++ {
|
|
r, g, b, _ := src.At(bounds.Min.X+int(float64(x)/scale), srcY).RGBA()
|
|
i := y*padW + x
|
|
pixels[i] = float32(b >> 8)
|
|
pixels[plane+i] = float32(g >> 8)
|
|
pixels[2*plane+i] = float32(r >> 8)
|
|
}
|
|
}
|
|
input, err := ort.NewTensor(ort.NewShape(1, 3, int64(padH), int64(padW)), pixels)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer input.Destroy()
|
|
|
|
outputs := make([]ort.Value, len(faceStrides)*3)
|
|
if err := sess.Run([]ort.Value{input}, outputs); err != nil {
|
|
return nil, err
|
|
}
|
|
defer func() {
|
|
for _, out := range outputs {
|
|
if out != nil {
|
|
out.Destroy()
|
|
}
|
|
}
|
|
}()
|
|
|
|
branch := func(i int) ([]float32, error) {
|
|
t, ok := outputs[i].(*ort.Tensor[float32])
|
|
if !ok {
|
|
return nil, fmt.Errorf("yunet output %d is not a float32 tensor", i)
|
|
}
|
|
return t.GetData(), nil
|
|
}
|
|
inferBounds := image.Rect(0, 0, inW, inH)
|
|
var dets []faceDetection
|
|
for si, stride := range faceStrides {
|
|
cls, err := branch(si)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
obj, err := branch(len(faceStrides) + si)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
box, err := branch(2*len(faceStrides) + si)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
cols, rows := padW/stride, padH/stride
|
|
for row := 0; row < rows; row++ {
|
|
for col := 0; col < cols; col++ {
|
|
idx := row*cols + col
|
|
score := float32(math.Sqrt(float64(clampUnit(cls[idx]) * clampUnit(obj[idx]))))
|
|
if score < faceScoreThreshold {
|
|
continue
|
|
}
|
|
cx := (float32(col) + box[idx*4]) * float32(stride)
|
|
cy := (float32(row) + box[idx*4+1]) * float32(stride)
|
|
w := float32(math.Exp(float64(box[idx*4+2]))) * float32(stride)
|
|
h := float32(math.Exp(float64(box[idx*4+3]))) * float32(stride)
|
|
rect := image.Rect(int(cx-w/2), int(cy-h/2), int(cx+w/2), int(cy+h/2)).Intersect(inferBounds)
|
|
if rect.Dx() > 0 && rect.Dy() > 0 {
|
|
dets = append(dets, faceDetection{rect: rect, score: score})
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
boxes := make([]image.Rectangle, 0, len(dets))
|
|
for _, d := range suppressOverlaps(dets, faceNMSIoU) {
|
|
rect := d.rect
|
|
if scale != 1 {
|
|
rect = image.Rect(
|
|
int(float64(rect.Min.X)/scale), int(float64(rect.Min.Y)/scale),
|
|
int(float64(rect.Max.X)/scale), int(float64(rect.Max.Y)/scale),
|
|
).Intersect(image.Rect(0, 0, bounds.Dx(), bounds.Dy()))
|
|
}
|
|
if rect.Dx() > 0 && rect.Dy() > 0 {
|
|
boxes = append(boxes, rect.Add(bounds.Min))
|
|
}
|
|
}
|
|
return boxes, nil
|
|
}
|
|
|
|
func clampUnit(v float32) float32 {
|
|
if v < 0 {
|
|
return 0
|
|
}
|
|
if v > 1 {
|
|
return 1
|
|
}
|
|
return v
|
|
}
|
|
|
|
// suppressOverlaps 按分数从高到低做 NMS，丢掉与已保留框 IoU 超过阈值的框。
|
|
func suppressOverlaps(dets []faceDetection, iouThreshold float64) []faceDetection {
|
|
sort.SliceStable(dets, func(i, j int) bool { return dets[i].score > dets[j].score })
|
|
kept := make([]faceDetection, 0, len(dets))
|
|
for _, d := range dets {
|
|
overlaps := false
|
|
for _, k := range kept {
|
|
if rectIoU(d.rect, k.rect) > iouThreshold {
|
|
overlaps = true
|
|
break
|
|
}
|
|
}
|
|
if !overlaps {
|
|
kept = append(kept, d)
|
|
}
|
|
}
|
|
return kept
|
|
}
|
|
|
|
func rectIoU(a, b image.Rectangle) float64 {
|
|
inter := a.Intersect(b)
|
|
if inter.Empty() {
|
|
return 0
|
|
}
|
|
interArea := float64(inter.Dx() * inter.Dy())
|
|
return interArea / (float64(a.Dx()*a.Dy()+b.Dx()*b.Dy()) - interArea)
|
|
}
|
|
|
|
// applyFaceNotice 给图中每张人脸盖一层红色网格点，返回 PNG。
|
|
// 没检出人脸（或不是可解码的图片）时返回 ErrNoFaceDetected，调用方应继续用原图；
|
|
// 其它错误说明检测器不可用，调用方不应把未打码的图上传。
|
|
func applyFaceNotice(b []byte) ([]byte, error) {
|
|
src, _, err := image.Decode(bytes.NewReader(b))
|
|
if err != nil {
|
|
return nil, ErrNoFaceDetected
|
|
}
|
|
boxes, err := detectFaces(src)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if len(boxes) == 0 {
|
|
return nil, ErrNoFaceDetected
|
|
}
|
|
|
|
w, h := src.Bounds().Dx(), src.Bounds().Dy()
|
|
orig := image.NewRGBA(image.Rect(0, 0, w, h))
|
|
draw.Draw(orig, orig.Bounds(), src, src.Bounds().Min, draw.Src)
|
|
|
|
offset := image.Pt(-src.Bounds().Min.X, -src.Bounds().Min.Y)
|
|
|
|
// 脸多（>4 张）时不做顶部条，直接给每张脸盖一层白点阵。
|
|
if len(boxes) > 4 {
|
|
dst := image.NewRGBA(image.Rect(0, 0, w, h))
|
|
draw.Draw(dst, dst.Bounds(), orig, image.Point{}, draw.Src)
|
|
for _, box := range boxes {
|
|
drawWhiteDots(dst, box.Add(offset).Intersect(dst.Bounds()))
|
|
}
|
|
var out bytes.Buffer
|
|
if err := png.Encode(&out, dst); err != nil {
|
|
return nil, err
|
|
}
|
|
return out.Bytes(), nil
|
|
}
|
|
|
|
// 紧贴脸的检出框：切片和盖黑都用它，两者一样大。
|
|
heads := make([]image.Rectangle, 0, len(boxes))
|
|
for _, box := range boxes {
|
|
heads = append(heads, box.Add(offset).Intersect(orig.Bounds()))
|
|
}
|
|
|
|
// 取面积最大的一张脸切成 4 块——切割线正好穿过脸中心，每块只有 1/4 张脸，
|
|
// 打乱后分开摆到图片上方新加的黑色条带里。
|
|
main := largestRect(heads)
|
|
quads := faceQuadrants(orig, main) // 已按打乱顺序排列的 4 块
|
|
gap := main.Dx() / 8
|
|
if gap < 8 {
|
|
gap = 8
|
|
}
|
|
stripH := main.Dy()/2 + 2*gap
|
|
|
|
black := image.NewUniform(color.RGBA{A: 255})
|
|
dst := image.NewRGBA(image.Rect(0, 0, w, h+stripH))
|
|
draw.Draw(dst, dst.Bounds(), black, image.Point{}, draw.Src) // 条带底色黑
|
|
draw.Draw(dst, image.Rect(0, stripH, w, h+stripH), orig, image.Point{}, draw.Src)
|
|
|
|
// 4 块横向排开，块间留黑缝，明显是碎片而非整脸。
|
|
x := gap
|
|
for _, q := range quads {
|
|
qw, qh := q.Bounds().Dx(), q.Bounds().Dy()
|
|
if x+qw > w {
|
|
break
|
|
}
|
|
draw.Draw(dst, image.Rect(x, gap, x+qw, gap+qh), q, image.Point{}, draw.Src)
|
|
x += qw + gap
|
|
}
|
|
|
|
// 原图里每张脸整块盖黑（坐标下移 stripH）。
|
|
for _, hd := range heads {
|
|
draw.Draw(dst, hd.Add(image.Pt(0, stripH)), black, image.Point{}, draw.Src)
|
|
}
|
|
|
|
var out bytes.Buffer
|
|
if err := png.Encode(&out, dst); err != nil {
|
|
return nil, err
|
|
}
|
|
return out.Bytes(), nil
|
|
}
|
|
|
|
// expandFaceBox 把 YuNet 的紧致人脸框向外扩到大致一个头部的范围（额头到下巴），
|
|
// 只挡眼睛不足以让 Adobe 认不出大图正脸。结果裁剪到图像范围内。
|
|
func expandFaceBox(b image.Rectangle, bounds image.Rectangle) image.Rectangle {
|
|
w := b.Dx()
|
|
h := b.Dy()
|
|
padX := w * 60 / 100
|
|
minX := b.Min.X - padX
|
|
maxX := b.Max.X + padX
|
|
width := maxX - minX
|
|
// 高度按脸宽推算（额头到下巴 ≈ 1.2×脸宽）。YuNet 框常只圈住眼部，
|
|
// 只按框高向下扩不足以盖住嘴和下巴。
|
|
height := width * 115 / 100
|
|
if hh := h * 180 / 100; hh > height {
|
|
height = hh
|
|
}
|
|
top := b.Min.Y - h*85/100
|
|
out := image.Rect(minX, top, maxX, top+height)
|
|
return out.Intersect(bounds)
|
|
}
|
|
|
|
// faceQuadrants 把 box 区域切成 2×2 四块，按固定置换打乱顺序后返回这 4 块子图。
|
|
func faceQuadrants(src *image.RGBA, box image.Rectangle) []*image.RGBA {
|
|
hw, hh := box.Dx()/2, box.Dy()/2
|
|
rects := [4]image.Rectangle{
|
|
image.Rect(box.Min.X, box.Min.Y, box.Min.X+hw, box.Min.Y+hh),
|
|
image.Rect(box.Min.X+hw, box.Min.Y, box.Max.X, box.Min.Y+hh),
|
|
image.Rect(box.Min.X, box.Min.Y+hh, box.Min.X+hw, box.Max.Y),
|
|
image.Rect(box.Min.X+hw, box.Min.Y+hh, box.Max.X, box.Max.Y),
|
|
}
|
|
perm := [4]int{3, 1, 2, 0} // 打乱顺序
|
|
rot := [4]int{1, 2, 3, 2} // 每块各自旋转 90°/180°/270°，打断人脸连续性
|
|
out := make([]*image.RGBA, 0, 4)
|
|
for i, s := range perm {
|
|
r := rects[s]
|
|
q := image.NewRGBA(image.Rect(0, 0, r.Dx(), r.Dy()))
|
|
draw.Draw(q, q.Bounds(), src, r.Min, draw.Src)
|
|
out = append(out, rotateRGBA(q, rot[i]))
|
|
}
|
|
return out
|
|
}
|
|
|
|
// rotateRGBA 顺时针旋转 90° 的 k 倍，返回旋转后的新图。
|
|
func rotateRGBA(src *image.RGBA, k int) *image.RGBA {
|
|
k = ((k % 4) + 4) % 4
|
|
if k == 0 {
|
|
return src
|
|
}
|
|
w, h := src.Bounds().Dx(), src.Bounds().Dy()
|
|
var dst *image.RGBA
|
|
if k == 2 {
|
|
dst = image.NewRGBA(image.Rect(0, 0, w, h))
|
|
for y := 0; y < h; y++ {
|
|
for x := 0; x < w; x++ {
|
|
dst.SetRGBA(w-1-x, h-1-y, src.RGBAAt(x, y))
|
|
}
|
|
}
|
|
return dst
|
|
}
|
|
dst = image.NewRGBA(image.Rect(0, 0, h, w))
|
|
for y := 0; y < h; y++ {
|
|
for x := 0; x < w; x++ {
|
|
if k == 1 {
|
|
dst.SetRGBA(h-1-y, x, src.RGBAAt(x, y))
|
|
} else {
|
|
dst.SetRGBA(y, w-1-x, src.RGBAAt(x, y))
|
|
}
|
|
}
|
|
}
|
|
return dst
|
|
}
|
|
|
|
// drawWhiteDots 在脸框内铺一层小白点阵，盖住五官（只盖脸、不外扩）。
|
|
func drawWhiteDots(dst *image.RGBA, box image.Rectangle) {
|
|
step := box.Dx() / 16
|
|
if step < 4 {
|
|
step = 4
|
|
}
|
|
r := step / 3
|
|
if r < 2 {
|
|
r = 2
|
|
}
|
|
white := color.RGBA{R: 255, G: 255, B: 255, A: 255}
|
|
for cy := box.Min.Y + step/2; cy < box.Max.Y; cy += step {
|
|
for cx := box.Min.X + step/2; cx < box.Max.X; cx += step {
|
|
fillDot(dst, cx, cy, r, box, white)
|
|
}
|
|
}
|
|
}
|
|
|
|
// largestRect 返回面积最大的矩形。
|
|
func largestRect(rs []image.Rectangle) image.Rectangle {
|
|
best := rs[0]
|
|
for _, r := range rs[1:] {
|
|
if r.Dx()*r.Dy() > best.Dx()*best.Dy() {
|
|
best = r
|
|
}
|
|
}
|
|
return best
|
|
}
|
|
|
|
// minInt 返回两数中较小者。
|
|
func minInt(a, b int) int {
|
|
if a < b {
|
|
return a
|
|
}
|
|
return b
|
|
}
|
|
|
|
// faceDotAlpha 网格点的不透明度（255 = 实心）。
|
|
const faceDotAlpha uint8 = 26
|
|
|
|
// blendPixel 把前景色按其 alpha 叠到底色上。
|
|
func blendPixel(bg, fg color.RGBA) color.RGBA {
|
|
a := int(fg.A)
|
|
mix := func(f, b uint8) uint8 {
|
|
return uint8((int(f)*a + int(b)*(255-a)) / 255)
|
|
}
|
|
return color.RGBA{R: mix(fg.R, bg.R), G: mix(fg.G, bg.G), B: mix(fg.B, bg.B), A: 255}
|
|
}
|
|
|
|
// fillDot 以 (cx,cy) 为心画一个半径 r 的圆，按 c.A 与底图混合，裁剪在 clip 内。
|
|
func fillDot(dst *image.RGBA, cx, cy, r int, clip image.Rectangle, c color.RGBA) {
|
|
r2 := r * r
|
|
for y := cy - r; y <= cy+r; y++ {
|
|
if y < clip.Min.Y || y >= clip.Max.Y {
|
|
continue
|
|
}
|
|
for x := cx - r; x <= cx+r; x++ {
|
|
if x < clip.Min.X || x >= clip.Max.X {
|
|
continue
|
|
}
|
|
dx := x - cx
|
|
dy := y - cy
|
|
if dx*dx+dy*dy <= r2 {
|
|
dst.SetRGBA(x, y, blendPixel(dst.RGBAAt(x, y), c))
|
|
}
|
|
}
|
|
}
|
|
}
|