valib/gosdk.git

package gosdk
 
/*
#cgo CFLAGS: -I. -I./sdk/face/include -I./sdk/darknet/include -I/usr/local/cuda/include -w -g
#cgo CXXFLAGS: -I. -I./sdk/face/include -I./sdk/darknet/include -I/usr/local/cuda/include -w -g -std=c++11
#cgo LDFLAGS: -L/usr/local/cuda/lib64 -L${SRCDIR}/sdk/face/lib/gpu -L${SRCDIR}/sdk/darknet/lib
#cgo LDFLAGS: -Wl,-rpath,${SRCDIR}/sdk/face/lib/gpu:${SRCDIR}/sdk/darknet/lib
#cgo LDFLAGS: -ldarknet -lTHFaceImage -lTHFeature -lTHFaceProperty -lTHFaceTracking -lcudart -lcublas -lcurand -lrt -ldl -lpthread
#include <stdlib.h>
#include "csdk.h"
*/
import "C"
import (
    "unsafe"
)
 
// YoloHandle wrap C
type YoloHandle struct {
    handle       C.YoloHandle
    LastYoloObjs []CObjTrackInfo //yolo跟踪的上一帧信息
    LastTrackID  uint64          //yolo 被使用的ID
}
 
const RatioInterTrack = 50 //跟踪判断重叠阈值
 
// SDKImage sdk image
type SDKImage struct {
    Data   []byte
    Width  int
    Height int
}
 
// InitYolo init yolo sdk
func InitYolo(fc, fw, fn string, gi int) *YoloHandle {
 
    c := C.CString(fc)
    defer C.free(unsafe.Pointer(c))
    w := C.CString(fw)
    defer C.free(unsafe.Pointer(w))
    n := C.CString(fn)
    defer C.free(unsafe.Pointer(n))
 
    g := C.int(gi)
 
    p := C.c_api_yolo_init(c, w, n, g)
    return &YoloHandle{handle: p}
}
 
// InitFaceDetector init face detector
func InitFaceDetector(tm, gi, minFaces, rollAngle int) {
 
    C.c_api_face_detector_init(C.int(tm), C.int(gi), C.int(minFaces), C.int(rollAngle))
}
 
// InitFaceProperty init face property
func InitFaceProperty(tm int) {
 
    C.c_api_face_property_init(C.int(tm))
}
 
// InitFaceExtractor init face extractor
func InitFaceExtractor(tm, gi int) {
    C.c_api_face_extractor_init(C.int(tm), C.int(gi))
}
 
// InitFaceTracker init face tracker
func InitFaceTracker(tm, gi, w, h, maxFaces, interval, sample int) {
 
    C.c_api_face_tracker_init(C.int(tm), C.int(gi), C.int(w), C.int(h), C.int(maxFaces), C.int(interval), C.int(sample))
}
 
// ResizeFaceTracker init face tracker
func ResizeFaceTracker(ch, w, h int) int {
 
    return int(C.c_api_face_track_resize(C.int(ch), C.int(w), C.int(h)))
}
 
// Free free sdk
func Free() {
    C.c_api_release()
}
 
// FaceDetect face detect
func FaceDetect(img SDKImage, ch int) []CFacePos {
 
    data := img.Data
    w := img.Width
    h := img.Height
 
    var count C.int
    cfpos := C.c_api_face_detect(&count, (*C.uchar)(unsafe.Pointer(&data[0])), C.int(w), C.int(h), C.int(ch))
    if cfpos != nil {
        defer C.free(unsafe.Pointer(cfpos))
        return CFacePosArrayToGoArray(unsafe.Pointer(cfpos), int(count))
    }
    return nil
}
 
// FaceProperty face property
func FaceProperty(fpos CFacePos, img SDKImage, ch int) CThftResult {
    data := img.Data
    w := img.Width
    h := img.Height
 
    pos := (*C.cFacePos)(unsafe.Pointer(&fpos))
    result := C.c_api_face_property(pos, (*C.uchar)(unsafe.Pointer(&data[0])), C.int(w), C.int(h), C.int(ch))
 
    return *((*CThftResult)(unsafe.Pointer(&result)))
}
 
// FaceExtract face extract feature
func FaceExtract(fpos CFacePos, img SDKImage, ch int) []byte {
 
    data := img.Data
    w := img.Width
    h := img.Height
 
    var featLen C.int
    pos := (*C.cFacePos)(unsafe.Pointer(&fpos))
 
    p := C.c_api_face_extract(&featLen, pos, (*C.uchar)(unsafe.Pointer(&data[0])), C.int(w), C.int(h), C.int(ch))
    defer C.free(unsafe.Pointer(p))
    b := C.GoBytes(unsafe.Pointer(p), featLen)
    return b
}
 
// FaceCompare face compare
func FaceCompare(feat1 []byte, feat2 []byte) float32 {
    res := C.c_api_face_compare((*C.uchar)(unsafe.Pointer(&feat1[0])), (*C.uchar)(unsafe.Pointer(&feat2[0])))
    return float32(res)
}
 
