SAHIのアイデアの実装をc++でざっくりやる

画像内の小さな物体を検出するための手法として「SHIA」なるものを見つけました。アイデア自体はいまやkaggleなどちょこちょこ目にするのですが、pythonばかり。c++で使いたい機会があったので、練習も兼ねてざっくりと実装します（厳密にリポジトリのプログラムを再現してはいません）。

SAHIでの大まかな流れは

1. （大きな）入力画像を、オーバーラップありで細かく分ける
2. それぞれをモデル検出にかける
3. bouding boxを統合する

です。

冒頭のリンクのリポジトリに含まれている「yolov5の例」をみてみると、get_sliced_prediction 関数を実行しており、このソースコードを見ることでなんとなく上記の流れが分かります。

c++のコードに起こしてみる

お題として、opencv4の「haarcascades」をつかって顔検出を行います。本当は物体検出みたいにラベルが複数出てくる方が好ましいですが、すぐ用意できなかったので。

条件や制約:

• モデル: haarcascade_frontalface_alt.xml
• 入力画像サイズ: 2560x1440 (HDの縦横2倍）
• クロップサイズ: 定数決め打ち。
• 検出対象: 顔のみ
• 検出範囲: 自分のユースケース的に端っこは不要だったので、画像の上と左右の端を捨ててます。
• opencvで縦横とxyがごっちゃになったので、命名で下記のルールがあります。許してください

  w: 幅、すなわち横方向 / h: 高さ、すなわち縦方向

main.cpp

#include <string>
#include <opencv2/opencv.hpp>
#include <opencv2/imgcodecs.hpp>
#include <opencv2/highgui.hpp>
#include <vector>
#include <iostream>
#include <random>
#include "util.cpp"


using namespace cv;
using namespace std;

struct DetectionBboxesResults {
    std::vector<BoundingBox> bboxes;
};

std::mt19937_64 mt64(0);  
std::uniform_real_distribution<double> uni(0, 1);


DetectionBboxesResults detect_faces(CascadeClassifier &cascade, cv::Mat& img){
    vector<Rect> faces;
    cascade.detectMultiScale(img, faces, 1.1, 3, 0, Size(20, 20)); //カスケードファイルに基づいて顔を検知する．検知した顔情報をベクトルfacesに格納
    vector<BoundingBox> detection_results;
    for(auto face: faces){
        // score取得方法がわからなかったので乱数でごまかす
        float score = uni(mt64);
        std::vector<float> coords{float(face.x), float(face.y), float(face.x+face.width), float(face.y+face.height)};
        BoundingBox res{76, score, coords};
        detection_results.push_back(res);
    }

    DetectionBboxesResults result{detection_results};
    return result;
}

struct TiledPrediction {
    int start_x;
    int start_y;
    DetectionBboxesResults prediction;
};


int main(int argc, char* argv[])
{
    Mat img = imread(argv[1], IMREAD_UNCHANGED); //入力画像読み込み
    
    CascadeClassifier cascade;

    //顔検出のためのxmlファイル読み込み
    cascade.load("/usr/share/opencv4/haarcascades/haarcascade_frontalface_alt.xml"); 

    int loaded_rows = img.rows;
    int loaded_cols = img.cols;

    cv::resize(img, img, cv::Size(2560,1440));

    int rows = img.rows;
    int cols = img.cols;
    

    // 画像を切り分ける
    int tile_size = 600;
    float over_lap_rate = 0.4; 

    int tile_overlap_size = int(tile_size * over_lap_rate);
    float non_overlap_rate = 1-over_lap_rate;
    int tile_slide = int(tile_size * non_overlap_rate);
    int num_w_tiles = (cols-tile_overlap_size) / tile_slide;
    int num_h_tiles = (rows-tile_overlap_size) / tile_slide;

    // 正方形のタイルを使いたいので、上と左右の端は捨てることにする。
    // SAHIリポジトリの手法では柔軟にやっており、正方形でもない。
    int h_start = rows - tile_slide * num_h_tiles - tile_overlap_size;
    int w_start = (cols - tile_slide * num_w_tiles - tile_overlap_size)/2;

    std::vector<TiledPrediction> tiled_predictions;

    // 画像をタイルに切り分けてモデル検出にかける
    for(int h_idx = 0; h_idx < num_h_tiles; h_idx++){
        for(int w_idx = 0; w_idx < num_w_tiles; w_idx++){
            cv::Mat cropped;
            int upper_left_x = w_start + tile_slide*w_idx;
            int upper_left_y = h_start + tile_slide*h_idx;
            // 切り分けた画像データ
            cropped = img(cv::Rect{
                upper_left_x,
                upper_left_y,
                tile_size,
                tile_size
            });
            // モデル推論実行
            DetectionBboxesResults res = detect_faces(cascade, cropped);
            std::cout << "num of detections at h_idx:"<<h_idx<<"/w_idx:" << w_idx << " is " << res.bboxes.size() << std::endl;

            // 自分のユースケース都合の型に合わせる
            for(auto &one_box: res.bboxes){
                one_box.box[0] += float(upper_left_x);
                one_box.box[1] += float(upper_left_y);
                one_box.box[2] += float(upper_left_x);
                one_box.box[3] += float(upper_left_y);
            }
            TiledPrediction tlp{upper_left_x, upper_left_y, res};
            tiled_predictions.push_back(tlp);
        }
    }

    std::cout << tiled_predictions.size() << std::endl;

    DetectionBboxesResults full_detection_result = tiled_predictions[0].prediction;
    for(int i=1; i<tiled_predictions.size(); i++){
        full_detection_result.bboxes.insert(
            full_detection_result.bboxes.end(), 
            tiled_predictions[i].prediction.bboxes.begin(), 
            tiled_predictions[i].prediction.bboxes.end()
        );
    }

    // tile の枠を描画するなら
    for(auto pred: tiled_predictions){
        int x = pred.start_x;
        int y = pred.start_y;
        int x2 =pred.start_x + tile_size;
        int y2 =pred.start_y + tile_size;
        cv::rectangle(img, cv::Point(x,y), cv::Point(x2,y2), cv::Scalar(255,0,255), 2);
    }    

    std::cout << "num of detected faces: " << full_detection_result.bboxes.size() << std::endl;
    for(int i=0; i<full_detection_result.bboxes.size(); i++){
        std::vector<float> obj_coords = full_detection_result.bboxes[i].box;
        cv::rectangle(img, cv::Point(int(obj_coords[0]), int(obj_coords[1])),cv::Point(int(obj_coords[2]),int(obj_coords[3])), cv::Scalar(0,255,0), 2);
        cv::putText(img, std::to_string(full_detection_result.bboxes[i].score).substr(0,4), cv::Point(int(obj_coords[0]), int(obj_coords[1])), cv::FONT_HERSHEY_SIMPLEX, 2.5, cv::Scalar(0,0,0),3 );
    }   

    // ダブっているbboxの統合
    std::vector<int> residual_idx;
    std::vector<DetectedBbox> full_detection_result_bboxes;
    for(auto &bbox_result: full_detection_result.bboxes){
        full_detection_result_bboxes.push_back(
            {
                bbox_result.label,
                bbox_result.score,
                bbox_result.box
            }
        );
    }
    std::sort(full_detection_result_bboxes.begin(), full_detection_result_bboxes.end(), has_smaller_score);

    // 統合結果をdetection_result_without_overlap に移していく
    std::vector<DetectedBbox> detection_result_without_overlap;
    while(!full_detection_result_bboxes.empty()){
        auto res = full_detection_result_bboxes.back();
        residual_idx.push_back(full_detection_result_bboxes.size()-1);

        std::vector<int> delete_idx;
        // iouが一定値(今回は0.05)以上のものは、同一の顔の検出であると判断する
        for(int i=full_detection_result_bboxes.size()-2; i>=0; i--){
            auto other_box = full_detection_result_bboxes[i];
            float iou = calc_bbox_iou(res, other_box);
            std::cout << iou << std::endl;
            if(iou > 0.05) {
                // 本来の物体検出では, label の一致をみるべき。
                res.box[0] = std::min(res.box[0], other_box.box[0]);
                res.box[1] = std::min(res.box[1], other_box.box[1]);
                res.box[2] = std::max(res.box[2], other_box.box[2]);
                res.box[3] = std::max(res.box[3], other_box.box[3]);
                delete_idx.push_back(i);
            }            
        }
        detection_result_without_overlap.push_back(res);
        full_detection_result_bboxes.pop_back();

        for(int i=0; i<delete_idx.size(); i++){
            full_detection_result_bboxes.erase(full_detection_result_bboxes.begin()+delete_idx[i]);
        }
    }

    std::cout << "my nums num:" << detection_result_without_overlap.size() << std::endl;

    for(auto &detected_face: detection_result_without_overlap){
        cv::rectangle(img, cv::Point(int(detected_face.box[0]), int(detected_face.box[1])),cv::Point(int(detected_face.box[2]),int(detected_face.box[3])), cv::Scalar(0,0,255), 2);
    }



    cv::imwrite("./result.jpg", img);
    cv::resize(img, img, cv::Size(loaded_cols, loaded_rows));
    imshow("detect face", img);
    waitKey(0);
}

util.cpp

#include <vector>
#include <algorithm>
#include <string>
#include <numeric>
#include <iostream>
#include <cmath>


struct DetectedBbox {
    // label name index. dummy this time.
    int label;
    // confidence score. dummy this time.
    float score;
    // x0,y0,x1,y1
    std::vector<float> box;
};

bool has_smaller_score(DetectedBbox const &a, DetectedBbox const &b){
    return a.score < b.score;
}

float calc_bbox_iou(DetectedBbox const &a, DetectedBbox const &b){
    
    float a_area = (a.box[2]-a.box[0]) * (a.box[3] - a.box[1]);
    float b_area = (b.box[2]-b.box[0]) * (b.box[3] - b.box[1]);

    float i_xmin = std::max(a.box[0], b.box[0]);
    float i_ymin = std::max(a.box[1], b.box[1]);
    float i_xmax = std::min(a.box[2], b.box[2]);
    float i_ymax = std::min(a.box[3], b.box[3]);
    float i_w = std::max(0.0f, i_xmax - i_xmin);
    float i_h = std::max(0.0f, i_ymax - i_ymin);
    float i_area = i_w * i_h;
    float u_area = a_area + b_area - i_area;

    return i_area / u_area;
}


struct BoundingBox {
    int label;
    float score;
    std::vector<float> box; // x1, y1, x2, y2
};

推論対象の画像は、パスをコマンドライン引数で渡す形式です。

あからさまな課題

「ダブっているbboxの統合」のところで、ベクトルの要素の削除を行い、統合版を別のベクトルに新規作成している個所があります。これはあるいみfor文で処理しながらベクトルの中身を消す方法が分からなかったからなのですが、iteratorを回して erase or インクリメントの処理を書けばできるそうです。(参考)

補足

• SAHIリポジトリのデモでは、defaultとして「greedy_nmm」なる関数が動いているようです。上記の実装と比べて、スコア(confidence)でのソート処理等が異なっています。
• ソースコードに書いていますが、正方形タイルでの検出範囲切り分けではありません。注意。
• どうもSAHIのdefaultで動く部分を読んでいる限り、いずれの検出範囲（タイル）にも全体が入らない/含まれないような物体は、検出されても枠が二つに分かれてしまう可能性があるようです。統合するかのiou閾値が低ければ、基本的にはマージされるかと思いますが。
• また、ラベルのことなる検出結果で、bboxがかぶっている場合は、confidenceスコアが大きい方だと判断するつくりのようです。