OzUVGL/ips/wb.py

import numpy as np
from utils.color import rgb2ycc

def illumination_parameters_estimation(current_image, illumination_estimation_option):
    ie_method = illumination_estimation_option.lower()

    if ie_method == "gw":
        ie = np.mean(current_image, axis=(0, 1))
        ie /= ie[1]
        return ie
    elif ie_method == "sog":
        sog_p = 4.
        ie = np.mean(current_image**sog_p, axis=(0, 1))**(1/sog_p)
        ie /= ie[1]
        return ie
    elif ie_method == "wp":
        ie = np.max(current_image, axis=(0, 1))
        ie /= ie[1]
        return ie
    elif ie_method == "iwp":
        samples_count = 10
        sample_size = 10
        rows, cols = current_image.shape[:2]
        data = np.reshape(current_image, (rows*cols, 3))
        maxima = np.zeros((samples_count, 3))
        for i in range(samples_count):
            maxima[i, :] = np.max(data[np.random.randint(low=0, high=rows*cols, size=(sample_size)), :], axis=0)
        ie = np.mean(maxima, axis=0)
        ie /= ie[1]
        return ie
    else:
        raise ValueError(
            'Bad illumination_estimation_option value! Use the following options: "gw", "wp", "sog", "iwp"')
      
      
def ratios2floats(ratios):
    floats = []
    for ratio in ratios:
        floats.append(float(ratio.num) / ratio.den)
    return floats