Kai - Golang实现的目标检测云服务

YOLO/Darknet是目前比较流行的Object Detection算法（后面统一称为Darknet），在GPU上的表现不但速度快而且准确率很高。但是使用起来不方便，只提供了命令行接口和简单的Python接口。所以我想用RESTful来实现一个云端的Darknet服务kai。

选择用Go的原因不是考虑并发，而是goroutine之间的同步能方便的处理，适合实现Pipeline的功能。问题来了，Darknet是c语言实现的，那Go必须得用cgo进行封装，才能调用c函数。目标是为了实现三个基本功能：1. 图片检测 2. 视频检测 3. 摄像头检测。为了方便使用我修改了Darknet的部分代码，然后重新定义下面几个函数：

// Set a gpu device
void set_gpu(int gpu);

// Recognize a image
void image_detector(char *datacfg, char *cfgfile, char *weightfile, char *filename,
    float thresh, float hier_thresh, char *outfile);

// Recognize a video
void video_detector(char *datacfg, char *cfgfile, char *weightfile, char *filename,
    float thresh, float hier_thresh, char *outfile);

// Recognize a camera stream
void camera_detector(char *datacfg, char *cfgfile, char *weightfile, int camindex,
        float thresh, float hier_thresh, char *outpath);

有了这几个函数，就好办了，下面用cgo导入相应的库和头文件即可：

// #cgo pkg-config: opencv
// #cgo linux LDFLAGS: -ldarknet -lm -L/usr/local/cuda/lib64 -lcuda -lcudart -lcublas -lcurand -lcudnn
// #cgo darwin LDFLAGS: -ldarknet
// #include "yolo.h"
import "C"

// SetGPU set a gpu device you want
func SetGPU(gpu int) {
    C.set_gpu(C.int(gpu))
}

// ImageDetector recognize a image
func ImageDetector(dc, cf, wf, fn string, t, ht float64, of ...string) {
    ...
}

// VideoDetector recognize a video
func VideoDetector(dc, cf, wf, fn string, t, ht float64, of ...string) {
    ...
}

// CameraDetector recognize a camera stream
func CameraDetector(dc, cf, wf string, i int, t, ht float64, of ...string) {
    ...
}

这样对Darknet的封装go-yolo就完成了。

下面进入主题，介绍一下kai的实现。

kai的设计目标如下：

• 后端基于Darknet（不支持训练）
• 提供RESTful接口进行图片和视频的检测
• 支持Amazon S3下载和上传
• 支持Ftp下载和上传
• 支持检测结果持久化到MongoDB

架构图是这样的

这里重点介绍一下Kai的Pipeline机制，这里的Pipeline包括下载（Download），检测（Yolo）和上（Upload）传这一系列流程。
先上个图：

这里的难点在于下载（Download），检测（Yolo）和上传（Upload）这三个步骤可以配置不同的Goroutine数量，而这三步之间是一个同步操作。

• 首先需要定义3个buffered channel来进行同步

// KaiServer represents the server for processing all job requests
type KaiServer struct {
    net.Listener
    logger        *logging.Logger
    config        types.ServerConfig
    listenAddr    string
    listenNetwork string
    router        *Router
    server        *http.Server
    db            db.Storage
    // jobDownBuff is the buffered channel for job downloading
    jobDownBuff chan types.Job
    // jobDownBuff is the buffered channel for job todo
    jobTodoBuff chan types.Job
    // jobDownBuff is the buffered channel for job done
    jobDoneBuff chan types.Job
}

• Pipeline的执行流程如下

// Pipeline contains downloading, processing and uploading a job
func Pipeline(logger *logging.Logger, config types.ServerConfig, dbInstance db.Storage, jobDownBuff chan types.Job,
    jobTodoBuff chan types.Job, jobDoneBuff chan types.Job, job types.Job) {
    logger.Infof("pipeline-job %+v", job)

    // download a job
    setupAndDownloadJob(logger, config.System, dbInstance, job, jobDownBuff)

    // jobDownBuff -> jobTodoBuff -> jobDoneBuff
    yoloJob(logger, config, dbInstance, jobDownBuff, jobTodoBuff, jobDoneBuff)

    // upload a job
    uploadJob(logger, dbInstance, jobDoneBuff)
}

• 下载（Download）

// setupAndDownloadJob setup and download jobs into jobDownBuff
func setupAndDownloadJob(logger *logging.Logger, config types.SystemConfig,
    dbInstance db.Storage, job types.Job, jobDownBuff chan<- types.Job) {

    go func() {
        logger.Infof("start setup and download a job: %+v", job)
        newJob, err := SetupJob(logger, job.ID, dbInstance, config)
        job = *newJob
        if err != nil {
            logger.Error("setup-job failed", err)
            return
        }

        downloadFunc := downloaders.GetDownloadFunc(job.Source)
        if err := downloadFunc(logger, config, dbInstance, job.ID); err != nil {
            logger.Error("download failed", err)
            job.Status = types.JobError
            job.Details = err.Error()
            dbInstance.UpdateJob(job.ID, job)
            return
        }

        jobDownBuff <- job
    }()
}

• 检测（Yolo）

func yoloJob(logger *logging.Logger, config types.ServerConfig, dbInstance db.Storage,
    jobDownBuff <-chan types.Job, jobTodoBuff chan types.Job, jobDoneBuff chan types.Job) {

    go func() {
        job, ok := <-jobDownBuff
        if !ok {
            logger.Info("job download buffer is closed")
            return
        }
        logger.Infof("start a yolo job: %+v", job)
        // limit the number of job in the jobTodoBuff
        jobTodoBuff <- job
        jobTodo, ok := <-jobTodoBuff
        if !ok {
            logger.Info("job todo buffer is closed")
            return
        }

        nGpu := config.System.NGpu
        t := yolo.NewTask(config.Yolo, jobTodo.Media.Cate, nGpu, jobTodo.LocalSource, jobTodo.LocalDestination)
        logger.Debugf("yolo task: %+v", *t)
        yolo.StartTask(t, logger, dbInstance, jobTodo.ID)
        jobDoneBuff <- job
    }()
}

• 上传（Upload）

func uploadJob(logger *logging.Logger, dbInstance db.Storage, jobDoneBuff <-chan types.Job) {
    go func() {
        jobDone, ok := <-jobDoneBuff
        if !ok {
            logger.Info("job done buffer is closed")
            return
        }
        logger.Infof("start a upload job: %+v", jobDone)

        uploadFunc := uploaders.GetUploadFunc(jobDone.Destination)
        if err := uploadFunc(logger, dbInstance, jobDone.ID); err != nil {
            logger.Error("upload failed", err)
            jobDone.Status = types.JobError
            jobDone.Details = err.Error()
            dbInstance.UpdateJob(jobDone.ID, jobDone)
            return
        }

        logger.Info("erasing temporary files")
        if err := CleanSwap(dbInstance, jobDone.ID); err != nil {
            logger.Error("erasing temporary files failed", err)
        }

        jobDone.Status = types.JobFinished
        dbInstance.UpdateJob(jobDone.ID, jobDone)

        logger.Infof("end a job: %+v", jobDone)
    }()
}

到此，这个项目主要机制都已经介绍完了，如果大家有兴趣的可以去点击下面的项目主页。

项目链接：
go-yolo: https://github.com/ZanLabs/go...
kai: https://github.com/ZanLabs/kai