DarkNetのYOLO
水中ロボコンでQRコードの検出のためにYOLOを使うことになったのですが、最初はnnablaというライブラリで学習を行おうとしたもののうまくいかず、本家DarkNetのYOLOを使うことのになりました。YOLOで自作のデータセットを学習させる方法はいくらでも見つかると思うのでココでは割愛します。
YOLOはCで書かれていますが、私が去年開発した、水中ロボコン用の遠隔通信ソフトは、C++とOpenCVで書かれていたため、DarkNetがC＋＋で書かれていることは、むしろ好都合でした。
本記事はで私のようにOpenCVで書かれたプログラムにDarkNetのYOLOを組み込みたい人のために、OpenCVでwebカメラから取得した画像をDarkNetで推論する例を示したいと思います。
YOLOの起動に必要なもの
YOLOを起動するためには以下のようなファイルが必要です。
- データ設定ファイル(*.data)
- 設定ファイル（*.cfg）
- 学習済みパラメータ(.weights )
まず学習済みのパラメータが必要です。これは自ら学習したものを利用しても良いですし、デフォルトとのものを使っても良いです。もう一つ重要なのが、YOLOに入力する画像のサイズです。コレも学習時のパラメータと一致させて下さい。 ファイルの読み込みは次のような関数をプログラムの起動時に呼び出しました。
list *options = read_data_cfg("自分のデータ設定ファル.data"); char *name_list = option_find_str(options, "names", "data/names.list"); char **names = get_labels(name_list); image **alphabet = load_alphabet(); network *net = load_network("自分の設定ファイル.cfg", "自分の学習済みでーた.weights", 0);
学習済のパラメータに関しては公式のドキュメントを見てもらうのが一番良いですが、僕は以下の記事を参照しました。
OpenCVからDarkNetのimgへの変換
OpenCVとDarkNetの画像形式は異なるため、cv::Matからdarknetまで変換する必要があります。DarkNetはImageという独自の構造体を持っているため、OpenCVをimage型に変換してあげる必要があります。
この変換は先人がやってくださっているので以下の記事を参照して下さい
僕はこの記事に出てきた変換を次のような関数にまとめました。
image Mat2Image(cv::Mat src)
{
cv::Mat FMat;
src.convertTo(FMat, CV_32FC3, 1.0 / 255);
std::vector<cv::Mat> tmp;
cv::split(FMat, tmp);
int size = FMat.size().width * FMat.size().height;
int fsize = size * sizeof(float);
image retval = make_image(FMat.size().width, FMat.size().height, 3);
float *p = retval.data;
memcpy((unsigned char *)p, tmp[2].data, fsize);
p += size;
memcpy((unsigned char *)p, tmp[1].data, fsize);
p += size;
memcpy((unsigned char *)p, tmp[0].data, fsize);
return retval;
}
推論
darknetの推論はnetwork_predict(net, X);という関数で行われます。netの方にニューラルネットワーク、Xの方に推論対象であるimageを入力します。画像サイズは自ら学習した画像のサイズに適宜変えて下さい、それと僕の学習したネットワークがモノクロを推論対象としていたため、画像をモノクロに変換する処理が入ってしまっていますがコレも必要に応じて変更して下さい。
cv::Mat feed(frame, rect);
cv::resize(feed, feed, cv::Size(), 506.0 / feed.rows, 506.0 / feed.rows);
image feed_img = Mat2Image(feed);
float *X = feed_img.data;
layer l = net->layers[net->n - 1];
network_predict(net, X);
int nboxes = 0;
detection *dets = get_network_boxes(net, feed_img.w, feed_img.h, thresh, hier_thresh, 0, 1, &nboxes);
printf("w is %d,h is%d, th is %f , hth is %f,box is%d\n", feed_img.w, feed_img.h, thresh, hier_thresh, nboxes);
実行
//gcc -Iinclude/ -Isrc/ -DOPENCV `pkg-config --cflags opencv` main.cpp -DGPU -I/usr/local/cuda/include/ -DCUDNN -Wall -Wno-unused-result -Wno-unknown-pragmas -Wfatal-errors -fPIC -Ofast -DOPENCV -DGPU -DCUDNN libdarknet.a -o darknet -lm -pthread `pkg-config --libs opencv` -lstdc++ -L/usr/local/cuda/lib64 -lcuda -lcudart -lcublas -lcurand -lcudnn -lstdc++ libdarknet.a #include <iostream> #include <string> #include <sstream> #include <iomanip> #include <stdexcept> #include <vector> #include <cmath> #include <thread> #include <boost/lockfree/queue.hpp> #include <opencv2/core/utility.hpp> #include <opencv2/core/core.hpp> #include <opencv2/highgui/highgui.hpp> #include <opencv2/imgproc/imgproc.hpp> #include <opencv2/calib3d/calib3d.hpp> #include <opencv2/stitching.hpp> #include <omp.h> #include <time.h> #include "include/darknet.h" image Mat2Image(cv::Mat src) { cv::Mat FMat; src.convertTo(FMat, CV_32FC3, 1.0 / 255); std::vector<cv::Mat> tmp; cv::split(FMat, tmp); int size = FMat.size().width * FMat.size().height; int fsize = size * sizeof(float); image retval = make_image(FMat.size().width, FMat.size().height, 3); float *p = retval.data; memcpy((unsigned char *)p, tmp[2].data, fsize); p += size; memcpy((unsigned char *)p, tmp[1].data, fsize); p += size; memcpy((unsigned char *)p, tmp[0].data, fsize); return retval; } cv::Mat draw_detections_mat(image im, cv::Mat src, detection *dets, int num, float thresh, char **names, image **alphabet, int classes) { int i, j; cv::Mat dst; for (i = 0; i < num; ++i) { char labelstr[4096] = {0}; int cls = -1; for (j = 0; j < classes; ++j) { if (dets[i].prob[j] > thresh) { if (cls < 0) { strcat(labelstr, names[j]); cls = j; } else { strcat(labelstr, ", "); strcat(labelstr, names[j]); } printf("%s: %.0f%%\n", names[j], dets[i].prob[j] * 100); } } if (cls >= 0) { int width = im.h * .006; /* if(0){ width = pow(prob, 1./2.)*10+1; alphabet = 0; } */ //printf("%d %s: %.0f%%\n", i, names[class], prob*100); int offset = cls * 123457 % classes; /* float red = get_color(2, offset, classes); float green = get_color(1, offset, classes); float blue = get_color(0, offset, classes); */ float rgb[3]; //width = prob*20+2; rgb[0] = 0; rgb[1] = 0; rgb[2] = 0; box b = dets[i].bbox; //printf("%f %f %f %f\n", b.x, b.y, b.w, b.h); int left = (b.x - b.w / 2.) * im.w; int right = (b.x + b.w / 2.) * im.w; int top = (b.y - b.h / 2.) * im.h; int bot = (b.y + b.h / 2.) * im.h; if (left < 0) left = 0; if (right > im.w - 1) right = im.w - 1; if (top < 0) top = 0; if (bot > im.h - 1) bot = im.h - 1; cv::rectangle(src, cv::Point(left, top), cv::Point(right, bot), cv::Scalar(255, 255, 0), 5, 8); } } return src; } int main(void) { std::cout << "start" << std::endl; float thresh = 0.9; float hier_thresh = 0.5; cv::VideoCapture streaming(0); cv::Mat frame; list *options = read_data_cfg("cfg/task/datasets.data"); char *name_list = option_find_str(options, "names", "data/names.list"); char **names = get_labels(name_list); image **alphabet = load_alphabet(); network *net = load_network("cfg/task/yolov3-voc.cfg", "cfg/task/backup/yolov3-voc.backup", 0); set_batch_network(net, 1); srand(2222222); double time; char buff[256]; char *input = buff; float nms = .45; streaming >> frame; std::cout << frame.cols << ":" << frame.rows << std::endl; cv::Rect rect = cv::Rect((frame.cols / 2) - frame.rows / 2, 0, frame.rows, frame.rows); while (1) { std::cout << "running" << std::endl; streaming >> frame; cv::cvtColor(frame, frame, cv::COLOR_RGB2GRAY); cv::cvtColor(frame, frame, cv::COLOR_GRAY2RGB); cv::Mat feed(frame, rect); cv::resize(feed, feed, cv::Size(), 506.0 / feed.rows, 506.0 / feed.rows); image feed_img = Mat2Image(feed); float *X = feed_img.data; layer l = net->layers[net->n - 1]; network_predict(net, X); int nboxes = 0; detection *dets = get_network_boxes(net, feed_img.w, feed_img.h, thresh, hier_thresh, 0, 1, &nboxes); printf("w is %d,h is%d, th is %f , hth is %f,box is%d\n", feed_img.w, feed_img.h, thresh, hier_thresh, nboxes); if (nms) do_nms_sort(dets, nboxes, l.classes, nms); draw_detections(feed_img, dets, nboxes, thresh, names, alphabet, l.classes); draw_detections_mat(feed_img, feed, dets, nboxes, thresh, names, alphabet, l.classes); cv::imshow("test", feed); free_detections(dets, nboxes); int key = cv::waitKey(1); if (key == 'q') { break; } free_image(feed_img); } return 0; }
