「HSV色空間での赤色・緑色の検出」で検出した赤色・緑色を基に信号機の色を判断します。画像で使われているすべての色に対し、k 平均法を適用して代表色を取得します。
色検出スクリプトの作成
信号機のアノテーション画像を使用して、赤色、緑色を色検出スクリプト「signaldetect」で判断します。
- 20行目の「kmeans」関数で「cv2.kmeans」関数を使って画像の代表色を取得します。
- 43行目の「visualization」関数でヒストグラムを表示します。83行目のコメントを外すことにより、ヒストグラムの表示が可能となります。
colorjudge.py
import os
import fnmatch
import csv
import cv2
import numpy as np
from IPython.display import Image, display
from matplotlib import pyplot as plt
INMASKIMAGE_DIR = 'out/'
COLOR_DIR = 'out/'
def imshow(img):
"""ndarray 配列をインラインで Notebook 上に表示する。
"""
ret, encoded = cv2.imencode(".jpg", img)
display(Image(encoded))
def kmeans(img):
# 画像で使用されている色一覧。(W * H, 3) の numpy 配列。
colors = img.reshape(-1, 3).astype(np.float32)
# クラスタ数
K = 5
# 最大反復回数: 10、移動量の閾値: 1.0
criteria = cv2.TERM_CRITERIA_MAX_ITER + cv2.TERM_CRITERIA_EPS, 10, 1.0
ret, labels, centers = cv2.kmeans(
colors, K, None, criteria, attempts=10, flags=cv2.KMEANS_RANDOM_CENTERS
)
print(f"ret: {ret:.2f}, label: {labels.shape}, center: {centers.shape}")
# ret: 127443.79220199585, label: (48380, 1), center: (8, 3)
labels = labels.squeeze(axis=1) # (N, 1) -> (N,)
centers = centers.astype(np.uint8) # float32 -> uint8
return labels, centers, K
def visualization(img, labels, centers, K):
# 各クラスタに属するサンプル数を計算する。
_, counts = np.unique(labels, axis=0, return_counts=True)
# 可視化する。
fig, [ax1, ax2] = plt.subplots(1, 2, figsize=(10, 3))
fig.subplots_adjust(wspace=0.5)
# matplotlib の引数の仕様上、[0, 1] にして、(R, G, B) の順番にする。
bar_color = centers[:, ::-1] / 255
bar_text = list(map(str, centers))
# 画像を表示する。
ax1.imshow(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
ax1.set_axis_off()
# ヒストグラムを表示する。
ax2.barh(np.arange(K), counts, color=bar_color, tick_label=bar_text)
print('arange:', np.arange(K), '\ncounts:', counts, ' \nbar_color:', bar_color, ' \nbar_text:', bar_text)
plt.show()
def main():
f = open('out.csv', 'w', newline="")
csv_data = ['ファイル名', 'count']
writer = csv.writer(f)
writer.writerow(['ファイル名', 'count'])
for fileName in fnmatch.filter(os.listdir(INMASKIMAGE_DIR), '*_img.jpg'):
print(fileName)
img = cv2.imread(INMASKIMAGE_DIR + fileName)
labels, centers, K = kmeans(img)
_, counts = np.unique(labels, axis=0, return_counts=True)
csv_data = [fileName, sorted(counts,reverse=True), sorted(counts,reverse=True)[1]]
writer.writerow(csv_data)
# visualization(img, labels, centers, K)
f.close()
if __name__ == "__main__":
main()
色検出スクリプトの実行
信号機の検出画像を入力として、赤色、緑色を色判断スクリプト「signaldetect」を実行します。次のように検出画像ごとの色の重みがCSV形式でファイル「out.csv」に保存されます。
signalA_green_0005_green_img.jpg ret: 71473.93, label: (4959, 1), center: (5, 3) signalA_green_0005_red_img.jpg ret: 8608.52, label: (4959, 1), center: (5, 3) signalA_green_0009_green_img.jpg ・・・ signalB_red_0231_red_img.jpg ret: 299920.61, label: (3306, 1), center: (5, 3) Process finished with exit code 0
csvファイル「out.csv」の内容を次に示します。2カラム目の「count」は検出された色ごとの重みを表します。一番大きい値はマスクを行ったためにマスクの黒色の重みを示します。次の重みがマスクした色に対応する赤色または緑色の重みを示します。二つのファイルの重みの大きい方が信号機の色となります。
ヒストグラム表示のコメントを外すと、次のようにヒストグラムの内容が表示されます。
signalA_green_0005_green_img.jpg ret: 71473.93, label: (4959, 1), center: (5, 3) arange: [0 1 2 3 4] counts: [ 41 4506 65 269 78] bar_color: [[0.10196078 0.38431373 0.30588235] [0. 0. 0. ] [0.02745098 0.21960784 0.16470588] [0.03921569 0.10196078 0.08627451] [0.11764706 0.14117647 0.1372549 ]] bar_text: ['[78 98 26]', '[0 0 0]', '[42 56 7]', '[22 26 10]', '[35 36 30]']
次のヒストグラムが表示されます。
