python颜色校正

def calibrateColor(img, refrgb, global_scale):
    # img = np.array(img)
    imgrgb = cv2.mean(img)

    # 1. 全局颜色校正
    if global_scale:
        scale = [refrgb[0] / imgrgb[0], refrgb[1] / imgrgb[1], refrgb[2] / imgrgb[2]]
        newImage = img.astype(np.float32)

        for i in range(3):
            newImage[:, :, i] = newImage[:, :, i] * scale[i]

        newImage = np.clip(newImage, 0, 255)
        newImage = newImage.astype(np.uint8)

    # 2. 分段scale
    else:
        newImage = img.astype(np.float32)
        for i in range(3):
            low_index = np.where(img[:, :, i] < imgrgb[i])
            hign_index = np.where(img[:, :, i] >= imgrgb[i])

            low_img = img[:, :, i].copy()
            high_img = img[:, :, i].copy()

            new_low_img = (low_img / imgrgb[i]) * refrgb[i]
            new_hign_img = ((high_img-imgrgb[i]) / (255-imgrgb[i])) * (255-refrgb[i]) + refrgb[i]

            new_low_img[hign_index] = 0
            new_hign_img[low_index] = 0

            newImage[:, :, i] = new_low_img + new_hign_img
        newImage = newImage.astype(np.uint8)

    return newImage