图像像素梯度

# 编辑梯度损失函数
import numpy as np
import matplotlib.pyplot as plt
def grad(img):
    b, h, w, c = img.shape
    g_y=img[:,1:,:,:]-img[:,:-1,:,:]  #计算高度方向梯度
    g_x=img[:,:,1:,:]-img[:,:,:-1,:]  #计算宽度方向梯度
    g_y1,g_y2=g_y[:,:,:-1,:],g_y[:,:,-1,:]
    g_x1,g_x2=g_x[:,:-1,:,:],g_x[:,-1,:,:]
    y1=np.arctan(g_y1 / (g_x1 + 1e-8))
    y2=np.arctan(g_y2/1e-8)
    plt.imshow((np.abs(y1)/1.6).reshape(h-1, w-1, c))
    plt.show()
    #
    return (np.sum(np.abs(y1))+np.sum(np.abs(y2)))/(b*h*w*c)

if __name__ == '__main__':

    img=plt.imread('z.jpg')
    h, w, c = img.shape
    img=np.reshape(img,(1,h, w, c))
    result = grad(img)
    print(result)