python实现简单的ps色阶调整过程

最近在做一个用unet神经网络识别种子出苗率的项目（种子贼小，分辨率还贼低），由于识别的效果不太好，只能对图像做预处理了。。。预处理选用的是ps中的色阶处理，可是cv2等图像库中没有现成的函数，只能自己一点一点搜资料，找到了ps中色阶调整的计算公式。
以下是ps中的色阶窗口

左边黑三角叫做黑场，中间叫做灰场，右边叫做白场,上面的图是分布直方图，表示每个色阶的数量，开始我天真的以为色阶调整是把低于黑场的像素值设置为0，高于白场的设为255就行，后来这么写出来后，发现事情并不简单，根本不是想要的效果，然后又去各种搜博客，终于，我发现了一位老哥的博客：https://blog.csdn.net/maozefa/article/details/43971063
太感谢这位老哥了，大家有公式不懂的可以去看看。
根据他给的公式，我用python实现了一下：

'''
    用于图像预处理，模拟ps的色阶调整
    img：传入的图片
    Highlight：白场(Shadow-255)
    Shadow：黑场(0-Highlight)
    0 <= Shadow < Highlight <= 255
    返回一张图片
'''
    def Whiite_Black_Level_Pretreatment(self, img, Shadow,Highlight):
        if Highlight > 255:
            Highlight = 255
        if Shadow < 0:
            Shadow = 0
        if Shadow >= Highlight:
            Shadow = Highlight - 2
        # 转类型
        img = np.array(img, dtype=int)
        # 计算白场黑场离差
        Diff = Highlight - Shadow
        # 计算系数
        coe = 255.0 / Diff
        rgbDiff = img - Shadow
        rgbDiff = np.maximum(rgbDiff, 0)
        img = rgbDiff * coe
        # 四舍五入到整数
        img = np.around(img, 0)
        # 变为int型
        img = img.astype(int)
        return img

注意：我这里没有关系到中间的灰场和输出色阶，我的灰场默认为1.0，如果需要用到灰场，需要实现这里：
效果：当黑场为81，白场为161时：
原图：

处理后：