2차원의 물체 검출 프레임워크를 3차원 대응으로 해 보았다

경위

지금까지 물체 검출이라고 하면 RGB 화상을 입력으로 한 수법이 주류였습니다. 하지만, realsense 등의 depth 데이터도 취할 수 있는 디바이스가 비교적 저렴하게 입수하게 되거나, depth estimation(RGB 화상으로부터 깊이 정보를 예측)의 정밀도가 올라온 등의 요인에 의해, 3D object detection 연구가 활발해지고 있습니다. 이번에는 기존의 2차원의 물체 검출 프레임워크를 3차원 대응으로 해 보았을 때의 순서를 공유합니다.

tkDNN이란?

이번 자신이 3차원 대응시킨 것은 tkDNN 이라는 darknet YOLO를 TensorRT 을 사용해 고속화시킨 프레임워크입니다. tkDNN의 절차를 소개하지만 다른 프레임 워크에서도하는 것은 거의 동일하다고 생각합니다.

3차원 대응 이미지

tkDNN과 darknet에서 RGB 이미지를 사용하는 이미지는 아래와 같은 느낌입니다. BGR에서 각 채널마다 0~255까지의 수치를 유지하고 있습니다만, DNN에 넣기 전에는 1차원 배열로 변형합니다. 이때 0~255를 정규화하여 0~1의 수치로 합니다.

그래서, BGR에 더해 depth도 더하고 싶을 때는 아래와 같은 이미지가 됩니다. 검정은 depth 데이터를 나타냅니다.

코드 변경

구체적으로 코드를 변경하는 방법을 설명합니다. 크게 나누면 두 파일을 변경합니다. tkDNN을 추론할 때 DetectionNN.h의 update 함수가 먼저 호출됩니다. 그래서 인수로 depth 데이터를 받을 수 있도록 변경합니다. 그런 다음 이미지 크기 조정과 같은 전처리를 수행하는 preprocess 함수에 depth 데이터를 전달합니다.

DetectionNN.h

// depthデータ(depth_frames)も受けとる
void update(std::vector<cv::Mat>& rgb_frames, std::vector<cv::Mat>& depth_frames, const int cur_batches=1, bool save_times=false, std::ofstream *times=nullptr, const bool mAP=false){
    ~~省略~~
    if(TKDNN_VERBOSE) printCenteredTitle(" TENSORRT detection ", '=', 30); 
    {
        TKDNN_TSTART
        for(int bi=0; bi<cur_batches;++bi){
            if(!frames[bi].data)
                FatalError("No image data feed to detection");
                originalSize.push_back(frames[bi].size());
                // 前処理にdepthデータを渡す
                preprocess(rgb_frames[bi], depth_frames[bi], bi);    
            }
            TKDNN_TSTOP
            if(save_times) *times<<t_ns<<";";
        }
    ~~省略~~
} 

다음으로 yolo에 대한 추론 처리가 작성된 Yolo3Detection.cpp를 변경합니다. 먼저 rgb와 depth 데이터를 각각 정규화합니다. 이번은 depth가 16bit이므로 65535(2의 16승-1)로 정규화하고 있습니다. 그런 다음 bgrd라는 배열에 b, g, r, d 순서로 저장합니다. 이것에 의해, RGB 데이터와 depth 데이터를 결합할 수 있었습니다.

yolo3Detection.cpp

void Yolo3Detection::preprocess(cv::Mat &rgb_frame, cv::Mat &depth_frame, const int bi){
~~省略~~
#else
    // 画像サイズをDNNの入力サイズに変更
    cv::resize(rgb_frame, rgb_frame, cv::Size(netRT->input_dim.w, netRT->input_dim.h));
    cv::resize(depth_frame, depth_frame, cv::Size(netRT->input_dim.w, netRT->input_dim.h));
    // 32float型にして0~1に正規化
    rgb_frame.convertTo(imagePreproc, CV_32FC3, 1/255.0);
    depth_frame.convertTo(depthPreproc, CV_32FC1, 1/65535.0); 

    //split channels & merge depth
    cv::split(imagePreproc,bgr);
    bgrd[0] = bgr[0];
    bgrd[1] = bgr[1];
    bgrd[2] = bgr[2];
    bgrd[3] = depthPreproc;

    //write channels
    for(int i=0; i<netRT->input_dim.c; i++) {
        int idx = i*imagePreproc.rows*imagePreproc.cols;
        int ch = netRT->input_dim.c-1 -i;
        memcpy((void*)&input[idx + netRT->input_dim.tot()*bi], (void*)bgrd[ch].data, imagePreproc.rows*imagePreproc.cols*sizeof(dnnType));     
    }
    checkCuda(cudaMemcpyAsync(input_d + netRT->input_dim.tot()*bi, input + netRT->input_dim.tot()*bi, netRT->input_dim.tot()*sizeof(dnnType), cudaMemcpyHostToDevice, netRT->stream));
#endif
}

변수의 선언등의 세세한 변경점은 생략했습니다만, 대략적으로는 이런 느낌으로 depth 대응을 할 수 있습니다.

요약

tkDNN을 만든 사람, 정말 상냥한 사람입니다.

실수나 질문, 의견등 있으면 부담없이 코멘트해 주세요. 열심히 대답하니까(웃음).