// FaceTrackOnly face tracker face only
func FaceTrackOnly(img SDKImage, ch int) []CRECT {
    data := img.Data
    w := img.Width
    h := img.Height
 
    var fCount C.int
    rect := C.c_api_face_track_only(&fCount, (*C.uchar)(unsafe.Pointer(&data[0])), C.int(w), C.int(h), C.int(ch))
    if rect != nil {
        defer C.free(unsafe.Pointer(rect))
        return CRECTArrayToGoArray(unsafe.Pointer(rect), int(fCount))
    }
    return nil
}
 
// FaceTrackDetect face tracker face detect
func FaceTrackDetect(img SDKImage, ch int) []CFaceInfo {
    data := img.Data
    w := img.Width
    h := img.Height
 
    var fCount C.int
    finfo := C.c_api_face_track_only(&fCount, (*C.uchar)(unsafe.Pointer(&data[0])), C.int(w), C.int(h), C.int(ch))
    if finfo != nil {
        defer C.free(unsafe.Pointer(finfo))
        return CFaceInfoArrayToGoArray(unsafe.Pointer(finfo), int(fCount))
    }
    return nil
}
 
// FaceTrackingInfo face track info
type FaceTrackingInfo struct {
    tracking bool
    faces    []CFaceInfo
}
 
// FaceTrack face tracking info
func FaceTrack(res map[int64]FaceTrackingInfo, img SDKImage, ch int) {
    data := img.Data
    w := img.Width
    h := img.Height
 
    var fCount C.int
    cFinfo := C.c_api_face_track(&fCount, (*C.uchar)(unsafe.Pointer(&data[0])), C.int(w), C.int(h), C.int(ch))
    if cFinfo == nil {
        return
    }
    defer C.free(unsafe.Pointer(cFinfo))
    goFinfo := CFaceInfoArrayToGoArray(unsafe.Pointer(cFinfo), int(fCount))
 
    // 空，添加
    if len(res) == 0 {
        for _, v := range goFinfo {
            var faces []CFaceInfo
            faces = append(faces, v)
            tInfo := FaceTrackingInfo{true, faces}
            res[v.NFaceID] = tInfo
        }
        // 非空
    } else {
        // 默认所有faces没有tracking
        for _, v := range res {
            v.tracking = false
        }
        // 查找是否存在tracking
        for _, v := range goFinfo {
            // 存在已有信息
            if info, ok := res[v.NFaceID]; ok {
                // tracking到face
                info.tracking = true
                // 如果tracking长度大于100，删除一些
                if len(info.faces) > 100 {
                    info.faces = info.faces[1:]
                }
                info.faces = append(info.faces, v)
                // 不存在，新face，加入
            } else {
                var faces []CFaceInfo
                faces = append(faces, v)
                tInfo := FaceTrackingInfo{true, faces}
                res[v.NFaceID] = tInfo
            }
        }
        // 删除不再tracking的faces
        for k, v := range res {
            if !v.tracking {
                delete(res, k)
            }
        }
    }
}
 
func FaceInfo2FacePos(face CFaceInfo) (fPos CFacePos) {
    fPos.RcFace = face.RcFace
    fPos.PtLeftEye = face.PtLeftEye
    fPos.PtRightEye = face.PtRightEye
    fPos.PtNose = face.PtNose
    fPos.PtMouth = face.PtMouth
    fPos.FAngle.Yaw = face.FAngle.Yaw
    fPos.FAngle.Pitch = face.FAngle.Pitch
    fPos.FAngle.Roll = face.FAngle.Roll
    fPos.FAngle.Confidence = face.FAngle.Confidence
 
    copy(fPos.PFacialData[:], face.PFacialData[:512])
 
    return fPos
}
 
// FaceTrackSimple face tracking info
func FaceTrackSimple(img SDKImage, ch int) (faces []CFaceInfo) {
    data := img.Data
    w := img.Width
    h := img.Height
 
    var fCount C.int
    cFinfo := C.c_api_face_track(&fCount, (*C.uchar)(unsafe.Pointer(&data[0])), C.int(w), C.int(h), C.int(ch))
    // fmt.Println("cFinfo detected:", cFinfo)
    if cFinfo == nil {
        return faces
    }
    defer C.free(unsafe.Pointer(cFinfo))
    faces = CFaceInfoArrayToGoArray(unsafe.Pointer(cFinfo), int(fCount))
    //if len(faces) > 0{
    // fmt.Println("faces detected:", len(faces))
    //}
 
    return faces
}
 
// YoloDetect yolo detect
func YoloDetect(handle *YoloHandle, img SDKImage, thrsh float32, umns int) []CObjInfo {
 
    data := img.Data
    w := img.Width
    h := img.Height
 
    var count C.int
 
    cobjinfo := C.c_api_yolo_detect(handle.handle, &count, (*C.uchar)(unsafe.Pointer(&data[0])), C.int(w), C.int(h), C.float(thrsh), C.int(umns))
 
    if cobjinfo != nil {
        defer C.free(unsafe.Pointer(cobjinfo))
        return CYoloObjInfoArrayToGoArray(unsafe.Pointer(cobjinfo), int(count))
    }
    return nil
}
 
