彩色图像的直方图python实现

彩色图像的直方图python实现

import torch
from torch import nn, sigmoid
from IPython.display import Image
from IPython.core.display import HTML 
import cv2
import matplotlib.pyplot as plt
import numpy as np

需要用到的函数

def np_to_torch(img, add_batch_dim=True):
    img = np.asarray(img)
    img = img.astype(np.float32).transpose((2, 0, 1))
    if add_batch_dim:
        img = img[np.newaxis, ...]
    img_torch = torch.from_numpy(img) / 255.0
    return img_torch
def phi_k(x, L, W):
    return sigmoid((x + (L / 2)) / W) - sigmoid((x - (L / 2)) / W)
def phi_step(x, L, W):
    print(x.shape)
    z1 = (x + (L / 2)) / W
    z2 = (x - (L / 2)) / W
    return sigmoid(1000*z1) - sigmoid(1000*z2)
def compute_pj(x, mu_k, K, L, W, kernel):
    # we assume that x has only one channel already
    # flatten spatial dims
    x = x.reshape(x.size(0), 1, -1)#(BATCH,C,H*W)
    
    x = x.repeat(1, K, 1)  # construct K channels
#     print(mu_k.shape)
    if kernel == "step":        
        return phi_step(x - mu_k, L, W)
    elif kernel == "sigmoid":
        return phi_k(x - mu_k, L, W)

主函数

class SingleDimHistLayer(nn.Module):
    def __init__(self, kernel):
        super().__init__()
        self.kernel = kernel
        self.K = 256
        self.L = 1 / self.K  # 2 / K -> if values in [-1,1] (Paper)
        self.W = self.L / 2.5
        #生成每个中点的值，将区间256等分后
        self.mu_k = (self.L * (torch.arange(self.K) + 0.5)).view(-1, 1).cuda()

    def forward(self, x):
        N = x.size(1) * x.size(2)
        print(x.shape)
      
        pj = compute_pj(x, self.mu_k, self.K, self.L, self.W, self.kernel)
    
        return pj.sum(dim=2) / N

效果展示：

img_src = cv2.imread("./1.png")
img_src = np_to_torch(img_src).cuda()
hist_layer = SingleDimHistLayer("step")
SS = hist_layer(img_src)
_, ax = plt.subplots(2)
ax[0].plot(SS[0].cpu().numpy())
# ax[1].plot(np.histogram(source[:, 0].view(-1).cpu().numpy(), bins=256)[0])
plt.show()

输入图像为lena，直方图为: