127 lines
2.3 KiB
Go
127 lines
2.3 KiB
Go
package main
|
|
|
|
import (
|
|
"bytes"
|
|
"encoding/base64"
|
|
"encoding/json"
|
|
"fmt"
|
|
"image"
|
|
"image/jpeg"
|
|
"io"
|
|
|
|
"os"
|
|
|
|
"github.com/Kagami/go-face"
|
|
"golang.org/x/image/draw"
|
|
)
|
|
|
|
type Result struct {
|
|
Success bool `json:"success"`
|
|
Descriptor face.Descriptor `json:"descriptor"`
|
|
Error string `json:"error"`
|
|
}
|
|
|
|
var res Result
|
|
|
|
func atExit() {
|
|
resString, err := json.Marshal(res)
|
|
if err != nil {
|
|
os.Exit(1)
|
|
}
|
|
|
|
fmt.Println(string(resString))
|
|
|
|
os.Exit(0)
|
|
}
|
|
|
|
func resizeIfNeeded(img image.Image, max int) image.Image {
|
|
b := img.Bounds()
|
|
w := b.Dx()
|
|
h := b.Dy()
|
|
|
|
if w <= max && h <= max {
|
|
return img
|
|
}
|
|
|
|
var scale float64
|
|
if w > h {
|
|
scale = float64(max) / float64(w)
|
|
} else {
|
|
scale = float64(max) / float64(h)
|
|
}
|
|
|
|
newW := int(float64(w) * scale)
|
|
newH := int(float64(h) * scale)
|
|
|
|
dst := image.NewRGBA(image.Rect(0, 0, newW, newH))
|
|
draw.BiLinear.Scale(dst, dst.Bounds(), img, b, draw.Over, nil)
|
|
|
|
return dst
|
|
}
|
|
|
|
func main() {
|
|
defer atExit()
|
|
|
|
// Read input image from stdin in Base64 encoding
|
|
input, err := io.ReadAll(os.Stdin)
|
|
if err != nil {
|
|
res.Success = false
|
|
res.Error = "[STDIN]: " + err.Error()
|
|
return
|
|
}
|
|
|
|
// Decode Base64
|
|
image, err := base64.StdEncoding.DecodeString(string(input))
|
|
if err != nil {
|
|
res.Success = false
|
|
res.Error = "[BASE64]: " + err.Error()
|
|
return
|
|
}
|
|
|
|
imageJPEG, err := jpeg.Decode(bytes.NewReader(image))
|
|
if err != nil {
|
|
res.Success = false
|
|
res.Error = "[JPEG][DECODE]: " + err.Error()
|
|
return
|
|
}
|
|
|
|
imageJPEG = resizeIfNeeded(imageJPEG, 600)
|
|
|
|
var buf bytes.Buffer
|
|
err = jpeg.Encode(&buf, imageJPEG, &jpeg.Options{Quality: 90})
|
|
if err != nil {
|
|
res.Success = false
|
|
res.Error = "[JPEG][ENCODE]: " + err.Error()
|
|
return
|
|
}
|
|
|
|
image = buf.Bytes()
|
|
|
|
// Create recognizer - provide ./models dir
|
|
rec, err := face.NewRecognizer("models")
|
|
if err != nil {
|
|
res.Success = false
|
|
res.Error = "[DLIB][INIT]: " + err.Error()
|
|
return
|
|
}
|
|
defer rec.Close()
|
|
|
|
faces, err := rec.RecognizeCNN(image)
|
|
if err != nil {
|
|
res.Success = false
|
|
res.Error = "[DLIB][CNN]: " + err.Error()
|
|
return
|
|
} else if len(faces) == 0 {
|
|
res.Success = false
|
|
res.Error = "[DLIB][CNN]: No face found in image."
|
|
return
|
|
} else if len(faces) != 1 {
|
|
res.Success = false
|
|
res.Error = "[DLIB][CNN]: Multiple faces found in image."
|
|
return
|
|
}
|
|
|
|
res.Success = true
|
|
res.Descriptor = faces[0].Descriptor
|
|
}
|