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

时间:2020-11-07 00:07:31

标签: python opencv matplotlib

我有一些来自交通信号灯的图像:

Red traffic light Green traffic light Yellow traffic light

如何仅使用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()

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

2 个答案:

答案 0 :(得分:1)

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

因此,使用np.count_nonzero()对区域进行遮罩,对结果进行遮罩并计算非零像素的数量。

输入1(红色)

this

输入2(绿色)

enter image description here

输入3(黄色)

enter image description here

pkill -P

红色面具:

enter image description here

绿色面具:

enter image description here

红色直方图:

enter image description here

绿色直方图:

enter image description here

黄色直方图:

enter image description here

结果:

$$

答案 1 :(得分:1)

我的完整解决方案:(感谢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()


相关问题