如何使用 OpenCV 使用直方图信息检测交通灯颜色？

Question

我有一些交通信号灯的图像：

如何仅使用 OpenCV 的颜色直方图信息知道它们的颜色？ 我有以下内容：

for img_file in glob.glob(f"{save_folder}/*"):
        img = cv2.imread(img_file)
        for i, col in enumerate(("b", "g", "r")):
            histr = cv2.calcHist([img], [i], None, [256], [0, 256])
            plt.plot(histr, color=col)
            plt.xlim([0, 256])
        plt.show()

上面的代码绘制了每个图像的颜色直方图。我真的不知道如何从这里开始。

Answer 1

这是在 Python/OpenCV 中使用 G 与 R 2D 直方图进行判断的一种方法。因此，如果二维直方图右上角有很多白色，则为红色，如果沿左下角，则为绿色，如果仅沿对角线，则为黄色。

因此，为区域制作掩码，对结果进行掩码并使用 np.count_nonzero() 计算非零像素的数量。

输入 1（红色）

输入 2（绿色）

输入 3（黄色）

import cv2
import numpy as np
import skimage.exposure as exposure
import os

filenames = ['traffic_light_red.jpg', 'traffic_light_green.jpg', 'traffic_light_yellow.jpg']

for filename in filenames:

    print(filename)

    # get name without suffix
    name = os.path.splitext(filename)[0]

    # read image
    img = cv2.imread(filename)

    # calculate 2D histograms for pairs of channels: GR
    histGR = cv2.calcHist([img], [1, 2], None, [256, 256], [0, 256, 0, 256])

    # histogram is float and counts need to be scale to range 0 to 255
    histScaled = exposure.rescale_intensity(histGR, in_range=(0,1), out_range=(0,255)).clip(0,255).astype(np.uint8)

    # make masks
    ww = 256
    hh = 256
    ww13 = ww // 3
    ww23 = 2 * ww13
    hh13 = hh // 3
    hh23 = 2 * hh13
    black = np.zeros_like(histScaled, dtype=np.uint8)
    # specify points in OpenCV x,y format
    ptsUR = np.array( [[[ww13,0],[ww-1,hh23],[ww-1,0]]], dtype=np.int32 )
    redMask = black.copy()
    cv2.fillPoly(redMask, ptsUR, (255,255,255))
    ptsBL = np.array( [[[0,hh13],[ww23,hh-1],[0,hh-1]]], dtype=np.int32 )
    greenMask = black.copy()
    cv2.fillPoly(greenMask, ptsBL, (255,255,255))

    #Test histogram against masks
    region = cv2.bitwise_and(histScaled,histScaled,mask=redMask)
    redCount = np.count_nonzero(region)
    region = cv2.bitwise_and(histScaled,histScaled,mask=greenMask)
    greenCount = np.count_nonzero(region)
    print('redCount:',redCount)
    print('greenCount:',greenCount)

    # Find color
    threshCount = 100
    if redCount > greenCount and redCount > threshCount:
        color = "red"
    elif greenCount > redCount and greenCount > threshCount:
        color = "green"
    elif redCount < threshCount and greenCount < threshCount:
        color = "yellow"
    else:
        color = "other"
    print("color: ",color)  

    # save result
    cv2.imwrite(name + '_histogram.jpg', histScaled)
    cv2.imwrite('redMask.jpg', redMask)
    cv2.imwrite('greenMask.jpg', greenMask)

    # view results
    cv2.imshow("hist", histScaled)
    cv2.imshow("redMask", redMask)
    cv2.imshow("greenMask", greenMask)
    cv2.waitKey(0)
    cv2.destroyAllWindows()

    print('')

红色面具：

绿色面具：

红色直方图：

绿色直方图：

黄色直方图：

结果：

traffic_light_red.jpg
redCount: 2627
greenCount: 0
color:  red

traffic_light_green.jpg
redCount: 0
greenCount: 1138
color:  green

traffic_light_yellow.jpg
redCount: 0
greenCount: 0
color:  yellow

Answer 2

我的完整解决方案：（感谢fmw42）

这可能不是最佳的，但它是一个很好的起点。

# Dependencies
import cv2
import glob
import numpy as np
import skimage.exposure as exposure

def get_mid_arr(arr: np.ndarray, k: int) -> np.ndarray:
    mid_arr = arr.copy()
    upper = np.triu_indices(mid_arr.shape[0], k=k)
    mid_arr[upper] = 0
    lower = np.tril_indices(mid_arr.shape[0], k=-k)
    mid_arr[lower] = 0
    return mid_arr


def get_upper_arr(arr, k: int) -> np.ndarray:
    upper_arr = arr.copy()
    lower_triangle_indices = np.tril_indices(upper_arr.shape[0], k= k - 1)

    upper_arr[lower_triangle_indices] = 0

    return upper_arr


def get_lower_arr(arr, k: int) -> np.ndarray:
    lower_arr = arr.copy()
    upper_triangle_indices = np.triu_indices(lower_arr.shape[0], k = - k + 1)
    lower_arr[upper_triangle_indices] = 0
    return lower_arr


class TraficLightHistogramClassifier:
    
    def __init__(
        self,
        hist_cutting_treshold=55,  # optimal value found for my dataset
        probability_boundary=0.09,  # optimal value found for my dataset
    ) -> None:
        self.hist_cutting_treshold = hist_cutting_treshold
        self.probability_boundary = probability_boundary

    
    def predict(self, images_folder: str) -> list[str]:
        # Create empty list for holding predictions
        predictions = []
        # Search every image in the save folder
        for img_file in glob.glob(f"{images_folder}/*"):
            # Read image
            img = cv2.imread(img_file)
            # calculate 2D histograms for pairs of channels: GR
            hist = cv2.calcHist([img], [1, 2], None, [256, 256], [0, 256, 0, 256])
            # hist is float and counts need to be scale to range 0 to 255
            scaled_hist = (
                exposure.rescale_intensity(hist, in_range=(0, 1), out_range=(0, 255))
                .clip(0, 255)
                .astype(np.float64)
            )

            # Split histogram into 3 regions
            (yellow_region, green_region, red_region) = (
                get_mid_arr(scaled_hist, self.hist_cutting_treshold),
                get_lower_arr(scaled_hist, self.hist_cutting_treshold),
                get_upper_arr(scaled_hist, self.hist_cutting_treshold),
            )

            # Count how many non zero values in each region
            (red_count, green_count, yellow_count) = (
                np.count_nonzero(red_region),
                np.count_nonzero(green_region),
                np.count_nonzero(yellow_region),
            )

            # Calculate total non-zero values
            total_count = red_count + green_count + yellow_count

            # Calculate red and green percentage
            red_percentage, green_percentage = (
                red_count / total_count,
                green_count / total_count,
            )

            # Logic for deciding color
            if green_percentage > self.probability_boundary:
                predict = "green"
            elif red_percentage > self.probability_boundary:
                predict = "red"
            else:
                predict = "yellow"

            # Append to predictions
            predictions.append(predict)
        return predictions


def main():

    y_true = [
        "green",
        "green",
        "green",
        "green",
        "red",
        "red",
        "yellow",
        "green",
        "green",
        "red",
        "red",
        "red",
        "red",
        "red",
        "green",
        "red",
        "red",
        "yellow",
        "green",
    ]

    # Create classifier
    clf = TraficLightHistogramClassifier()

    # "Predict" lights
    y_pred = clf.predict("Cropped Images")

    # Print true "labels"
    print(y_true)
    # Print predict "labels"
    print(y_pred)


if __name__ == "__main__":
    main()