yolo跟踪代码及调试

chenshijun

2019-07-10 deb2e7127b5e9315eaad9de95b47fa82e0c8dffd

yolo跟踪代码及调试

 go2c.go

@@ -58,4 +58,10 @@
    RcObj    CRECT
    Typ    int32
    Prob    float32
}

//track yolo objs info
type CObjTrackInfo struct {
    ObjInfo CObjInfo
    ID      uint64
}

 gosdk.go

@@ -26,6 +26,10 @@
    Height int
}

var LastYoloObjs []CObjTrackInfo //yolo跟踪的上一帧信息
var LastTrackID uint64 = 0       //yolo 被使用的ID
const RatioInterTrack = 50       //跟踪判断重叠阈值

// InitYolo init yolo sdk
func InitYolo(fc, fw, fn string, gi int) *YoloHandle {

@@ -236,3 +240,94 @@

    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 i := 0; i < len(detectObjs); i++ {
        if detectObjs[i].Typ != 0 {
            detectObjs = append(detectObjs[:i], detectObjs[i+1:]...) //从检测目标里删除已经查到的跟踪目标
            i--
        }
    }

    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 //上一帧跟踪的目标已经找到，无需往下处理其他检测目标
                }
            }
        }
    }

    //处理新出现的目标
    if len(detectObjs) > 0 {
        for _, vAdd := range detectObjs {
            tmp.ObjInfo = vAdd
            tmp.ID = LastTrackID
            LastTrackID++

            allObjs = append(allObjs, tmp)
            newObjs = append(newObjs, tmp)
        }
    }

    //刷新上一帧的跟踪目标
    LastYoloObjs = allObjs

    return allObjs, newObjs
}