diff --git a/README.md b/README.md index 21563c0..59c75e9 100644 --- a/README.md +++ b/README.md @@ -33,6 +33,7 @@ cmake .. # optional -DNCNN_VULKAN=OFF -DCMAKE_CXX_COMPILER=clang++ -DCMAKE_C_COM - scrfd [Google Drive](https://drive.google.com/drive/folders/1XPjfsuXGj9rXqAmo1K70BsqWmHvoYQv_?usp=sharing) - tracker (for face IOU calculation bettween frames) - hopenet (for head pose detection) [Google Drive](https://drive.google.com/drive/folders/1zLam-8s9ZMPDUxUEtNU2F9yFTDRM5fk-?usp=sharing) + - hair (for hair segmentation) [Google Drive](https://drive.google.com/drive/folders/14DOBaFrxTL1k4T1ved5qfRUUziurItT8?usp=sharing) - pose - detector (for pose detection/estimation) - ultralight [Google Drive](https://drive.google.com/drive/folders/15b-I5HDyGe2WLb-TO85SJYmnYONvGOKh?usp=sharing) diff --git a/go/README.md b/go/README.md index b593adc..ba72e51 100644 --- a/go/README.md +++ b/go/README.md @@ -45,6 +45,7 @@ make -j 4 - scrfd [Google Drive](https://drive.google.com/drive/folders/1XPjfsuXGj9rXqAmo1K70BsqWmHvoYQv_?usp=sharing) - tracker (for face IOU calculation bettween frames) - hopenet (for head pose detection) [Google Drive](https://drive.google.com/drive/folders/1zLam-8s9ZMPDUxUEtNU2F9yFTDRM5fk-?usp=sharing) + - hair (for hair segmentation) [Google Drive](https://drive.google.com/drive/folders/14DOBaFrxTL1k4T1ved5qfRUUziurItT8?usp=sharing) - pose - detector (for pose detection/estimation) - ultralight [Google Drive](https://drive.google.com/drive/folders/15b-I5HDyGe2WLb-TO85SJYmnYONvGOKh?usp=sharing) diff --git a/go/common/image.go b/go/common/image.go index 43e57b3..a37a1da 100644 --- a/go/common/image.go +++ b/go/common/image.go @@ -26,6 +26,9 @@ type Image struct { // NewImage returns a new Image func NewImage(img image.Image) *Image { buf := new(bytes.Buffer) + if img == nil { + return &Image{buffer: buf} + } Image2RGB(buf, img) return &Image{ Image: img, @@ -33,6 +36,14 @@ func NewImage(img image.Image) *Image { } } +// Write write bytes to buffer +func (i *Image) Write(b []byte) { + if i.buffer == nil { + return + } + i.buffer.Write(b) +} + // Bytes returns image bytes in rgb func (i Image) Bytes() []byte { if i.buffer == nil { @@ -74,20 +85,23 @@ func NewCImage() *C.Image { return ret } +// FreeCImage free C.Image func FreeCImage(c *C.Image) { C.FreeImage(c) C.free(unsafe.Pointer(c)) } -func GoImage(c *C.Image) (image.Image, error) { +// GoImage returns Image from C.Image +func GoImage(c *C.Image, out *Image) { w := int(c.width) h := int(c.height) channels := int(c.channels) data := C.GoBytes(unsafe.Pointer(c.data), C.int(w*h*channels)*C.sizeof_uchar) - return NewImageFromBytes(data, w, h, channels) + NewImageFromBytes(data, w, h, channels, out) } -func NewImageFromBytes(data []byte, w int, h int, channels int) (image.Image, error) { +// NewImageFromBytes returns Image by []byte +func NewImageFromBytes(data []byte, w int, h int, channels int, out *Image) { img := image.NewRGBA(image.Rect(0, 0, w, h)) for y := 0; y < h; y++ { for x := 0; x < w; x++ { @@ -97,9 +111,10 @@ func NewImageFromBytes(data []byte, w int, h int, channels int) (image.Image, er alpha = data[pos+3] } img.SetRGBA(x, y, color.RGBA{uint8(data[pos]), uint8(data[pos+1]), uint8(data[pos+2]), uint8(alpha)}) + out.Write([]byte{byte(data[pos]), byte(data[pos+1]), byte(data[pos+2]), byte(alpha)}) } } - return img, nil + out.Image = img } // Image2RGB write image rgbdata to buffer diff --git a/go/error.go b/go/error.go index b897658..4ce1551 100644 --- a/go/error.go +++ b/go/error.go @@ -56,6 +56,12 @@ var ( Message: "detect head pose failed", } } + HairMattingError = func(code int) Error { + return Error{ + Code: code, + Message: "hair matting failed", + } + } DetectHandError = func(code int) Error { return Error{ Code: code, diff --git a/go/examples/aligner/main.go b/go/examples/aligner/main.go index 9a3b7e2..f40d477 100644 --- a/go/examples/aligner/main.go +++ b/go/examples/aligner/main.go @@ -50,8 +50,9 @@ func align(d detecter.Detecter, a *aligner.Aligner, imgPath string, filename str if err != nil { log.Fatalln(err) } + aligned := common.NewImage(nil) for idx, face := range faces { - aligned, err := a.Align(img, face) + err := a.Align(img, face, aligned) if err != nil { log.Fatalln(err) } diff --git a/go/examples/hair/main.go b/go/examples/hair/main.go new file mode 100644 index 0000000..805eac8 --- /dev/null +++ b/go/examples/hair/main.go @@ -0,0 +1,94 @@ +package main + +import ( + "bytes" + "fmt" + "image" + "image/jpeg" + "log" + "os" + "os/user" + "path/filepath" + "strings" + + "github.com/bububa/openvision/go/common" + "github.com/bububa/openvision/go/face/hair" +) + +func main() { + wd, _ := os.Getwd() + dataPath := cleanPath(wd, "~/go/src/github.com/bububa/openvision/data") + imgPath := filepath.Join(dataPath, "./images") + modelPath := filepath.Join(dataPath, "./models") + common.CreateGPUInstance() + defer common.DestroyGPUInstance() + d := estimator(modelPath) + defer d.Destroy() + matting(d, imgPath, "hair1.jpg") +} + +func estimator(modelPath string) *hair.Hair { + modelPath = filepath.Join(modelPath, "hair") + d := hair.NewHair() + if err := d.LoadModel(modelPath); err != nil { + log.Fatalln(err) + } + return d +} + +func matting(d *hair.Hair, imgPath string, filename string) { + inPath := filepath.Join(imgPath, filename) + imgLoaded, err := loadImage(inPath) + if err != nil { + log.Fatalln("load image failed,", err) + } + img := common.NewImage(imgLoaded) + out := common.NewImage(nil) + if err := d.Matting(img, out); err != nil { + log.Fatalln(err) + } + outPath := filepath.Join(imgPath, "./results", fmt.Sprintf("hair-matting-%s", filename)) + + if err := saveImage(out, outPath); err != nil { + log.Fatalln(err) + } + +} + +func loadImage(filePath string) (image.Image, error) { + fn, err := os.Open(filePath) + if err != nil { + return nil, err + } + defer fn.Close() + img, _, err := image.Decode(fn) + if err != nil { + return nil, err + } + return img, nil +} + +func saveImage(img image.Image, filePath string) error { + buf := new(bytes.Buffer) + if err := jpeg.Encode(buf, img, nil); err != nil { + return err + } + fn, err := os.Create(filePath) + if err != nil { + return err + } + defer fn.Close() + fn.Write(buf.Bytes()) + return nil +} + +func cleanPath(wd string, path string) string { + usr, _ := user.Current() + dir := usr.HomeDir + if path == "~" { + return dir + } else if strings.HasPrefix(path, "~/") { + return filepath.Join(dir, path[2:]) + } + return filepath.Join(wd, path) +} diff --git a/go/examples/poseseg/main.go b/go/examples/poseseg/main.go index 24a28c3..77a7649 100644 --- a/go/examples/poseseg/main.go +++ b/go/examples/poseseg/main.go @@ -77,8 +77,8 @@ func videomatting(seg segmentor.Segmentor, imgPath string, filename string, idx log.Fatalln("load image failed,", err) } img := common.NewImage(imgLoaded) - out, err := seg.Matting(img) - if err != nil { + out := common.NewImage(nil) + if err := seg.Matting(img, out); err != nil { log.Fatalln(err) } outPath := filepath.Join(imgPath, "./results/videomatting", fmt.Sprintf("%d.jpeg", idx)) @@ -95,8 +95,8 @@ func matting(seg segmentor.Segmentor, imgPath string, filename string, idx int) log.Fatalln("load image failed,", err) } img := common.NewImage(imgLoaded) - out, err := seg.Matting(img) - if err != nil { + out := common.NewImage(nil) + if err := seg.Matting(img, out); err != nil { log.Fatalln(err) } outPath := filepath.Join(imgPath, "./results", fmt.Sprintf("poseseg-matting-%d-%s", idx, filename)) @@ -119,8 +119,8 @@ func merge(seg segmentor.Segmentor, imgPath string, filename string, bgFilename log.Fatalln("load bg image failed,", err) } bg := common.NewImage(bgLoaded) - out, err := seg.Merge(img, bg) - if err != nil { + out := common.NewImage(nil) + if err := seg.Merge(img, bg, out); err != nil { log.Fatalln(err) } outPath := filepath.Join(imgPath, "./results", fmt.Sprintf("poseseg-merge-%d-%s", idx, filename)) diff --git a/go/examples/styletransfer/main.go b/go/examples/styletransfer/main.go index f8fd553..9f9983e 100644 --- a/go/examples/styletransfer/main.go +++ b/go/examples/styletransfer/main.go @@ -49,15 +49,14 @@ func transform(transfer styletransfer.StyleTransfer, imgPath string, filename st log.Fatalln("load image failed,", err) } img := common.NewImage(imgLoaded) - out, err := transfer.Transform(img) - if err != nil { + out := common.NewImage(nil) + if err := transfer.Transform(img, out); err != nil { log.Fatalln(err) } outPath := filepath.Join(imgPath, "./results", fmt.Sprintf("%s-%s", modelName, filename)) if err := saveImage(out, outPath); err != nil { log.Fatalln(err) } - } func loadImage(filePath string) (image.Image, error) { diff --git a/go/face/aligner/aligner.go b/go/face/aligner/aligner.go index bfaf2f5..cfe8948 100644 --- a/go/face/aligner/aligner.go +++ b/go/face/aligner/aligner.go @@ -8,7 +8,6 @@ package aligner */ import "C" import ( - "image" "unsafe" openvision "github.com/bububa/openvision/go" @@ -39,7 +38,7 @@ func (a *Aligner) SetThreads(n int) { } // Align face -func (a *Aligner) Align(img *common.Image, faceInfo face.FaceInfo) (image.Image, error) { +func (a *Aligner) Align(img *common.Image, faceInfo face.FaceInfo, out *common.Image) error { imgWidth := img.WidthF64() imgHeight := img.HeightF64() data := img.Bytes() @@ -61,7 +60,8 @@ func (a *Aligner) Align(img *common.Image, faceInfo face.FaceInfo) (image.Image, (*C.Image)(unsafe.Pointer(outImgC)), ) if errCode != 0 { - return nil, openvision.AlignFaceError(int(errCode)) + return openvision.AlignFaceError(int(errCode)) } - return common.GoImage(outImgC) + common.GoImage(outImgC, out) + return nil } diff --git a/go/face/hair/cgo.go b/go/face/hair/cgo.go new file mode 100644 index 0000000..727fd8a --- /dev/null +++ b/go/face/hair/cgo.go @@ -0,0 +1,11 @@ +// +build !vulkan + +package hair + +/* +#cgo CXXFLAGS: --std=c++11 -fopenmp +#cgo CPPFLAGS: -I ${SRCDIR}/../../../include -I /usr/local/include +#cgo LDFLAGS: -lstdc++ -lncnn -lomp -lopenvision +#cgo LDFLAGS: -L /usr/local/lib -L ${SRCDIR}/../../../lib +*/ +import "C" diff --git a/go/face/hair/cgo_vulkan.go b/go/face/hair/cgo_vulkan.go new file mode 100644 index 0000000..9d7735a --- /dev/null +++ b/go/face/hair/cgo_vulkan.go @@ -0,0 +1,11 @@ +// +build vulkan + +package hair + +/* +#cgo CXXFLAGS: --std=c++11 -fopenmp +#cgo CPPFLAGS: -I ${SRCDIR}/../../../include -I /usr/local/include +#cgo LDFLAGS: -lstdc++ -lncnn -lomp -lopenvision -lglslang -lvulkan -lSPIRV -lOGLCompiler -lMachineIndependent -lGenericCodeGen -lOSDependent +#cgo LDFLAGS: -L /usr/local/lib -L ${SRCDIR}/../../../lib +*/ +import "C" diff --git a/go/face/hair/doc.go b/go/face/hair/doc.go new file mode 100644 index 0000000..a4254a3 --- /dev/null +++ b/go/face/hair/doc.go @@ -0,0 +1,2 @@ +// Package hair include hair segmentation +package hair diff --git a/go/face/hair/hair.go b/go/face/hair/hair.go new file mode 100644 index 0000000..edd9bef --- /dev/null +++ b/go/face/hair/hair.go @@ -0,0 +1,61 @@ +package hair + +/* +#include +#include +#include "openvision/face/hair.h" +*/ +import "C" +import ( + "unsafe" + + openvision "github.com/bububa/openvision/go" + "github.com/bububa/openvision/go/common" +) + +// Hair represents Hair segmentor +type Hair struct { + d C.IHair +} + +// NewHair returns a new Hair +func NewHair() *Hair { + return &Hair{ + d: C.new_hair(), + } +} + +// Pointer implement Estimator interface +func (h *Hair) Pointer() unsafe.Pointer { + return unsafe.Pointer(h.d) +} + +// LoadModel load detecter model +func (h *Hair) LoadModel(modelPath string) error { + return common.EstimatorLoadModel(h, modelPath) +} + +// Destroy destroy C.IHair +func (h *Hair) Destroy() { + common.DestroyEstimator(h) +} + +// Matting returns hair matting image +func (h *Hair) Matting(img *common.Image, out *common.Image) error { + imgWidth := img.WidthF64() + imgHeight := img.HeightF64() + data := img.Bytes() + outImgC := common.NewCImage() + defer common.FreeCImage(outImgC) + errCode := C.hair_matting( + (C.IHair)(h.Pointer()), + (*C.uchar)(unsafe.Pointer(&data[0])), + C.int(imgWidth), C.int(imgHeight), + (*C.Image)(unsafe.Pointer(outImgC)), + ) + if errCode != 0 { + return openvision.HairMattingError(int(errCode)) + } + common.GoImage(outImgC, out) + return nil +} diff --git a/go/face/hopenet/hopenet.go b/go/face/hopenet/hopenet.go index 16dd55b..6161589 100644 --- a/go/face/hopenet/hopenet.go +++ b/go/face/hopenet/hopenet.go @@ -34,7 +34,6 @@ func (h *Hopenet) Pointer() unsafe.Pointer { // LoadModel load detecter model func (h *Hopenet) LoadModel(modelPath string) error { return common.EstimatorLoadModel(h, modelPath) - return nil } // Destroy destroy C.IHopeNet diff --git a/go/pose/segmentor/deeplabv3plus.go b/go/pose/segmentor/deeplabv3plus.go index 5f819c4..1e90e0e 100644 --- a/go/pose/segmentor/deeplabv3plus.go +++ b/go/pose/segmentor/deeplabv3plus.go @@ -7,7 +7,6 @@ package segmentor */ import "C" import ( - "image" "unsafe" "github.com/bububa/openvision/go/common" @@ -41,11 +40,11 @@ func (d *Deeplabv3plus) LoadModel(modelPath string) error { } // Matting implement Segmentor interface -func (d *Deeplabv3plus) Matting(img *common.Image) (image.Image, error) { - return Matting(d, img) +func (d *Deeplabv3plus) Matting(img *common.Image, out *common.Image) error { + return Matting(d, img, out) } // Merge implement Segmentor interface -func (d *Deeplabv3plus) Merge(img *common.Image, bg *common.Image) (image.Image, error) { - return Merge(d, img, bg) +func (d *Deeplabv3plus) Merge(img *common.Image, bg *common.Image, out *common.Image) error { + return Merge(d, img, bg, out) } diff --git a/go/pose/segmentor/erdnet.go b/go/pose/segmentor/erdnet.go index c269bbf..b6472d5 100644 --- a/go/pose/segmentor/erdnet.go +++ b/go/pose/segmentor/erdnet.go @@ -7,7 +7,6 @@ package segmentor */ import "C" import ( - "image" "unsafe" "github.com/bububa/openvision/go/common" @@ -41,11 +40,11 @@ func (d *ERDNet) LoadModel(modelPath string) error { } // Matting implement Segmentor interface -func (d *ERDNet) Matting(img *common.Image) (image.Image, error) { - return Matting(d, img) +func (d *ERDNet) Matting(img *common.Image, out *common.Image) error { + return Matting(d, img, out) } // Merge implement Segmentor interface -func (d *ERDNet) Merge(img *common.Image, bg *common.Image) (image.Image, error) { - return Merge(d, img, bg) +func (d *ERDNet) Merge(img *common.Image, bg *common.Image, out *common.Image) error { + return Merge(d, img, bg, out) } diff --git a/go/pose/segmentor/rvm.go b/go/pose/segmentor/rvm.go index a50bc5b..32243c0 100644 --- a/go/pose/segmentor/rvm.go +++ b/go/pose/segmentor/rvm.go @@ -7,7 +7,6 @@ package segmentor */ import "C" import ( - "image" "unsafe" "github.com/bububa/openvision/go/common" @@ -44,11 +43,11 @@ func (d *RVM) LoadModel(modelPath string) error { } // Matting implement Segmentor interface -func (d *RVM) Matting(img *common.Image) (image.Image, error) { - return Matting(d, img) +func (d *RVM) Matting(img *common.Image, out *common.Image) error { + return Matting(d, img, out) } // Merge implement Segmentor interface -func (d *RVM) Merge(img *common.Image, bg *common.Image) (image.Image, error) { - return Merge(d, img, bg) +func (d *RVM) Merge(img *common.Image, bg *common.Image, out *common.Image) error { + return Merge(d, img, bg, out) } diff --git a/go/pose/segmentor/segmentor.go b/go/pose/segmentor/segmentor.go index 07ff69e..2ee75b3 100644 --- a/go/pose/segmentor/segmentor.go +++ b/go/pose/segmentor/segmentor.go @@ -8,7 +8,6 @@ package segmentor */ import "C" import ( - "image" "unsafe" openvision "github.com/bububa/openvision/go" @@ -18,12 +17,12 @@ import ( // Segmentor represents segmentor interface type Segmentor interface { common.Estimator - Matting(img *common.Image) (image.Image, error) - Merge(img *common.Image, bg *common.Image) (image.Image, error) + Matting(img *common.Image, out *common.Image) error + Merge(img *common.Image, bg *common.Image, out *common.Image) error } // Matting returns pose segment matting image -func Matting(d Segmentor, img *common.Image) (image.Image, error) { +func Matting(d Segmentor, img *common.Image, out *common.Image) error { imgWidth := img.WidthF64() imgHeight := img.HeightF64() data := img.Bytes() @@ -36,13 +35,14 @@ func Matting(d Segmentor, img *common.Image) (image.Image, error) { C.int(imgHeight), (*C.Image)(unsafe.Pointer(outImgC))) if errCode != 0 { - return nil, openvision.DetectPoseError(int(errCode)) + return openvision.DetectPoseError(int(errCode)) } - return common.GoImage(outImgC) + common.GoImage(outImgC, out) + return nil } // Merge merge pose with background -func Merge(d Segmentor, img *common.Image, bg *common.Image) (image.Image, error) { +func Merge(d Segmentor, img *common.Image, bg *common.Image, out *common.Image) error { imgWidth := img.WidthF64() imgHeight := img.HeightF64() data := img.Bytes() @@ -59,7 +59,8 @@ func Merge(d Segmentor, img *common.Image, bg *common.Image) (image.Image, error C.int(bgWidth), C.int(bgHeight), (*C.Image)(unsafe.Pointer(outImgC))) if errCode != 0 { - return nil, openvision.DetectPoseError(int(errCode)) + return openvision.DetectPoseError(int(errCode)) } - return common.GoImage(outImgC) + common.GoImage(outImgC, out) + return nil } diff --git a/go/styletransfer/animegan2.go b/go/styletransfer/animegan2.go index 3f649f4..1154927 100644 --- a/go/styletransfer/animegan2.go +++ b/go/styletransfer/animegan2.go @@ -7,7 +7,6 @@ package styletransfer */ import "C" import ( - "image" "unsafe" "github.com/bububa/openvision/go/common" @@ -41,6 +40,6 @@ func (d *AnimeGan2) LoadModel(modelPath string) error { } // Transform implement StyleTransfer interface -func (d *AnimeGan2) Transform(img *common.Image) (image.Image, error) { - return Transform(d, img) +func (d *AnimeGan2) Transform(img *common.Image, out *common.Image) error { + return Transform(d, img, out) } diff --git a/go/styletransfer/transfer.go b/go/styletransfer/transfer.go index c1d09f9..d4c7b42 100644 --- a/go/styletransfer/transfer.go +++ b/go/styletransfer/transfer.go @@ -8,7 +8,6 @@ package styletransfer */ import "C" import ( - "image" "unsafe" openvision "github.com/bububa/openvision/go" @@ -18,11 +17,11 @@ import ( // StyleTransfer represents Style Transfer interface type StyleTransfer interface { common.Estimator - Transform(img *common.Image) (image.Image, error) + Transform(img *common.Image, out *common.Image) error } // Transform returns style transform image -func Transform(d StyleTransfer, img *common.Image) (image.Image, error) { +func Transform(d StyleTransfer, img *common.Image, out *common.Image) error { imgWidth := img.WidthF64() imgHeight := img.HeightF64() data := img.Bytes() @@ -35,7 +34,8 @@ func Transform(d StyleTransfer, img *common.Image) (image.Image, error) { C.int(imgHeight), (*C.Image)(unsafe.Pointer(outImgC))) if errCode != 0 { - return nil, openvision.DetectPoseError(int(errCode)) + return openvision.DetectPoseError(int(errCode)) } - return common.GoImage(outImgC) + common.GoImage(outImgC, out) + return nil } diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 5084975..6b87012 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -57,7 +57,7 @@ target_include_directories(openvision $ $ - + $ $ $ @@ -91,6 +91,7 @@ file(COPY ${CMAKE_CURRENT_SOURCE_DIR}/face/tracker.h ${CMAKE_CURRENT_SOURCE_DIR}/face/hopenet.h ${CMAKE_CURRENT_SOURCE_DIR}/face/aligner.h + ${CMAKE_CURRENT_SOURCE_DIR}/face/hair.h DESTINATION ${INCLUDE_OUTPUT_PATH}/openvision/face ) diff --git a/src/face/hair.h b/src/face/hair.h new file mode 100644 index 0000000..891f598 --- /dev/null +++ b/src/face/hair.h @@ -0,0 +1,17 @@ +#ifndef _FACE_HAIR_C_H_ +#define _FACE_HAIR_C_H_ + +#include "common.h" + +#ifdef __cplusplus +#include "hair/hair.hpp" +extern "C" { +#endif +typedef void *IHair; +IHair new_hair(); +int hair_matting(IHair d, const unsigned char *rgbdata, int img_width, + int img_height, Image *out); +#ifdef __cplusplus +} +#endif +#endif // !_FACE_HAIR_C_H_ diff --git a/src/face/hair/hair.cpp b/src/face/hair/hair.cpp new file mode 100644 index 0000000..1356f74 --- /dev/null +++ b/src/face/hair/hair.cpp @@ -0,0 +1,64 @@ +#include "../hair.h" + +#ifdef OV_VULKAN +#include "gpu.h" +#endif // OV_VULKAN + +IHair new_hair() { return new ovface::Hair(); } + +int hair_matting(IHair d, const unsigned char *rgbdata, int img_width, + int img_height, Image *out) { + int ret = static_cast(d)->Matting(rgbdata, img_width, + img_height, out); + if (ret != 0) { + return ret; + } + return 0; +} + +namespace ovface { + +int Hair::Matting(const unsigned char *rgbdata, int img_width, int img_height, + Image *out) { + if (!initialized_) { + return 10000; + } + if (rgbdata == 0) { + return 10001; + } + + ncnn::Extractor ex = net_->create_extractor(); + ex.set_light_mode(light_mode_); + ex.set_num_threads(num_threads); + + ncnn::Mat ncnn_in = + ncnn::Mat::from_pixels_resize(rgbdata, ncnn::Mat::PIXEL_RGB, img_width, + img_height, target_size, target_size); + ncnn::Mat matting = ncnn::Mat(target_size, target_size, 3); + + ncnn_in.substract_mean_normalize(mean_vals, norm_vals); + ex.input("input", ncnn_in); + ncnn::Mat ncnn_out; + ex.extract("1006", ncnn_out); + + for (int c = 0; c < 3; ++c) { + float *pImage = matting.channel(c); + for (int i = 0; i < target_size * target_size; i++) { + const float alpha = ncnn_out[i]; + float value = 255 * alpha; + value = std::max(std::min(value, 255.f), 0.f); + pImage[i] = value; + } + } + ncnn::Mat outimg; + ncnn::resize_bicubic(matting, outimg, img_width, img_height); + + out->width = outimg.w; + out->height = outimg.h; + out->channels = outimg.c; + out->data = (unsigned char *)malloc(outimg.total()); + outimg.to_pixels(out->data, ncnn::Mat::PIXEL_RGB); + return 0; +} + +} // namespace ovface diff --git a/src/face/hair/hair.hpp b/src/face/hair/hair.hpp new file mode 100644 index 0000000..a1dd486 --- /dev/null +++ b/src/face/hair/hair.hpp @@ -0,0 +1,18 @@ +#ifndef _FACE_HAIR_H_ +#define _FACE_HAIR_H_ + +#include "../common/common.h" + +namespace ovface { +class Hair : public ov::Estimator { +public: + int Matting(const unsigned char *rgbdata, int img_width, int img_height, + Image *out); + +private: + const int target_size = 288; + const float mean_vals[3] = {123.675f, 116.28f, 103.53f}; + const float norm_vals[3] = {0.01712475f, 0.0175f, 0.01742919f}; +}; +} // namespace ovface +#endif // !_FACE_HAIR_H_ diff --git a/src/face/hopenet.h b/src/face/hopenet.h index 71dd3cb..bc29506 100644 --- a/src/face/hopenet.h +++ b/src/face/hopenet.h @@ -7,9 +7,10 @@ #include "hopenet/hopenet.hpp" extern "C" { #endif - typedef void* IHopenet; - IHopenet new_hopenet(); - int hopenet_detect(IHopenet d, const unsigned char* rgbdata, int img_width, int img_height, const Rect* roi, HeadPose* euler_angles); +typedef void *IHopenet; +IHopenet new_hopenet(); +int hopenet_detect(IHopenet d, const unsigned char *rgbdata, int img_width, + int img_height, const Rect *roi, HeadPose *euler_angles); #ifdef __cplusplus } #endif diff --git a/src/face/hopenet/hopenet.cpp b/src/face/hopenet/hopenet.cpp index 3dd3987..ca10b81 100644 --- a/src/face/hopenet/hopenet.cpp +++ b/src/face/hopenet/hopenet.cpp @@ -32,7 +32,7 @@ int Hopenet::LoadModel(const char *root_path) { } int Hopenet::Detect(const unsigned char *rgbdata, int img_width, int img_height, - Rect roi, HeadPose *head_angles) { + ov::Rect roi, HeadPose *head_angles) { float diff = fabs(roi.height - roi.width); if (roi.height > roi.width) { roi.x -= diff / 2; diff --git a/src/face/hopenet/hopenet.hpp b/src/face/hopenet/hopenet.hpp index 2ed3744..5721f92 100644 --- a/src/face/hopenet/hopenet.hpp +++ b/src/face/hopenet/hopenet.hpp @@ -7,16 +7,14 @@ namespace ovface { class Hopenet : public ov::Estimator { public: - int LoadModel(const char* root_path); - int Detect(const unsigned char* rgbdata, - int img_width, int img_height, - Rect roi, HeadPose* euler_angles); + int LoadModel(const char *root_path); + int Detect(const unsigned char *rgbdata, int img_width, int img_height, + ov::Rect roi, HeadPose *euler_angles); private: - float idx_tensor[66]; - void softmax(float* z, size_t el); - double getAngle(float* prediction, size_t len); - + float idx_tensor[66]; + void softmax(float *z, size_t el); + double getAngle(float *prediction, size_t len); }; -} +} // namespace ovface #endif // !_HEAD_HOPENET_H_