手动实现直方图匹配（python）

# -*- coding: utf-8 -*-
"""
Created on Fri Jun 11 06:24:33 2021

@author: 陨星落云
"""
import imageio
import cupy as cp
import numpy as np
import matplotlib.pylab as plt
import matplotlib
matplotlib.rcParams['font.family']='SimHei'
import time 


def compute_histogram(img):
    # 计算直方图
    tab = np.zeros(256)
    for i in np.unique(img):
        # print(i)
        tab[i] = np.bincount(img.flatten())[i]
    return tab.astype(np.uint64).tolist()

def hist_matching_cpu(img,target):
    # 直方图匹配CPU版 img:原图,targettraget：匹配的图像
    hist0,_ = np.histogram(target.flatten(),256,range=(0,255))
    cdf_dest = hist0.cumsum()/target.size

    imhist,bins = np.histogram(img.flatten(),len(cdf_dest),density=True)
    cdf = imhist.cumsum()
    cdf = cdf/cdf[-1]
    im2 = np.interp(img.flatten(),bins[:-1],cdf)
    im3 = np.interp(im2,cdf_dest,bins[:-1])
    imres = im3.reshape(img.shape[:2])
    return (imres*255).astype(np.uint8) 


def hist_matching_gpu(img,target):
    # 直方图匹配GPU版
    img,target = cp.array(img),cp.array(target)        
    hist0,_ = cp.histogram(target.flatten(),256,range=(0,255))
    cdf_dest = hist0.cumsum()/target.size

    imhist,bins = np.histogram(img.flatten(),len(cdf_dest),density=True)
    cdf = imhist.cumsum()
    cdf = cdf/cdf[-1]
    im2 = cp.interp(img.flatten(),bins[:-1],cdf)
    im3 = cp.interp(im2,cdf_dest,bins[:-1])
    imres = cp.asnumpy(im3.reshape(img.shape[:2]))
    return (imres*255).astype(np.uint8) 

if __name__ == "__main__":

    # 读取图像
    img = imageio.imread("Unequalized.jpg")
    # target = (np.ones(100).reshape(10,10)*205).astype(np.uint8)
    target = imageio.imread("gray_lena.jpg")

    # 直方图匹配CPU版
    t0 = time.time()
    img1 = hist_matching_cpu(img,target)
    print("直方图匹配CPU所需时间：",time.time()-t0)

    # 直方图匹配GPU版
    t1 = time.time()
    img2 = hist_matching_gpu(img,target)

    print("直方图匹配GPU所需时间：",time.time()-t1)

    # 显示图像
    plt.figure(figsize=(20,16))

    plt.subplot(231)
    plt.imshow(img,cmap='gray')
    plt.title('原图')
    plt.subplot(232)
    plt.imshow(target,cmap='gray')
    plt.title('参考图')
    plt.subplot(233)
    plt.imshow(img2,cmap='gray')
    plt.title('结果图')

    plt.subplot(234)
    # print(compute_histogram(img))
    plt.plot(compute_histogram(img))
    plt.title('原图的直方图')
    plt.subplot(235)
    # print(compute_histogram(image_stretch(img)))
    plt.plot(compute_histogram(target))
    plt.title('参考的直方图')
    plt.subplot(236)
    # print(compute_histogram(image_stretch(img)))
    plt.plot(compute_histogram(img2))
    plt.title('结果的直方图')
    plt.show()

结果：

直方图匹配CPU所需时间： 0.05585050582885742
直方图匹配GPU所需时间： 0.0329129695892334