Source code for pyctools.components.photo.colourcorrect

#  Pyctools - a picture processing algorithm development kit.
#  http://github.com/jim-easterbrook/pyctools
#  Copyright (C) 2016  Pyctools contributors
#
#  This program is free software: you can redistribute it and/or
#  modify it under the terms of the GNU General Public License as
#  published by the Free Software Foundation, either version 3 of the
#  License, or (at your option) any later version.
#
#  This program is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
#  General Public License for more details.
#
#  You should have received a copy of the GNU General Public License
#  along with this program.  If not, see
#  <http://www.gnu.org/licenses/>.

from __future__ import print_function

__all__ = ['ColourCorrect']
__docformat__ = 'restructuredtext en'

import numpy

from pyctools.core.config import ConfigFloat
from pyctools.core.base import Transformer
from pyctools.core.types import pt_float

[docs] class ColourCorrect(Transformer): """Colour correction. Adjust hue and saturation of R, G & B channels separately. =========== ===== ==== Config =========== ===== ==== ``gain`` float Adjust overall gain. ``R_hue`` float Adjust hue of red primary. ``R_sat`` float Adjust saturation of red primary. ``G_hue`` float Adjust hue of green primary. ``G_sat`` float Adjust saturation of green primary. ``B_hue`` float Adjust hue of blue primary. ``B_sat`` float Adjust saturation of blue primary. =========== ===== ==== """ def initialise(self): self.config['gain'] = ConfigFloat(value=1.0, decimals=2) self.config['R_hue'] = ConfigFloat(decimals=2) self.config['R_sat'] = ConfigFloat(value=1.0, decimals=2) self.config['G_hue'] = ConfigFloat(decimals=2) self.config['G_sat'] = ConfigFloat(value=1.0, decimals=2) self.config['B_hue'] = ConfigFloat(decimals=2) self.config['B_sat'] = ConfigFloat(value=1.0, decimals=2) def transform(self, in_frame, out_frame): # compute colour matrix self.update_config() gain = self.config['gain'] R_hue = self.config['R_hue'] R_sat = self.config['R_sat'] G_hue = self.config['G_hue'] G_sat = self.config['G_sat'] B_hue = self.config['B_hue'] B_sat = self.config['B_sat'] # 'hue' matrix hue_matrix = numpy.array( [[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]], dtype=pt_float) hue_matrix += numpy.array( [[ 0.0, 0.0, 0.0], [ 0.5, 0.0, 0.0], [-0.5, 0.0, 0.0]], dtype=pt_float) * pt_float(R_hue * R_sat) hue_matrix += numpy.array( [[0.0, 0.5, 0.0], [0.0, 0.0, 0.0], [0.0, -0.5, 0.0]], dtype=pt_float) * pt_float(G_hue * G_sat) hue_matrix += numpy.array( [[0.0, 0.0, 0.5], [0.0, 0.0, -0.5], [0.0, 0.0, 0.0]], dtype=pt_float) * pt_float(B_hue * B_sat) # adjust to preserve white balance hue_matrix /= numpy.array( [[sum(hue_matrix[0])], [sum(hue_matrix[1])], [sum(hue_matrix[2])]], dtype=pt_float) # 'saturation' matrix - uses BT.709 RGB->Y as specified by sRGB sat_matrix = numpy.array( [[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]], dtype=pt_float) sat_matrix += numpy.array( [[0.2126 - 1.0, 0.0, 0.0], [0.2126, 0.0, 0.0], [0.2126, 0.0, 0.0]], dtype=pt_float) * pt_float(1.0 - R_sat) sat_matrix += numpy.array( [[0.0, 0.7152, 0.0], [0.0, 0.7152 - 1.0, 0.0], [0.0, 0.7152, 0.0]], dtype=pt_float) * pt_float(1.0 - G_sat) sat_matrix += numpy.array( [[0.0, 0.0, 0.0722], [0.0, 0.0, 0.0722], [0.0, 0.0, 0.0722 - 1.0]], dtype=pt_float) * pt_float(1.0 - B_sat) matrix = numpy.dot(hue_matrix, sat_matrix) * pt_float(gain) # apply matrix in_data = in_frame.as_numpy(dtype=pt_float) out_frame.data = numpy.dot(in_data, matrix.T) # add audit audit = out_frame.metadata.get('audit') audit += 'data = ColourCorrect(data)\n' if gain != 1.0: audit += ' gain: {}\n'.format(gain) if R_hue != 0.0: audit += ' R_hue: {}\n'.format(R_hue) if R_sat != 1.0: audit += ' R_sat: {}\n'.format(R_sat) if G_hue != 0.0: audit += ' G_hue: {}\n'.format(G_hue) if G_sat != 1.0: audit += ' G_sat: {}\n'.format(G_sat) if B_hue != 0.0: audit += ' B_hue: {}\n'.format(B_hue) if B_sat != 1.0: audit += ' B_sat: {}\n'.format(B_sat) out_frame.metadata.set('audit', audit) return True