white_balancer.py

import numpy as np
from py_profiler import profiler


class WhiteBalancer:

    @profiler()
    def image_agcwd(self, img, a=0.25, truncated_cdf=False):
        h, w = img.shape[:2]
        hist, bins = np.histogram(img.flatten(), 256, [0, 256])
        cdf = hist.cumsum()
        cdf_normalized = cdf / cdf.max()
        prob_normalized = hist / hist.sum()

        unique_intensity = np.unique(img)
        intensity_max = unique_intensity.max()
        intensity_min = unique_intensity.min()
        prob_min = prob_normalized.min()
        prob_max = prob_normalized.max()

        pn_temp = (prob_normalized - prob_min) / (prob_max - prob_min)
        pn_temp[pn_temp > 0] = prob_max * (pn_temp[pn_temp > 0] ** a)
        pn_temp[pn_temp < 0] = prob_max * (-((-pn_temp[pn_temp < 0]) ** a))
        prob_normalized_wd = pn_temp / pn_temp.sum()  # normalize to [0,1]
        cdf_prob_normalized_wd = prob_normalized_wd.cumsum()

        if truncated_cdf:
            inverse_cdf = np.maximum(0.5, 1 - cdf_prob_normalized_wd)
        else:
            inverse_cdf = 1 - cdf_prob_normalized_wd

        img_new = img.copy()
        for i in unique_intensity:
            img_new[img == i] = np.round(255 * (i / 255) ** inverse_cdf[i])

        return img_new

    @profiler()
    def process_bright(self, img):
        img_negative = 255 - img
        agcwd = self.image_agcwd(img_negative, a=0.25, truncated_cdf=False)
        reversed = 255 - agcwd
        return reversed

    @profiler()
    def process_dimmed(self, img):
        agcwd = self.image_agcwd(img, a=0.75, truncated_cdf=True)
        return agcwd