shinestacker 1.1.0__py3-none-any.whl → 1.2.0__py3-none-any.whl

shinestacker/algorithms/balance.py

@@ -1,4 +1,5 @@
-# pylint: disable=C0114, C0115, C0116, E1101, R0902, E1128, E0606, W0640, R0913, R0917
+# pylint: disable=C0114, C0115, C0116, E1101, R0902, E1128, E0606, W0640, R0913, R0917, R0914
+import math
 import numpy as np
 import cv2
 import matplotlib.pyplot as plt
@@ -7,10 +8,188 @@ from scipy.interpolate import interp1d
 from .. config.constants import constants
 from .. core.exceptions import InvalidOptionError
 from .. core.colors import color_str
-from .utils import read_img, save_plot, img_subsample
+from .utils import (read_img, save_plot, img_subsample, bgr_to_hsv, bgr_to_hls,
+                    hsv_to_bgr, hls_to_bgr, bgr_to_lab, lab_to_bgr)
 from .stack_framework import SubAction
 
 
+class BaseHistogrammer:
+    def __init__(self, dtype, num_pixel_values, max_pixel_value, channels,
+                 plot_histograms, plot_summary, process=None):
+        self.dtype = dtype
+        self.num_pixel_values = num_pixel_values
+        self.max_pixel_value = max_pixel_value
+        self.channels = channels
+        self.plot_histograms = plot_histograms
+        self.plot_summary = plot_summary
+        self.process = process
+        self.corrections = None
+        self.figsize = (10, 5)
+
+    def begin(self, size):
+        self.corrections = np.ones((size, self.channels))
+
+    def add_correction(self, idx, correction):
+        if idx != self.process.ref_idx:
+            self.corrections[idx] = correction
+
+    def histo_plot(self, ax, hist, x_label, color, alpha=1):
+        ax.set_ylabel("# of pixels")
+        ax.set_xlabel(x_label)
+        ax.set_xlim([0, self.max_pixel_value])
+        ax.set_yscale('log')
+        x_values = np.linspace(0, self.max_pixel_value, len(hist))
+        ax.plot(x_values, hist, color=color, alpha=alpha)
+
+    def save_plot(self, idx):
+        idx_str = f"{idx:04d}"
+        plot_path = f"{self.process.working_path}/{self.process.plot_path}/" \
+                    f"{self.process.name}-hist-{idx_str}.pdf"
+        save_plot(plot_path)
+        plt.close('all')
+        self.process.callback(
+            'save_plot',
+            self.process.id, f"{self.process.name}: balance\nframe {idx_str}",
+            plot_path
+        )
+
+    def save_summary_plot(self, name='balance'):
+        plot_path = f"{self.process.working_path}/{self.process.plot_path}/" \
+                    f"{self.process.name}-{name}.pdf"
+        save_plot(plot_path)
+        plt.close('all')
+        self.process.callback(
+            'save_plot', self.process.id,
+            f"{self.process.name}: {name}", plot_path
+        )
+
+
+class LumiHistogrammer(BaseHistogrammer):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.colors = ("r", "g", "b")
+
+    def generate_frame_plot(self, idx, hist, chans, calc_hist_func):
+        _fig, axs = plt.subplots(1, 2, figsize=self.figsize, sharey=True)
+        self.histo_plot(axs[0], hist, "pixel luminosity", 'black')
+        for (chan, color) in zip(chans, self.colors):
+            hist_col = calc_hist_func(chan)
+            self.histo_plot(axs[1], hist_col, "R, G, B intensity", color, alpha=0.5)
+        plt.xlim(0, self.max_pixel_value)
+        self.save_plot(idx)
+
+    def generate_summary_plot(self, ref_idx):
+        plt.figure(figsize=self.figsize)
+        x = np.arange(0, len(self.corrections), dtype=int)
+        y = self.corrections
+        plt.plot([ref_idx, ref_idx], [0, np.max(y)], color='cornflowerblue',
+                 linestyle='--', label='reference frame')
+        plt.plot([x[0], x[-1]], [1, 1], color='lightgray', linestyle='--',
+                 label='no correction')
+        plt.plot(x, y, color='navy', label='luminosity correction')
+        plt.xlabel('frame')
+        plt.ylabel('correction')
+        plt.legend()
+        plt.xlim(x[0], x[-1])
+        plt.ylim(0, np.max(y) * 1.1)
+        self.save_summary_plot()
+
+
+class RGBHistogrammer(BaseHistogrammer):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.colors = ("r", "g", "b")
+
+    def generate_frame_plot(self, idx, hists):
+        _fig, axs = plt.subplots(1, 3, figsize=self.figsize, sharey=True)
+        for c in [2, 1, 0]:
+            self.histo_plot(axs[c], hists[c], self.colors[c] + " luminosity", self.colors[c])
+        plt.xlim(0, self.max_pixel_value)
+        self.save_plot(idx)
+
+    def generate_summary_plot(self, ref_idx):
+        plt.figure(figsize=self.figsize)
+        x = np.arange(0, len(self.corrections), dtype=int)
+        y = self.corrections
+        max_val = np.max(y) if np.any(y) else 1.0
+        plt.plot([ref_idx, ref_idx], [0, max_val], color='cornflowerblue',
+                 linestyle='--', label='reference frame')
+        plt.plot([x[0], x[-1]], [1, 1], color='lightgray', linestyle='--',
+                 label='no correction')
+        plt.plot(x, y[:, 0], color='r', label='R correction')
+        plt.plot(x, y[:, 1], color='g', label='G correction')
+        plt.plot(x, y[:, 2], color='b', label='B correction')
+        plt.xlabel('frame')
+        plt.ylabel('correction')
+        plt.legend()
+        plt.xlim(x[0], x[-1])
+        plt.ylim(0, max_val * 1.1)
+        self.save_summary_plot()
+
+
+class Ch1Histogrammer(BaseHistogrammer):
+    def __init__(self, labels, colors, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.labels = labels
+        self.colors = colors
+
+    def generate_frame_plot(self, idx, hists):
+        _fig, axs = plt.subplots(1, 3, figsize=self.figsize, sharey=True)
+        for c in range(3):
+            self.histo_plot(axs[c], hists[c], self.labels[c], self.colors[c])
+        plt.xlim(0, self.max_pixel_value)
+        self.save_plot(idx)
+
+    def generate_summary_plot(self, ref_idx):
+        plt.figure(figsize=self.figsize)
+        x = np.arange(0, len(self.corrections), dtype=int)
+        y = self.corrections
+        max_val = np.max(y) if np.any(y) else 1.0
+        plt.plot([ref_idx, ref_idx], [0, max_val], color='cornflowerblue',
+                 linestyle='--', label='reference frame')
+        plt.plot([x[0], x[-1]], [1, 1], color='lightgray', linestyle='--',
+                 label='no correction')
+        plt.plot(x, y[:, 0], color=self.colors[0], label=self.labels[0] + ' correction')
+        plt.xlabel('frame')
+        plt.ylabel('correction')
+        plt.legend()
+        plt.xlim(x[0], x[-1])
+        plt.ylim(0, max_val * 1.1)
+        self.save_summary_plot()
+
+
+class Ch2Histogrammer(BaseHistogrammer):
+    def __init__(self, labels, colors, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.labels = labels
+        self.colors = colors
+
+    def generate_frame_plot(self, idx, hists):
+        _fig, axs = plt.subplots(1, 3, figsize=self.figsize, sharey=True)
+        for c in range(3):
+            self.histo_plot(axs[c], hists[c], self.labels[c], self.colors[c])
+        plt.xlim(0, self.max_pixel_value)
+        self.save_plot(idx)
+
+    def generate_summary_plot(self, ref_idx):
+        plt.figure(figsize=self.figsize)
+        x = np.arange(0, len(self.corrections), dtype=int)
+        y = self.corrections
+        max_val = np.max(y) if np.any(y) else 1.0
+        plt.plot([ref_idx, ref_idx], [0, max_val], color='cornflowerblue',
+                 linestyle='--', label='reference frame')
+        plt.plot([x[0], x[-1]], [1, 1], color='lightgray', linestyle='--',
+                 label='no correction')
+        plt.plot(x, y[:, 0], color=self.colors[1], label=self.labels[1] + ' correction')
+        plt.plot(x, y[:, 1], color=self.colors[2], label=self.labels[2] + ' correction')
+        plt.xlabel('frame')
+        plt.ylabel('correction')
+        plt.legend()
+        plt.xlim(x[0], x[-1])
+        plt.ylim(0, max_val * 1.1)
+        self.save_summary_plot()
+
+
 class CorrectionMapBase:
     def __init__(self, dtype, ref_hist, intensity_interval=None):
         intensity_interval = {**constants.DEFAULT_INTENSITY_INTERVAL, **(intensity_interval or {})}
@@ -37,13 +216,11 @@ class CorrectionMapBase:
         chans = cv2.split(image)
         if self.channels == 2:
             ch_out = [chans[0]] + [self.apply_lut(
-                correction[c - 1],
-                self.reference[c - 1], chans[c]
+                correction[c - 1], self.reference[c - 1], chans[c]
             ) for c in range(1, 3)]
         elif self.channels == 3:
             ch_out = [self.apply_lut(
-                correction[c],
-                self.reference[c], chans[c]
+                correction[c], self.reference[c], chans[c]
             ) for c in range(3)]
         return cv2.merge(ch_out)
 
@@ -53,7 +230,7 @@ class CorrectionMapBase:
 
 class MatchHist(CorrectionMapBase):
     def __init__(self, dtype, ref_hist, intensity_interval=None):
-        CorrectionMapBase.__init__(self, dtype, ref_hist, intensity_interval)
+        super().__init__(dtype, ref_hist, intensity_interval)
         self.reference = self.cumsum(ref_hist)
         self.reference_mean = [r.mean() for r in self.reference]
         self.values = [*range(self.num_pixel_values)]
@@ -86,7 +263,7 @@ class MatchHist(CorrectionMapBase):
 
 class CorrectionMap(CorrectionMapBase):
     def __init__(self, dtype, ref_hist, intensity_interval=None):
-        CorrectionMapBase.__init__(self, dtype, ref_hist, intensity_interval)
+        super().__init__(dtype, ref_hist, intensity_interval)
         self.reference = [self.mid_val(self.id_lut, h) for h in ref_hist]
 
     def mid_val(self, lut, h):
@@ -94,9 +271,6 @@ class CorrectionMap(CorrectionMapBase):
 
 
 class GammaMap(CorrectionMap):
-    def __init__(self, dtype, ref_hist, intensity_interval=None):
-        CorrectionMap.__init__(self, dtype, ref_hist, intensity_interval)
-
     def correction(self, hist):
         return [bisect(lambda x: self.mid_val(self.lut(x), h) - r, 0.1, 5)
                 for h, r in zip(hist, self.reference)]
@@ -109,9 +283,6 @@ class GammaMap(CorrectionMap):
 
 
 class LinearMap(CorrectionMap):
-    def __init__(self, dtype, ref_hist, intensity_interval=None):
-        CorrectionMap.__init__(self, dtype, ref_hist, intensity_interval)
-
     def lut(self, correction, _reference=None):
         ar = np.arange(0, self.num_pixel_values)
         return np.clip(ar * correction, 0, self.max_pixel_value).astype(self.dtype)
@@ -122,42 +293,56 @@ class LinearMap(CorrectionMap):
 
 class Correction:
     def __init__(self, channels, mask_size=0, intensity_interval=None,
-                 subsample=-1, fast_subsampling=constants.DEFAULT_BALANCE_FAST_SUBSAMPLING,
+                 subsample=constants.DEFAULT_BALANCE_SUBSAMPLE,
+                 fast_subsampling=constants.DEFAULT_BALANCE_FAST_SUBSAMPLING,
                  corr_map=constants.DEFAULT_CORR_MAP,
                  plot_histograms=False, plot_summary=False):
         self.mask_size = mask_size
         self.intensity_interval = intensity_interval
-        self.plot_histograms = plot_histograms
-        self.plot_summary = plot_summary
-        self.subsample = constants.DEFAULT_BALANCE_SUBSAMPLE if subsample == -1 else subsample
+        self.subsample = subsample
         self.fast_subsampling = fast_subsampling
         self.corr_map = corr_map
         self.channels = channels
+        self.plot_histograms = plot_histograms
+        self.plot_summary = plot_summary
         self.dtype = None
         self.num_pixel_values = None
-        self. max_pixel_value = None
-        self.corrections = None
+        self.max_pixel_value = None
+        self.corr_map_obj = None
         self.process = None
+        self.histogrammer = None
 
     def begin(self, ref_image, size, ref_idx):
         self.dtype = ref_image.dtype
-        self.num_pixel_values = constants.NUM_UINT8 if ref_image.dtype == np.uint8 \
-            else constants.NUM_UINT16
+        self.num_pixel_values = constants.NUM_UINT8 \
+            if ref_image.dtype == np.uint8 else constants.NUM_UINT16
         self.max_pixel_value = self.num_pixel_values - 1
+        self._create_histogrammer()
+        self.histogrammer.process = self.process
         hist = self.get_hist(self.preprocess(ref_image), ref_idx)
         if self.corr_map == constants.BALANCE_LINEAR:
-            self.corr_map = LinearMap(self.dtype, hist, self.intensity_interval)
+            self.corr_map_obj = LinearMap(self.dtype, hist, self.intensity_interval)
         elif self.corr_map == constants.BALANCE_GAMMA:
-            self.corr_map = GammaMap(self.dtype, hist, self.intensity_interval)
+            self.corr_map_obj = GammaMap(self.dtype, hist, self.intensity_interval)
         elif self.corr_map == constants.BALANCE_MATCH_HIST:
-            self.corr_map = MatchHist(self.dtype, hist, self.intensity_interval)
+            self.corr_map_obj = MatchHist(self.dtype, hist, self.intensity_interval)
         else:
             raise InvalidOptionError("corr_map", self.corr_map)
-        self.corrections = np.ones((size, self.channels))
+        self.histogrammer.begin(size)
+
+    def _create_histogrammer(self):
+        raise NotImplementedError("Subclasses must implement _create_histogrammer")
 
     def calc_hist_1ch(self, image):
+        if self.subsample > 0:
+            subsample = self.subsample
+        else:
+            h, w = image.shape[:2]
+            img_res = float(h) * float(w) / constants.ONE_MEGA
+            target_res = constants.DEFAULT_BALANCE_RES_TARGET_MPX
+            subsample = int(1 + math.floor(img_res / target_res))
         img_sub = image if self.subsample == 1 \
-            else img_subsample(image, self.subsample, self.fast_subsampling)
+            else img_subsample(image, subsample, self.fast_subsampling)
         if self.mask_size == 0:
             image_sel = img_sub
         else:
@@ -170,16 +355,13 @@ class Correction:
             image_sel = img_sub[
                 (xv - width / 2) ** 2 + (yv - height / 2) ** 2 <= mask_radius ** 2
             ]
-        hist, _bins = np.histogram(
-            image_sel,
-            bins=np.linspace(-0.5, self.num_pixel_values - 0.5,
-                             self.num_pixel_values + 1)
-        )
+        bins = np.linspace(-0.5, self.num_pixel_values - 0.5, self.num_pixel_values + 1)
+        hist, _bins = np.histogram(image_sel, bins=bins)
         return hist
 
     def balance(self, image, idx):
-        correction = self.corr_map.correction(self.get_hist(image, idx))
-        return correction, self.corr_map.adjust(image, correction)
+        correction = self.corr_map_obj.correction(self.get_hist(image, idx))
+        return correction, self.corr_map_obj.adjust(image, correction)
 
     def get_hist(self, _image, _idx):
         return None
@@ -188,10 +370,12 @@ class Correction:
         pass
 
     def apply_correction(self, idx, image):
+        if idx == self.process.ref_idx:
+            return image
         image = self.preprocess(image)
         correction, image = self.balance(image, idx)
         image = self.postprocess(image)
-        self.corrections[idx] = self.corr_map.correction_size(correction)
+        self.histogrammer.add_correction(idx, self.corr_map_obj.correction_size(correction))
         return image
 
     def preprocess(self, image):
@@ -200,167 +384,171 @@ class Correction:
     def postprocess(self, image):
         return image
 
-    def histo_plot(self, ax, hist, x_label, color, alpha=1):
-        ax.set_ylabel("# of pixels")
-        ax.set_xlabel(x_label)
-        ax.set_xlim([0, self.num_pixel_values])
-        ax.set_yscale('log')
-        ax.plot(hist, color=color, alpha=alpha)
-
-    def save_plot(self, idx):
-        idx_str = f"{idx:04d}"
-        plot_path = f"{self.process.working_path}/" \
-            f"{self.process.plot_path}/{self.process.name}-hist-{idx_str}.pdf"
-        save_plot(plot_path)
-        plt.close('all')
-        self.process.callback(
-            'save_plot',
-            self.process.id, f"{self.process.name}: balance\nframe {idx_str}",
-            plot_path
-        )
-
-    def save_summary_plot(self, name='balance'):
-        plot_path = f"{self.process.working_path}/" \
-            f"{self.process.plot_path}/{self.process.name}-{name}.pdf"
-        save_plot(plot_path)
-        plt.close('all')
-        self.process.callback(
-            'save_plot', self.process.id,
-            f"{self.process.name}: {name}", plot_path
-        )
-
 
 class LumiCorrection(Correction):
     def __init__(self, **kwargs):
-        Correction.__init__(self, 1, **kwargs)
+        super().__init__(1, **kwargs)
+
+    def _create_histogrammer(self):
+        self.histogrammer = LumiHistogrammer(
+            dtype=self.dtype,
+            num_pixel_values=self.num_pixel_values,
+            max_pixel_value=self.max_pixel_value,
+            channels=1,
+            plot_histograms=self.plot_histograms,
+            plot_summary=self.plot_summary
+        )
 
     def get_hist(self, image, idx):
         hist = self.calc_hist_1ch(cv2.cvtColor(image, cv2.COLOR_BGR2GRAY))
-        chans = cv2.split(image)
-        colors = ("r", "g", "b")
-        if self.plot_histograms:
-            _fig, axs = plt.subplots(1, 2, figsize=(10, 5), sharey=True)
-            self.histo_plot(axs[0], hist, "pixel luminosity", 'black')
-            for (chan, color) in zip(chans, colors):
-                hist_col = self.calc_hist_1ch(chan)
-                self.histo_plot(axs[1], hist_col, "r,g,b luminosity", color, alpha=0.5)
-            plt.xlim(0, self.max_pixel_value)
-            self.save_plot(idx)
+        if self.histogrammer.plot_histograms:
+            chans = cv2.split(image)
+            self.histogrammer.generate_frame_plot(idx, hist, chans, self.calc_hist_1ch)
         return [hist]
 
     def end(self, ref_idx):
-        if self.plot_summary:
-            plt.figure(figsize=(10, 5))
-            x = np.arange(1, len(self.corrections) + 1, dtype=int)
-            y = self.corrections
-            plt.plot([ref_idx + 1, ref_idx + 1], [0, 1], color='cornflowerblue',
-                     linestyle='--', label='reference frame')
-            plt.plot([x[0], x[-1]], [1, 1], color='lightgray', linestyle='--',
-                     label='no correction')
-            plt.plot(x, y, color='navy', label='luminosity correction')
-            plt.xlabel('frame')
-            plt.ylabel('correction')
-            plt.legend()
-            plt.xlim(x[0], x[-1])
-            plt.ylim(0)
-            self.save_summary_plot()
+        if self.histogrammer.plot_summary:
+            self.histogrammer.generate_summary_plot(ref_idx)
 
 
 class RGBCorrection(Correction):
     def __init__(self, **kwargs):
-        Correction.__init__(self, 3, **kwargs)
+        super().__init__(3, **kwargs)
+
+    def _create_histogrammer(self):
+        self.histogrammer = RGBHistogrammer(
+            dtype=self.dtype,
+            num_pixel_values=self.num_pixel_values,
+            max_pixel_value=self.max_pixel_value,
+            channels=3,
+            plot_histograms=self.plot_histograms,
+            plot_summary=self.plot_summary
+        )
 
     def get_hist(self, image, idx):
         hist = [self.calc_hist_1ch(chan) for chan in cv2.split(image)]
-        colors = ("r", "g", "b")
-        if self.plot_histograms:
-            _fig, axs = plt.subplots(1, 3, figsize=(10, 5), sharey=True)
-            for c in [2, 1, 0]:
-                self.histo_plot(axs[c], hist[c], colors[c] + " luminosity", colors[c])
-            plt.xlim(0, self.max_pixel_value)
-            self.save_plot(idx)
+        if self.histogrammer.plot_histograms:
+            self.histogrammer.generate_frame_plot(idx, hist)
         return hist
 
     def end(self, ref_idx):
-        if self.plot_summary:
-            plt.figure(figsize=(10, 5))
-            x = np.arange(1, len(self.corrections) + 1, dtype=int)
-            y = self.corrections
-            plt.plot([ref_idx + 1, ref_idx + 1], [0, 1], color='cornflowerblue',
-                     linestyle='--', label='reference frame')
-            plt.plot([x[0], x[-1]], [1, 1], color='lightgray', linestyle='--',
-                     label='no correction')
-            plt.plot(x, y[:, 0], color='r', label='R correction')
-            plt.plot(x, y[:, 1], color='g', label='G correction')
-            plt.plot(x, y[:, 2], color='b', label='B correction')
-            plt.xlabel('frame')
-            plt.ylabel('correction')
-            plt.legend()
-            plt.xlim(x[0], x[-1])
-            plt.ylim(0)
-            self.save_summary_plot()
+        if self.histogrammer.plot_summary:
+            self.histogrammer.generate_summary_plot(ref_idx)
+
+
+class Ch1Correction(Correction):
+    def __init__(self, **kwargs):
+        super().__init__(1, **kwargs)
+        self.labels = None
+        self.colors = None
+
+    def preprocess(self, image):
+        raise NotImplementedError('abstract method')
+
+    def get_hist(self, image, idx):
+        hist = [self.calc_hist_1ch(chan) for chan in cv2.split(image)]
+        if self.histogrammer.plot_histograms:
+            self.histogrammer.generate_frame_plot(idx, hist)
+        return [hist[0]]
+
+    def end(self, ref_idx):
+        if self.histogrammer.plot_summary:
+            self.histogrammer.generate_summary_plot(ref_idx)
 
 
 class Ch2Correction(Correction):
     def __init__(self, **kwargs):
-        Correction.__init__(self, 2, **kwargs)
+        super().__init__(2, **kwargs)
+        self.labels = None
+        self.colors = None
 
     def preprocess(self, image):
-        assert False, 'abstract method'
+        raise NotImplementedError('abstract method')
 
     def get_hist(self, image, idx):
         hist = [self.calc_hist_1ch(chan) for chan in cv2.split(image)]
-        if self.plot_histograms:
-            _fig, axs = plt.subplots(1, 3, figsize=(10, 5), sharey=True)
-            for c in range(3):
-                self.histo_plot(axs[c], hist[c], self.labels[c], self.colors[c])
-            plt.xlim(0, self.max_pixel_value)
-            self.save_plot(idx)
+        if self.histogrammer.plot_histograms:
+            self.histogrammer.generate_frame_plot(idx, hist)
         return hist[1:]
 
     def end(self, ref_idx):
-        if self.plot_summary:
-            plt.figure(figsize=(10, 5))
-            x = np.arange(1, len(self.corrections) + 1, dtype=int)
-            y = self.corrections
-            plt.plot([ref_idx + 1, ref_idx + 1], [0, 1], color='cornflowerblue',
-                     linestyle='--', label='reference frame')
-            plt.plot([x[0], x[-1]], [1, 1], color='lightgray', linestyle='--',
-                     label='no correction')
-            plt.plot(x, y[:, 0], color=self.colors[1], label=self.labels[1] + ' correction')
-            plt.plot(x, y[:, 1], color=self.colors[2], label=self.labels[2] + ' correction')
-            plt.xlabel('frame')
-            plt.ylabel('correction')
-            plt.legend()
-            plt.xlim(x[0], x[-1])
-            plt.ylim(0)
-            self.save_summary_plot()
+        if self.histogrammer.plot_summary:
+            self.histogrammer.generate_summary_plot(ref_idx)
 
 
 class SVCorrection(Ch2Correction):
     def __init__(self, **kwargs):
-        Ch2Correction.__init__(self, **kwargs)
+        super().__init__(**kwargs)
         self.labels = ("H", "S", "V")
         self.colors = ("hotpink", "orange", "navy")
 
+    def _create_histogrammer(self):
+        self.histogrammer = Ch2Histogrammer(
+            dtype=self.dtype,
+            num_pixel_values=self.num_pixel_values,
+            max_pixel_value=self.max_pixel_value,
+            channels=2,
+            plot_histograms=self.plot_histograms,
+            plot_summary=self.plot_summary,
+            labels=self.labels,
+            colors=self.colors
+        )
+
     def preprocess(self, image):
-        return cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
+        return bgr_to_hsv(image)
 
     def postprocess(self, image):
-        return cv2.cvtColor(image, cv2.COLOR_HSV2BGR)
+        return hsv_to_bgr(image)
 
 
 class LSCorrection(Ch2Correction):
     def __init__(self, **kwargs):
-        Ch2Correction.__init__(self, **kwargs)
+        super().__init__(**kwargs)
         self.labels = ("H", "L", "S")
         self.colors = ("hotpink", "navy", "orange")
 
+    def _create_histogrammer(self):
+        self.histogrammer = Ch2Histogrammer(
+            dtype=self.dtype,
+            num_pixel_values=self.num_pixel_values,
+            max_pixel_value=self.max_pixel_value,
+            channels=2,
+            plot_histograms=self.plot_histograms,
+            plot_summary=self.plot_summary,
+            labels=self.labels,
+            colors=self.colors
+        )
+
+    def preprocess(self, image):
+        return bgr_to_hls(image)
+
+    def postprocess(self, image):
+        return hls_to_bgr(image)
+
+
+class LABCorrection(Ch1Correction):
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.labels = ("L", "A", "B")
+        self.colors = ("black", "yellow", "red")
+
+    def _create_histogrammer(self):
+        self.histogrammer = Ch1Histogrammer(
+            dtype=self.dtype,
+            num_pixel_values=self.num_pixel_values,
+            max_pixel_value=self.max_pixel_value,
+            channels=2,
+            plot_histograms=self.plot_histograms,
+            plot_summary=self.plot_summary,
+            labels=self.labels,
+            colors=self.colors
+        )
+
     def preprocess(self, image):
-        return cv2.cvtColor(image, cv2.COLOR_BGR2HLS)
+        return bgr_to_lab(image)
 
     def postprocess(self, image):
-        return cv2.cvtColor(image, cv2.COLOR_HLS2BGR)
+        return lab_to_bgr(image)
 
 
 class BalanceFrames(SubAction):
@@ -368,26 +556,37 @@ class BalanceFrames(SubAction):
         super().__init__(enabled=enabled)
         self.process = None
         self.shape = None
-        corr_map = kwargs.get('corr_map', constants.DEFAULT_CORR_MAP)
-        subsample = kwargs.get('subsample', constants.DEFAULT_BALANCE_SUBSAMPLE)
+        self.corr_map = kwargs.get(
+            'corr_map', constants.DEFAULT_CORR_MAP)
+        self.subsample = kwargs.get(
+            'subsample', constants.DEFAULT_BALANCE_SUBSAMPLE)
         self.fast_subsampling = kwargs.get(
             'fast_subsampling', constants.DEFAULT_BALANCE_FAST_SUBSAMPLING)
-        channel = kwargs.pop('channel', constants.DEFAULT_CHANNEL)
-        kwargs['subsample'] = (
-            1 if corr_map == constants.BALANCE_MATCH_HIST
-            else constants.DEFAULT_BALANCE_SUBSAMPLE) if subsample == -1 else subsample
+        self.channel = kwargs.get(
+            'channel', constants.DEFAULT_CHANNEL)
         self.mask_size = kwargs.get('mask_size', 0)
         self.plot_summary = kwargs.get('plot_summary', False)
-        if channel == constants.BALANCE_LUMI:
-            self.correction = LumiCorrection(**kwargs)
-        elif channel == constants.BALANCE_RGB:
-            self.correction = RGBCorrection(**kwargs)
-        elif channel == constants.BALANCE_HSV:
-            self.correction = SVCorrection(**kwargs)
-        elif channel == constants.BALANCE_HLS:
-            self.correction = LSCorrection(**kwargs)
-        else:
-            raise InvalidOptionError("channel", channel)
+        self.plot_histograms = kwargs.get('plot_histograms', False)
+        if self.subsample == -1:
+            self.subsample = (1 if self.corr_map == constants.BALANCE_MATCH_HIST
+                              else constants.DEFAULT_BALANCE_SUBSAMPLE)
+        correction_class = {
+            constants.BALANCE_LUMI: LumiCorrection,
+            constants.BALANCE_RGB: RGBCorrection,
+            constants.BALANCE_HSV: SVCorrection,
+            constants.BALANCE_HLS: LSCorrection,
+            constants.BALANCE_LAB: LABCorrection,
+        }.get(self.channel, None)
+        if correction_class is None:
+            raise InvalidOptionError("channel", self.channel)
+        self.correction = correction_class(
+            mask_size=self.mask_size,
+            subsample=self.subsample,
+            fast_subsampling=self.fast_subsampling,
+            corr_map=self.corr_map,
+            plot_histograms=self.plot_histograms,
+            plot_summary=self.plot_summary
+        )
 
     def begin(self, process):
         self.process = process
@@ -407,10 +606,10 @@ class BalanceFrames(SubAction):
             plt.figure(figsize=(10, 5))
             plt.title('Mask')
             plt.imshow(img, 'gray')
-            self.correction.save_summary_plot("mask")
+            self.correction.histogrammer.save_summary_plot("mask")
 
     def run_frame(self, idx, _ref_idx, image):
         if idx != self.process.ref_idx:
             self.process.sub_message_r(color_str(': balance image', constants.LOG_COLOR_LEVEL_3))
-            image = self.correction.apply_correction(idx, image)
+        image = self.correction.apply_correction(idx, image)
         return image