package main
|
|
/*
|
#cgo CFLAGS: -I${SRCDIR}/sdk/include -w -g
|
#cgo CXXFLAGS: -I${SRCDIR}/sdk/include -w -g -std=c++11
|
#cgo LDFLAGS: -L/usr/local/cuda-8.0/lib64 -L${SRCDIR}/sdk/lib
|
#cgo LDFLAGS: -Wl,-rpath,${SRCDIR}/sdk/lib
|
#cgo LDFLAGS: -ldarknet -lcudart -lcublas -lcurand -lcudnn -lrt -ldl -lpthread
|
#include <stdlib.h>
|
#include "cyolo.h"
|
*/
|
import "C"
|
import (
|
"unsafe"
|
|
"basic.com/libgowrapper/sdkstruct.git"
|
)
|
|
type trackInfo struct {
|
lastTrackObjs []sdkstruct.CObjTrackInfo
|
lastTrackID uint64
|
}
|
|
// YoloHandle wrap C
|
type YoloHandle struct {
|
handle unsafe.Pointer
|
tracker map[string]*trackInfo
|
}
|
|
// RatioInterTrack 跟踪判断重叠阈值
|
const RatioInterTrack = 50 //跟踪判断重叠阈值
|
|
// NewSDK init yolo sdk
|
func NewSDK(fc, fw, fn string, gi int) interface{} {
|
|
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.init(c, w, n, g)
|
return &YoloHandle{
|
handle: p,
|
tracker: make(map[string]*trackInfo),
|
}
|
}
|
|
// Free free
|
func Free(i interface{}) {
|
y := i.(*YoloHandle)
|
if y != nil && y.handle != nil {
|
C.release(y.handle)
|
}
|
}
|
|
// CYoloObjInfoArrayToGoArray convert cObjInfo array to go
|
func CYoloObjInfoArrayToGoArray(cArray unsafe.Pointer, count int) (goArray []sdkstruct.CObjInfo) {
|
p := uintptr(cArray)
|
|
for i := 0; i < count; i++ {
|
j := *(*sdkstruct.CObjInfo)(unsafe.Pointer(p))
|
goArray = append(goArray, j)
|
p += unsafe.Sizeof(j)
|
}
|
return
|
}
|
|
// YoloDetect yolo detect
|
func YoloDetect(y *YoloHandle, data []byte, w, h, c int, thrsh float32, umns int) []sdkstruct.CObjInfo {
|
|
var count C.int
|
var cobjinfo unsafe.Pointer
|
|
ret := C.detect(y.handle,
|
unsafe.Pointer(&data[0]), C.int(w), C.int(h), C.int(c),
|
C.float(thrsh), C.int(umns),
|
&cobjinfo, &count)
|
|
if ret == 0 {
|
defer C.free(cobjinfo)
|
return CYoloObjInfoArrayToGoArray(unsafe.Pointer(cobjinfo), int(count))
|
}
|
return nil
|
}
|
|
// YoloObjName obj name by type
|
func YoloObjName(typ int) string {
|
p := C.obj_name_by_type(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 sdkstruct.CRECT, rect2 sdkstruct.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
|
}
|
|
// YoloDetectTrack2 yolo detect (只识别人)
|
func YoloDetectTrack2(y *YoloHandle, LastYoloObjs []sdkstruct.CObjTrackInfo, LastTrackID *uint64, data []byte, w, h, c int, thrsh float32, umns int) (allObjs []sdkstruct.CObjTrackInfo, newObjs []sdkstruct.CObjTrackInfo) {
|
|
var tmp sdkstruct.CObjTrackInfo
|
//LastYoloObjs
|
detectObjs := YoloDetect(y, data, w, h, c, 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
|
}
|
|
// Run yolo detect (只识别人)
|
func Run(i interface{}, id string, data []byte, w, h, c int, thrsh float32, umns int) ([]sdkstruct.CObjTrackInfo, []sdkstruct.CObjTrackInfo) {
|
if data == nil || w <= 0 || h <= 0 {
|
return nil, nil
|
}
|
y := i.(*YoloHandle)
|
|
channel := c
|
if channel == 0 {
|
channel = 3
|
}
|
|
v, ok := y.tracker[id]
|
if !ok {
|
i := &trackInfo{}
|
y.tracker[id] = i
|
v = i
|
}
|
whole, recent := YoloDetectTrack2(y, v.lastTrackObjs, &v.lastTrackID, data, w, h, channel, thrsh, umns)
|
y.tracker[id].lastTrackObjs = whole
|
y.tracker[id].lastTrackID = v.lastTrackID
|
|
return whole, recent
|
}
|
