diff --git a/README.md b/README.md index 65103d7..7db3eed 100644 --- a/README.md +++ b/README.md @@ -58,7 +58,7 @@ This repository contains a suite of face AI models designed for various applicat ### ArcFace for Feature Embedding -- Model Name: [arcface](model/arcface/arcface.go) +- Model Name: [arcface_w600k_r50](model/arcface/arcface.go) - Description: Generates feature embeddings for faces, useful for identity verification and facial recognition tasks. - Download Link: [Download ArcFace Model](https://github.com/facefusion/facefusion-assets/releases/download/models/arcface_w600k_r50.onnx) @@ -69,6 +69,14 @@ This repository contains a suite of face AI models designed for various applicat | | 0.29 | 0.00 | 0.45 | | | 0.48 | 0.45 | 0.00 | +### Face Occluder Detection +- Model Name: [face_occluder](model/faceoccluder/faceoccluder.go) +- Description: Detects parts of a face that are not occluded by objects, providing insights into visible facial features. +- Download Link: [Download Face Occluder Model](https://github.com/facefusion/facefusion-assets/releases/download/models/face_occluder.onnx) + + + + ### Gender and Age Estimation - Model Name: [gender_age](model/genderage/genderage.go) - Description: Detects gender and estimates the age of detected faces. diff --git a/docs/face_occluder_1.jpg b/docs/face_occluder_1.jpg new file mode 100644 index 0000000..07c7408 Binary files /dev/null and b/docs/face_occluder_1.jpg differ diff --git a/docs/face_occluder_2.jpg b/docs/face_occluder_2.jpg new file mode 100644 index 0000000..459d0a9 Binary files /dev/null and b/docs/face_occluder_2.jpg differ diff --git a/model/faceoccluder/faceoccluder.go b/model/faceoccluder/faceoccluder.go new file mode 100644 index 0000000..a0c2aee --- /dev/null +++ b/model/faceoccluder/faceoccluder.go @@ -0,0 +1,46 @@ +package faceoccluder + +import ( + "github.com/dev6699/face/model" + "gocv.io/x/gocv" +) + +type Model struct { + cropSize model.Size + cropVisionFrame gocv.Mat + affineMatrix gocv.Mat + boxMask gocv.Mat +} + +type Input struct { + Img gocv.Mat + FaceLandmark5 []gocv.Point2f +} + +type Output struct { + CropVisionFrame gocv.Mat + AffineMatrix gocv.Mat + CropMask gocv.Mat +} + +type TModel = model.Model[*Input, *Output] + +var _ TModel = &Model{} + +func NewFactory() func() TModel { + return func() TModel { + return New() + } +} + +func New() *Model { + return &Model{} +} + +func (m *Model) ModelName() string { + return "face_occluder" +} + +func (m *Model) ModelVersion() string { + return "1" +} diff --git a/model/faceoccluder/post.go b/model/faceoccluder/post.go new file mode 100644 index 0000000..7f76746 --- /dev/null +++ b/model/faceoccluder/post.go @@ -0,0 +1,85 @@ +package faceoccluder + +import ( + "image" + + "github.com/dev6699/face/model" + "gocv.io/x/gocv" +) + +func (m *Model) PostProcess(rawOutputContents [][]byte) (*Output, error) { + // "outputs": [ + // { + // "name": "out_mask:0", + // "datatype": "FP32", + // "shape": [ + // -1, + // 256, + // 256, + // 1 + // ] + // } + // ] + + outMask, err := model.BytesToFloat32Slice(rawOutputContents[0]) + if err != nil { + return nil, err + } + + rows := 256 + cols := 256 + maskMat := gocv.NewMatWithSize(rows, cols, gocv.MatTypeCV32F) + for i := 0; i < rows; i++ { + for j := 0; j < cols; j++ { + idx := i*cols + j + maskMat.SetFloatAt(i, j, outMask[idx]) + } + } + + model.ClipMat(maskMat, 0, 1) + gocv.Resize(maskMat, &maskMat, image.Point{X: m.cropSize.Width, Y: m.cropSize.Height}, 0, 0, gocv.InterpolationDefault) + gocv.GaussianBlur(maskMat, &maskMat, image.Point{0, 0}, 5.0, 0, gocv.BorderDefault) + + model.ClipMat(maskMat, 0.5, 1) + model.MatSubtract(maskMat, 0.5) + maskMat.MultiplyFloat(2) + cropMask := reduceMinimum([]gocv.Mat{m.boxMask, maskMat}) + model.ClipMat(cropMask, 0, 1) + + defer m.boxMask.Close() + defer maskMat.Close() + + return &Output{ + CropVisionFrame: m.cropVisionFrame, + AffineMatrix: m.affineMatrix, + CropMask: cropMask, + }, nil +} + +// reduceMinimum finds the element-wise minimum of a list of gocv.Mat +func reduceMinimum(mats []gocv.Mat) gocv.Mat { + if len(mats) == 0 { + return gocv.NewMat() + } + + // Start with the first matrix as the initial minimum + minMat := mats[0].Clone() + rows, cols := minMat.Rows(), minMat.Cols() + + // Iterate over the remaining matrices + for i := 1; i < len(mats); i++ { + for row := 0; row < rows; row++ { + for col := 0; col < cols; col++ { + currentMin := minMat.GetFloatAt(row, col) + newValue := mats[i].GetFloatAt(row, col) + + // Update the minimum value + if newValue < currentMin { + minMat.SetFloatAt(row, col, newValue) + } + } + } + } + + return minMat +} diff --git a/model/faceoccluder/pre.go b/model/faceoccluder/pre.go new file mode 100644 index 0000000..14dab24 --- /dev/null +++ b/model/faceoccluder/pre.go @@ -0,0 +1,81 @@ +package faceoccluder + +import ( + "image" + "math" + + "github.com/dev6699/face/model" + "github.com/dev6699/face/protobuf" + "gocv.io/x/gocv" +) + +func (m *Model) PreProcess(i *Input) ([]*protobuf.InferTensorContents, error) { + cropSize := model.Size{Width: 128, Height: 128} + m.cropSize = cropSize + cropVisionFrame, affineMatrix := model.WarpFaceByFaceLandmark5(i.Img, i.FaceLandmark5, arcface_128_v2, cropSize) + m.cropVisionFrame = cropVisionFrame + m.affineMatrix = affineMatrix + + boxMask := createStaticBoxMask(model.Size{Width: 128, Height: 128}, 0.3, Padding{Top: 0, Right: 0, Bottom: 0, Left: 0}) + m.boxMask = boxMask + + resizedFrame := gocv.NewMat() + defer resizedFrame.Close() + gocv.Resize(cropVisionFrame, &resizedFrame, image.Point{X: 256, Y: 256}, 0, 0, gocv.InterpolationDefault) + data, _ := resizedFrame.DataPtrUint8() + + d := make([]float32, len(data)) + for i, a := range data { + d[i] = float32(a) / 255.0 + } + + contents := &protobuf.InferTensorContents{ + Fp32Contents: d, + } + return []*protobuf.InferTensorContents{contents}, nil +} + +// Padding represents the padding values for the mask. +type Padding struct { + Top, Right, Bottom, Left float64 +} + +// Create a static box mask with specified size, blur, and padding. +func createStaticBoxMask(cropSize model.Size, faceMaskBlur float64, faceMaskPadding Padding) gocv.Mat { + blurAmount := int(float64(cropSize.Width) * 0.5 * faceMaskBlur) + blurArea := int(math.Max(float64(blurAmount/2), 1)) + + // Create a box mask initialized to ones. + boxMask := gocv.NewMatWithSize(cropSize.Height, cropSize.Width, gocv.MatTypeCV32F) + boxMask.SetTo(gocv.NewScalar(1.0, 1.0, 1.0, 1.0)) // Fill the entire matrix with ones. + + // Calculate padding values. + padTop := int(math.Max(float64(blurArea), float64(cropSize.Height)*faceMaskPadding.Top/100)) + padBottom := int(math.Max(float64(blurArea), float64(cropSize.Height)*faceMaskPadding.Bottom/100)) + padLeft := int(math.Max(float64(blurArea), float64(cropSize.Width)*faceMaskPadding.Left/100)) + padRight := int(math.Max(float64(blurArea), float64(cropSize.Width)*faceMaskPadding.Right/100)) + + // Set padding areas to zero. + topRegion := boxMask.Region(image.Rect(0, 0, cropSize.Width, padTop)) + defer topRegion.Close() + topRegion.SetTo(gocv.NewScalar(0, 0, 0, 0)) + + bottomRegion := boxMask.Region(image.Rect(0, cropSize.Height-padBottom, cropSize.Width, cropSize.Height)) + defer bottomRegion.Close() + bottomRegion.SetTo(gocv.NewScalar(0, 0, 0, 0)) + + leftRegion := boxMask.Region(image.Rect(0, 0, padLeft, cropSize.Height)) + defer leftRegion.Close() + leftRegion.SetTo(gocv.NewScalar(0, 0, 0, 0)) + + rightRegion := boxMask.Region(image.Rect(cropSize.Width-padRight, 0, cropSize.Width, cropSize.Height)) + defer rightRegion.Close() + rightRegion.SetTo(gocv.NewScalar(0, 0, 0, 0)) + + // Apply Gaussian blur if required. + if blurAmount > 0 { + gocv.GaussianBlur(boxMask, &boxMask, image.Point{0, 0}, float64(blurAmount)*0.25, 0, gocv.BorderDefault) + } + + return boxMask +} diff --git a/model/faceoccluder/template.go b/model/faceoccluder/template.go new file mode 100644 index 0000000..c71bb8a --- /dev/null +++ b/model/faceoccluder/template.go @@ -0,0 +1,11 @@ +package faceoccluder + +import "gocv.io/x/gocv" + +var arcface_128_v2 = []gocv.Point2f{ + {X: 0.36167656, Y: 0.40387734}, + {X: 0.63696719, Y: 0.40235469}, + {X: 0.50019687, Y: 0.56044219}, + {X: 0.38710391, Y: 0.72160547}, + {X: 0.61507734, Y: 0.72034453}, +} diff --git a/model/util.go b/model/util.go index 870025d..bb9f79e 100644 --- a/model/util.go +++ b/model/util.go @@ -4,6 +4,7 @@ import ( "bytes" "encoding/binary" "image" + "image/color" "io" "math" @@ -88,3 +89,77 @@ func CalculateMean2f(points []gocv.Point2f) gocv.Point2f { return gocv.Point2f{X: meanX, Y: meanY} } + +// Size represents the width and height dimensions. +type Size struct { + Width, Height int +} + +func WarpFaceByFaceLandmark5(visionFrame gocv.Mat, faceLandmark5 []gocv.Point2f, warpTemplate []gocv.Point2f, cropSize Size) (gocv.Mat, gocv.Mat) { + affineMatrix := estimateMatrixByFaceLandmark5(faceLandmark5, warpTemplate, cropSize) + cropVisionFrame := gocv.NewMat() + + gocv.WarpAffineWithParams( + visionFrame, + &cropVisionFrame, + affineMatrix, + image.Pt(cropSize.Width, cropSize.Height), + gocv.InterpolationArea, + gocv.BorderReplicate, + color.RGBA{}, + ) + + return cropVisionFrame, affineMatrix +} + +func estimateMatrixByFaceLandmark5(faceLandmark5 []gocv.Point2f, warpTemplate []gocv.Point2f, cropSize Size) gocv.Mat { + normedWarpTemplate := normalizeWarpTemplate(warpTemplate, cropSize) + pvsrc := gocv.NewPoint2fVectorFromPoints(faceLandmark5) + pvdst := gocv.NewPoint2fVectorFromPoints(normedWarpTemplate) + inliers := gocv.NewMat() + defer inliers.Close() + method := 8 + ransacProjThreshold := 100.0 + maxiters := uint(2000) + confidence := 0.99 + refineIters := uint(10) + // https://docs.opencv.org/4.x/d9/d0c/group__calib3d.html#gad767faff73e9cbd8b9d92b955b50062d + affineMatrix := gocv.EstimateAffinePartial2DWithParams(pvsrc, pvdst, inliers, method, ransacProjThreshold, maxiters, confidence, refineIters) + return affineMatrix +} + +// normalizeWarpTemplate scales the warp template according to the crop size. +func normalizeWarpTemplate(warpTemplate []gocv.Point2f, cropSize Size) []gocv.Point2f { + normedWarpTemplate := make([]gocv.Point2f, len(warpTemplate)) + for i, pt := range warpTemplate { + normedWarpTemplate[i] = gocv.Point2f{ + X: pt.X * float32(cropSize.Width), + Y: pt.Y * float32(cropSize.Height), + } + } + return normedWarpTemplate +} + +// MatSubtract subtract value v fom mat +func MatSubtract(mat gocv.Mat, v float64) { + constantValue := float64(v) + constantMat := gocv.NewMatWithSizeFromScalar(gocv.NewScalar(constantValue, constantValue, constantValue, 0), mat.Rows(), mat.Cols(), mat.Type()) + defer constantMat.Close() + gocv.Subtract(mat, constantMat, &mat) +} + +// ClipMat clips the values of a gocv.Mat within a specified range. +// mat need to be float32 type +func ClipMat(mat gocv.Mat, minVal, maxVal float32) { + for row := 0; row < mat.Rows(); row++ { + for col := 0; col < mat.Cols(); col++ { + value := mat.GetFloatAt(row, col) + if value < minVal { + value = minVal + } else if value > maxVal { + value = maxVal + } + mat.SetFloatAt(row, col, value) + } + } +} diff --git a/model_repository/arcface_w600k_r50/config.pbtxt b/model_repository/arcface_w600k_r50/config.pbtxt new file mode 100644 index 0000000..c2b7103 --- /dev/null +++ b/model_repository/arcface_w600k_r50/config.pbtxt @@ -0,0 +1,2 @@ +name: "arcface_w600k_r50" +platform: "onnxruntime_onnx" \ No newline at end of file diff --git a/model_repository/face_occluder/config.pbtxt b/model_repository/face_occluder/config.pbtxt new file mode 100644 index 0000000..d4c75c3 --- /dev/null +++ b/model_repository/face_occluder/config.pbtxt @@ -0,0 +1,2 @@ +name: "face_occluder" +platform: "onnxruntime_onnx" \ No newline at end of file