直方圖均衡化-python實作

"""
@author: LiShiHang
@software: PyCharm
@file: 5.1.直方圖均衡化.py
@time: 2018/12/24 16:02
@desc:
"""
import cv2 # 僅用于讀取圖像矩陣
import matplotlib.pyplot as plt  
import numpy as np

gray_level = 256  # 灰階級


def pixel_probability(img):
    """
    計算像素值出現機率
    :param img:
    :return:
    """
    assert isinstance(img, np.ndarray)

    prob = np.zeros(shape=(256))

    for rv in img:
        for cv in rv:
            prob[cv] += 1

    r, c = img.shape
    prob = prob / (r * c)

    return prob


def probability_to_histogram(img, prob):
    """
    根據像素機率将原始圖像直方圖均衡化
    :param img:
    :param prob:
    :return: 直方圖均衡化後的圖像
    """
    prob = np.cumsum(prob)  # 累計機率

    img_map = [int(i * prob[i]) for i in range(256)]  # 像素值映射

   # 像素值替換
    assert isinstance(img, np.ndarray)
    r, c = img.shape
    for ri in range(r):
        for ci in range(c):
            img[ri, ci] = img_map[img[ri, ci]]

    return img


def plot(y, name):
    """
    畫直方圖，len(y)==gray_level
    :param y: 機率值
    :param name:
    :return:
    """
    plt.figure(num=name)
    plt.bar([i for i in range(gray_level)], y, width=1)


if __name__ == '__main__':

    img = cv2.imread("source.jpg", 0)  # 讀取灰階圖

    prob = pixel_probability(img)
    plot(prob, "原圖直方圖")

    # 直方圖均衡化
    img = probability_to_histogram(img, prob)
    cv2.imwrite("source_hist.jpg", img)  # 儲存圖像

    prob = pixel_probability(img)
    plot(prob, "直方圖均衡化結果")

    plt.show()