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時：
原圖：

處理後：