Python實現圖像直方圖均衡化算法

效果圖

代碼

#!/usr/bin/env python3
# coding=utf-8

import matplotlib.image as mpimg
from matplotlib import pyplot as plt
import sys
import numpy as np


def equalization(gray_value):
    """
    傳入灰度值，對灰度值作均衡化，不須要返回，直接修改傳入的參數
    :param gray_value:
    """
    # 統計灰度直方圖
    gray = np.zeros(256)
    row, column = gray_value.shape
    for i in range(row):
        for j in range(column):
            gray[gray_value[i][j]] += 1

    # 計算灰度佔比
    gray /= (row * column)
    # 顯示灰度直方圖
    plt.subplot(2, 2, 2)
    plt.plot(gray)

    cumsum = np.cumsum(gray)  # 計算累積和

    # 均衡化
    # equa_t[i]=j表示原灰度值i通過均衡化後轉化爲灰度值j
    # 255×累積和四捨五入爲int型
    equa_t = np.array((255 * cumsum + 0.5)).astype(np.int32)
    # 統計均衡化後的灰度數量
    equa_gray = np.zeros(256)
    for i in range(256):
        equa_gray[equa_t[i]] += gray[i]
    # 顯示均衡化後的直方圖
    plt.subplot(2, 2, 4)
    plt.plot(equa_gray)
    # 對原灰度矩陣作均衡化
    for i in range(row):
        for j in range(column):
            gray_value[i][j] = equa_t[gray_value[i][j]]


def run(img_path):
    img_array = mpimg.imread(img_path)
    plt.subplot(2, 2, 1)
    plt.imshow(img_array)
    img_array *= 255
    img_array = img_array.astype(np.int32)
    equalization(img_array[:, :, 0])
    equalization(img_array[:, :, 1])
    equalization(img_array[:, :, 2])
    img_array = img_array.astype(np.float64)
    img_array /= 255
    plt.subplot(2, 2, 3)
    plt.imshow(img_array)


if __name__ == "__main__":
    if sys.argv.__len__() <= 1:
        png = input("請輸入要處理的圖片名:\n")
    else:
        png = sys.argv[1]
    run(png)
    plt.show()