diff --git a/README.md b/README.md index 593fc78..d9023e2 100644 --- a/README.md +++ b/README.md @@ -69,6 +69,16 @@ 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 Enhancer with GFPGAN +- Model Name: [gfpgan_1.4](model/gfpgan/gfpgan.go) +- Description: Enhances facial features and improves image quality, often used for face restoration and super-resolution tasks. +- Download Link: [Download GFPGAN Model](https://github.com/facefusion/facefusion-assets/releases/download/models/gfpgan_1.4.onnx) + + | Input | Output | + | :---: | :----: | + | | | + + ### 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. diff --git a/docs/gfpgan_1.jpg b/docs/gfpgan_1.jpg new file mode 100644 index 0000000..9e3176e Binary files /dev/null and b/docs/gfpgan_1.jpg differ diff --git a/docs/gfpgan_2.jpg b/docs/gfpgan_2.jpg new file mode 100644 index 0000000..adc8ab7 Binary files /dev/null and b/docs/gfpgan_2.jpg differ diff --git a/model/faceoccluder/post.go b/model/faceoccluder/post.go index 7f76746..5033b87 100644 --- a/model/faceoccluder/post.go +++ b/model/faceoccluder/post.go @@ -43,7 +43,7 @@ func (m *Model) PostProcess(rawOutputContents [][]byte) (*Output, error) { model.ClipMat(maskMat, 0.5, 1) model.MatSubtract(maskMat, 0.5) maskMat.MultiplyFloat(2) - cropMask := reduceMinimum([]gocv.Mat{m.boxMask, maskMat}) + cropMask := model.ReduceMinimum([]gocv.Mat{m.boxMask, maskMat}) model.ClipMat(cropMask, 0, 1) defer m.boxMask.Close() @@ -55,31 +55,3 @@ func (m *Model) PostProcess(rawOutputContents [][]byte) (*Output, error) { 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 index 14dab24..eff54e8 100644 --- a/model/faceoccluder/pre.go +++ b/model/faceoccluder/pre.go @@ -2,7 +2,6 @@ package faceoccluder import ( "image" - "math" "github.com/dev6699/face/model" "github.com/dev6699/face/protobuf" @@ -16,7 +15,7 @@ func (m *Model) PreProcess(i *Input) ([]*protobuf.InferTensorContents, error) { 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}) + boxMask := model.CreateStaticBoxMask(cropSize, 0.3, model.Padding{Top: 0, Right: 0, Bottom: 0, Left: 0}) m.boxMask = boxMask resizedFrame := gocv.NewMat() @@ -34,48 +33,3 @@ func (m *Model) PreProcess(i *Input) ([]*protobuf.InferTensorContents, error) { } 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/gfpgan/gfpgan.go b/model/gfpgan/gfpgan.go new file mode 100644 index 0000000..13a40e9 --- /dev/null +++ b/model/gfpgan/gfpgan.go @@ -0,0 +1,47 @@ +package gfpgan + +import ( + "github.com/dev6699/face/model" + "gocv.io/x/gocv" +) + +type Model struct { + faceEnhancerBlend float64 + + img gocv.Mat + cropSize model.Size + affineMatrix gocv.Mat +} + +type Input struct { + Img gocv.Mat + FaceLandmark5 []gocv.Point2f +} + +type Output struct { + OutFrame gocv.Mat +} + +type ModelT = model.Model[*Input, *Output] + +var _ ModelT = &Model{} + +func NewFactory(faceEnhancerBlend float64) func() ModelT { + return func() ModelT { + return New(faceEnhancerBlend) + } +} + +func New(faceEnhancerBlend float64) *Model { + return &Model{ + faceEnhancerBlend: faceEnhancerBlend, + } +} + +func (m *Model) ModelName() string { + return "gfpgan_1.4" +} + +func (m *Model) ModelVersion() string { + return "1" +} diff --git a/model/gfpgan/post.go b/model/gfpgan/post.go new file mode 100644 index 0000000..cd4e905 --- /dev/null +++ b/model/gfpgan/post.go @@ -0,0 +1,81 @@ +package gfpgan + +import ( + "math" + + "github.com/dev6699/face/model" + "gocv.io/x/gocv" +) + +func (m *Model) PostProcess(rawOutputContents [][]byte) (*Output, error) { + // "outputs": [ + // { + // "name": "output", + // "datatype": "FP32", + // "shape": [ + // 1, + // 3, + // 512, + // 512 + // ] + // } + // ] + output, err := model.BytesToFloat32Slice(rawOutputContents[0]) + if err != nil { + return nil, err + } + + d := make([]uint8, len(output)) + j := 0 + for i := 0; i < len(output); i += 3 { + d[i+2] = uint8(math.Round(255.0 * float64(clip(output[j], -1, 1)+1.0) / 2.0)) + d[i+1] = uint8(math.Round(255.0 * float64(clip(output[len(d)/3+j], -1, 1)+1.0) / 2.0)) + d[i] = uint8(math.Round(255.0 * float64(clip(output[len(d)/3*2+j], -1, 1)+1.0) / 2.0)) + j++ + } + + width := 512 + height := 512 + imgType := gocv.MatTypeCV8UC3 + + mat, err := gocv.NewMatFromBytes(height, width, imgType, d) + if err != nil { + return nil, err + } + defer mat.Close() + + boxMask := model.CreateStaticBoxMask(m.cropSize, 0.3, model.Padding{Top: 0, Right: 0, Bottom: 0, Left: 0}) + defer boxMask.Close() + + cropMask := model.ReduceMinimum([]gocv.Mat{boxMask}) + defer cropMask.Close() + model.ClipMat(cropMask, 0, 1) + + outMat := model.PasteBack(m.img, mat, cropMask, m.affineMatrix) + defer outMat.Close() + defer m.affineMatrix.Close() + + return &Output{ + OutFrame: m.blendFrame(m.img, outMat), + }, nil +} + +// blendFrame blends two frame (gocv.Mat) images with a specified blending factor. +func (m *Model) blendFrame(tempVisionFrame, pasteVisionFrame gocv.Mat) gocv.Mat { + faceEnhancerBlendRatio := 1.0 - (m.faceEnhancerBlend / 100.0) + outputFrame := gocv.NewMat() + gocv.AddWeighted(tempVisionFrame, faceEnhancerBlendRatio, pasteVisionFrame, 1-faceEnhancerBlendRatio, 0, &outputFrame) + return outputFrame +} + +func clip(v float32, min float32, max float32) float32 { + if v < min { + return min + } + + if v > max { + return max + } + + return v +} diff --git a/model/gfpgan/pre.go b/model/gfpgan/pre.go new file mode 100644 index 0000000..94f0547 --- /dev/null +++ b/model/gfpgan/pre.go @@ -0,0 +1,44 @@ +package gfpgan + +import ( + "github.com/dev6699/face/model" + "github.com/dev6699/face/protobuf" + "gocv.io/x/gocv" +) + +func (m *Model) PreProcess(i *Input) ([]*protobuf.InferTensorContents, error) { + m.img = i.Img + cropSize := model.Size{Width: 512, Height: 512} + m.cropSize = cropSize + + cropVisionFrame, affineMatrix := model.WarpFaceByFaceLandmark5(i.Img, i.FaceLandmark5, ffhq_512, cropSize) + m.affineMatrix = affineMatrix + defer cropVisionFrame.Close() + + d := []float32{} + cropVisionFrame.ConvertTo(&cropVisionFrame, gocv.MatTypeCV32F) + cropVisionFrame.DivideFloat(255.0) + model.MatSubtract(cropVisionFrame, 0.5) + cropVisionFrame.DivideFloat(0.5) + + rgbChannels := gocv.Split(cropVisionFrame) + b := rgbChannels[2] + defer b.Close() + bd, _ := b.DataPtrFloat32() + d = append(d, bd...) + + g := rgbChannels[1] + defer g.Close() + gd, _ := g.DataPtrFloat32() + d = append(d, gd...) + + r := rgbChannels[0] + defer r.Close() + rd, _ := r.DataPtrFloat32() + d = append(d, rd...) + + contents := &protobuf.InferTensorContents{ + Fp32Contents: d, + } + return []*protobuf.InferTensorContents{contents}, nil +} diff --git a/model/gfpgan/template.go b/model/gfpgan/template.go new file mode 100644 index 0000000..6e211f2 --- /dev/null +++ b/model/gfpgan/template.go @@ -0,0 +1,11 @@ +package gfpgan + +import "gocv.io/x/gocv" + +var ffhq_512 = []gocv.Point2f{ + {X: 0.37691676, Y: 0.46864664}, + {X: 0.62285697, Y: 0.46912813}, + {X: 0.50123859, Y: 0.61331904}, + {X: 0.39308822, Y: 0.72541100}, + {X: 0.61150205, Y: 0.72490465}, +} diff --git a/model/util.go b/model/util.go index bb9f79e..0501db3 100644 --- a/model/util.go +++ b/model/util.go @@ -163,3 +163,127 @@ func ClipMat(mat gocv.Mat, minVal, maxVal float32) { } } } + +// Padding represents the padding values for the mask. +type Padding struct { + Top, Right, Bottom, Left float64 +} + +// CreateStaticBoxMask create a static box mask with specified size, blur, and padding. +func CreateStaticBoxMask(cropSize 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 +} + +// 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 +} + +// PasteBack paste cropVisionFrame back to targetVisionFrame +func PasteBack(targetVisionFrame gocv.Mat, cropVisionFrame gocv.Mat, cropMask gocv.Mat, affineMatrix gocv.Mat) gocv.Mat { + inverseMatrix := InvertAffineMatrix(affineMatrix) + defer inverseMatrix.Close() + tempSize := image.Pt(targetVisionFrame.Cols(), targetVisionFrame.Rows()) + cropVisionFrame.ConvertTo(&cropVisionFrame, gocv.MatTypeCV64F) + inverseVisionFrame := getInverseVisionFrame(cropVisionFrame, inverseMatrix, tempSize) + defer inverseVisionFrame.Close() + + inverseMask := getInverseMask(inverseMatrix, cropMask, tempSize) + defer inverseMask.Close() + inverseMaskData, _ := inverseMask.DataPtrFloat64() + inverseVisionFrameData, _ := inverseVisionFrame.DataPtrFloat64() + + data, _ := targetVisionFrame.DataPtrUint8() + d := make([]uint8, len(data)) + j := 0 + for i := 0; i < len(data); i += 3 { + inverseM := inverseMaskData[j] + d[i] = uint8(inverseM*inverseVisionFrameData[i] + (1-inverseM)*float64(data[i])) + d[i+1] = uint8(inverseM*inverseVisionFrameData[i+1] + (1-inverseM)*float64(data[i+1])) + d[i+2] = uint8(inverseM*inverseVisionFrameData[i+2] + (1-inverseM)*float64(data[i+2])) + j++ + } + + mat, _ := gocv.NewMatFromBytes(targetVisionFrame.Rows(), targetVisionFrame.Cols(), gocv.MatTypeCV8UC3, d) + return mat +} + +func getInverseMask(inverseMatrix gocv.Mat, cropMask gocv.Mat, tempSize image.Point) gocv.Mat { + inverseMask := gocv.NewMat() + gocv.WarpAffine( + cropMask, + &inverseMask, + inverseMatrix, + tempSize, + ) + inverseMask.ConvertTo(&inverseMask, gocv.MatTypeCV32F) + ClipMat(inverseMask, 0, 1) + inverseMask.ConvertTo(&inverseMask, gocv.MatTypeCV64F) + return inverseMask +} + +func getInverseVisionFrame(cropVisionFrame gocv.Mat, inverseMatrix gocv.Mat, tempSize image.Point) gocv.Mat { + inverseVisionFrame := gocv.NewMat() + cropVisionFrame.ConvertTo(&cropVisionFrame, gocv.MatTypeCV64F) + gocv.WarpAffineWithParams(cropVisionFrame, &inverseVisionFrame, inverseMatrix, tempSize, gocv.InterpolationLinear, gocv.BorderReplicate, color.RGBA{}) + inverseVisionFrame.ConvertTo(&inverseVisionFrame, gocv.MatTypeCV64F) + return inverseVisionFrame +} diff --git a/model_repository/gfpgan_1.4/config.pbtxt b/model_repository/gfpgan_1.4/config.pbtxt new file mode 100644 index 0000000..14454ea --- /dev/null +++ b/model_repository/gfpgan_1.4/config.pbtxt @@ -0,0 +1,2 @@ +name: "gfpgan_1.4" +platform: "onnxruntime_onnx" \ No newline at end of file