// YoloObjName obj name by type
func YoloObjName(typ int) string {
    p := C.c_api_yolo_obj_name(C.int(typ))
 
    return C.GoString(p)
}
 
func max(a, b int32) int32 {
    if a < b {
        return b
    }
    return a
}
 
func min(a, b int32) int32 {
    if a < b {
        return a
    }
    return b
}
 
func countInterAreaOfTwoRect(rect1 CRECT, rect2 CRECT) int32 {
    xMin := min(rect1.Left, rect2.Left)
    yMin := min(rect1.Top, rect2.Top)
    xMax := max(rect1.Right, rect2.Right)
    yMax := max(rect1.Bottom, rect2.Bottom)
 
    wRect1 := rect1.Right - rect1.Left
    hRect1 := rect1.Bottom - rect1.Top
 
    wRect2 := rect2.Right - rect2.Left
    hRect2 := rect2.Bottom - rect2.Top
 
    wInter := wRect1 + wRect2 - (xMax - xMin)
    hInter := hRect1 + hRect2 - (yMax - yMin)
 
    if (wInter <= 0) || (hInter <= 0) {
        return 0
    }
 
    areaInter := wInter * hInter
    areaRect1 := wRect1 * hRect1
    areaRect2 := wRect2 * hRect2
    ratio := areaInter * 100 / min(areaRect1, areaRect2)
 
    return ratio
}
 
// YoloDetect yolo detect   (只识别人)
func YoloDetectTrack(handle *YoloHandle, img SDKImage, thrsh float32, umns int) (allObjs []CObjTrackInfo, newObjs []CObjTrackInfo) {
 
    var tmp CObjTrackInfo
    //LastYoloObjs
    detectObjs := YoloDetect(handle, img, thrsh, umns)
 
    for _, vLast := range handle.LastYoloObjs {
        for i := 0; i < len(detectObjs); i++ {
            //fmt.Println("vNew.Typ:", vNew.Typ)
            if vLast.ObjInfo.Typ == detectObjs[i].Typ { //同一类别，比如都是人体
                ratio := countInterAreaOfTwoRect(vLast.ObjInfo.RcObj, detectObjs[i].RcObj)
                if ratio >= RatioInterTrack {
                    //update LastYoloObjs
                    vLast.ObjInfo.RcObj = detectObjs[i].RcObj
                    vLast.ObjInfo.Prob = detectObjs[i].Prob
 
                    allObjs = append(allObjs, vLast)
                    detectObjs = append(detectObjs[:i], detectObjs[i+1:]...) //从检测目标里删除已经查到的跟踪目标
                    i--
                    break //上一帧跟踪的目标已经找到，无需往下处理其他检测目标
                }
            }
        }
    }
 
    //处理新出现的目标
    if len(detectObjs) > 0 {
        for _, vAdd := range detectObjs {
            tmp.ObjInfo = vAdd
            tmp.ID = handle.LastTrackID
            handle.LastTrackID++
 
            allObjs = append(allObjs, tmp)
            newObjs = append(newObjs, tmp)
        }
    }
 
    //刷新上一帧的跟踪目标
    handle.LastYoloObjs = allObjs
 
    return allObjs, newObjs
}
 
// YoloDetectTrack2 yolo detect   (只识别人)
func YoloDetectTrack2(handle *YoloHandle, LastYoloObjs []CObjTrackInfo, LastTrackID *uint64, img SDKImage, thrsh float32, umns int) (allObjs []CObjTrackInfo, newObjs []CObjTrackInfo) {
 
    var tmp CObjTrackInfo
    //LastYoloObjs
    detectObjs := YoloDetect(handle, img, thrsh, umns)
 
    for _, vLast := range LastYoloObjs {
        for i := 0; i < len(detectObjs); i++ {
            //fmt.Println("vNew.Typ:", vNew.Typ)
            if vLast.ObjInfo.Typ == detectObjs[i].Typ { //同一类别，比如都是人体
                ratio := countInterAreaOfTwoRect(vLast.ObjInfo.RcObj, detectObjs[i].RcObj)
                if ratio >= RatioInterTrack {
                    //update LastYoloObjs
                    vLast.ObjInfo.RcObj = detectObjs[i].RcObj
                    vLast.ObjInfo.Prob = detectObjs[i].Prob
 
                    allObjs = append(allObjs, vLast)
                    detectObjs = append(detectObjs[:i], detectObjs[i+1:]...) //从检测目标里删除已经查到的跟踪目标
                    i--
                    break //上一帧跟踪的目标已经找到，无需往下处理其他检测目标
                }
            }
        }
    }
 
    //处理新出现的目标
    id := *LastTrackID
    if len(detectObjs) > 0 {
        for _, vAdd := range detectObjs {
            tmp.ObjInfo = vAdd
            tmp.ID = id
            id++
 
            allObjs = append(allObjs, tmp)
            newObjs = append(newObjs, tmp)
        }
    }
    *LastTrackID = id
    return allObjs, newObjs
}