信号機が何色かを判別してみた。

おじさん、老後の趣味の一つで勝手に自動運転の仕組みを勉強しています。
自動運転って仕組みを一つ一つ理解していけばいつかできそうですね。

前回は信号機自体がどこにあるのかを認識しました。
信号機までの距離は認識した信号機の大きさでたぶんわかるので省略。
今回はその認識した信号の色を調べます。

信号の色ってどうやってプログラムで識別するんだろう。
いろいろ調べてみると、どうやら色空間をHSVに変換して、HSV空間の特定の領域に入っているかどうかを調べればよいみたいです。

映画やテレビでよく使われているクロマキーとかと同じ手法ですね。
特定の色のところだけを抜き出して、その割合でその色があるかどうか判別します。

https://algorithm.joho.info/programming/python/opencv-color-tracking-py/

信号機の色判別なんてどこにもソースコードがないので今回上のサイトをもとに自分でごりごり書きます。

こんなかんじです。

1.前回認識した信号機の部分

2.赤色の部分をマスク

3.黄色の部分をマスク

4.緑の部分をマスク

HSV空間の特定の領域にある部分だけをマスクすると、このように信号機が緑色であることがわかります。作るのもとっても簡単。

--------------------------
#include <stdio.h>
#include <vector>

#define STB_IMAGE_IMPLEMENTATION
#define STB_IMAGE_WRITE_IMPLEMENTATION

#include "core.hpp"
#include "imgproc.hpp"
#include "objdetect.hpp"

#include "stb_image.h"
#include "stb_image_write.h"

#define X_DIV 10
#define Y_DIV 10



using namespace std;



void changeb_g(unsigned char* p, int x, int y, int c)
{
int ct = x*y;
int i;
unsigned char t;
for (i = 0; i < ct; i++) {
t = p[0];
p[0] = p[2];
p[2] = t;
p[3] = 255;
p += c;
}
}


static double calc_mask(unsigned char* buf, int x, int y, int s)
{
int i, c;
int ct = 0;
double d;
c = x*y;
for (i = 0; i < c; i++) {
if (*buf >= 128)ct++;
buf += s;
}
d = (double)ct / (double)c*100.0;
//printf("ct=%d c=%d %f\n",ct,c,d);
return d;
}

int get_traffic_light_color(const char* fn)
{
unsigned char*p;
int x = -1, y = -1;
int n, m;
int i;
double dr, dg, dy;

//load image
p = stbi_load(fn, &x, &y, &n, 4);
if (p == NULL || x < 1 || y < 1)return 0;

// R<-->B
changeb_g(p, x, y, 4);

cv::Mat color = cv::Mat(y, x, CV_8UC4);
cv::Mat hsv;
cv::Mat mask_r1, mask_r2, mask_r, color_r;
cv::Mat mask_g, color_g;
cv::Mat mask_y, color_y;

//copy data
memcpy(color.data, p, x * 4 * y);
stbi_image_free(p);

cv::cvtColor(color, hsv, CV_BGR2HSV);

//r
cv::inRange(hsv, cv::Scalar(0, 128, 0), cv::Scalar(10, 255, 255), mask_r1);
cv::inRange(hsv, cv::Scalar(150, 128, 0), cv::Scalar(179, 255, 255), mask_r2);
cv::bitwise_or(mask_r1, mask_r2, mask_r);
cv::cvtColor(mask_r, color_r, CV_GRAY2RGBA);
stbi_write_png("result_r.png", x, y, 4, color_r.data, 4 * x);
dr=calc_mask(color_r.data, x, y, 4);
//printf("red=%f\n", dr);

//y
cv::inRange(hsv, cv::Scalar(10, 64, 0), cv::Scalar(40, 255, 255), mask_y);
cv::cvtColor(mask_y, color_y, CV_GRAY2RGBA);
stbi_write_png("result_y.png", x, y, 4, color_y.data, 4 * x);
dy = calc_mask(color_y.data, x, y, 4);
//printf("yellow=%f\n", dy);

//g
cv::inRange(hsv, cv::Scalar(40, 128, 0), cv::Scalar(100, 255, 255), mask_g);
cv::cvtColor(mask_g, color_g, CV_GRAY2RGBA);
stbi_write_png("result_g.png", x, y, 4, color_g.data, 4 * x);
dg = calc_mask(color_g.data, x, y, 4);
//printf("green=%f\n", dg);

if (dr >= dy && dr >= dg)return 1; //red
if (dy >= dr && dy >= dg)return 2; //yellow
if (dg >= dy && dg >= dr)return 3; //green

return 0;
}


int main()
{
unsigned char*p;
int i;
int color;

char* fn_list[] = {
"red1.jpg",
"yellow1.jpg",
"green1.jpg",
"green2.jpg",
NULL
};
char* color_list[] =
{
"black",
"red",
"yellow",
"green",
};

for (i = 0; i < 10000; i++) {
if (fn_list[i] == NULL)break;
color = get_traffic_light_color(fn_list[i]);
printf("color=%s  file=%s  \n",color_list[color],fn_list[i]);
}

return 0;
}

--------------------------

信号機の色判別できた！