apps/CameraITS/tools/dng_noise_model.py - platform/cts - Git at Google

 # Copyright 2014 The Android Open Source Project
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
 # You may obtain a copy of the License at
 #
 #      http://www.apache.org/licenses/LICENSE-2.0
 #
 # Unless required by applicable law or agreed to in writing, software
 # distributed under the License is distributed on an "AS IS" BASIS,
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """CameraITS script to generate noise models."""

 import logging
 import math
 import os.path
 import pathlib
 import pickle
 import tempfile
 import textwrap

 import capture_read_noise_utils
 import its_base_test
 import its_session_utils
 from matplotlib import pylab
 import matplotlib.pyplot as plt
 import matplotlib.ticker
 from mobly import test_runner
 import noise_model_utils
 import numpy as np

 _IS_QUAD_BAYER = False  # A manual flag to choose standard or quad Bayer noise
                         # model generation.
 if _IS_QUAD_BAYER:
   _COLOR_CHANNEL_NAMES = noise_model_utils.QUAD_BAYER_COLORS
   _PLOT_COLORS = noise_model_utils.QUAD_BAYER_PLOT_COLORS
   _TILE_SIZE = 64  # Tile size to compute mean/variance. Large tiles may have
                    # their variance corrupted by low freq image changes.
   _STATS_FORMAT = 'raw10QuadBayerStats'  # rawQuadBayerStats|raw10QuadBayerStats
   _READ_NOISE_RAW_FORMAT = 'raw10QuadBayer'  # rawQuadBayer|raw10QuadBayer
 else:
   _COLOR_CHANNEL_NAMES = noise_model_utils.BAYER_COLORS
   _PLOT_COLORS = noise_model_utils.BAYER_PLOT_COLORS
   _TILE_SIZE = 32  # Tile size to compute mean/variance. Large tiles may have
                    # their variance corrupted by low freq image changes.
   _STATS_FORMAT = 'rawStats'  # rawStats|raw10Stats
   _READ_NOISE_RAW_FORMAT = 'raw'  # raw|raw10

 _STATS_CONFIG = {
     'format': _STATS_FORMAT,
     'gridWidth': _TILE_SIZE,
     'gridHeight': _TILE_SIZE,
 }
 _BRACKET_MAX = 8  # Exposure bracketing range in stops
 _BRACKET_FACTOR = math.pow(2, _BRACKET_MAX)
 _ISO_MAX_VALUE = None  # ISO range max value, uses sensor max if None
 _ISO_MIN_VALUE = None  # ISO range min value, uses sensor min if None
 _MAX_SCALE_FUDGE = 1.1
 _MAX_SIGNAL_VALUE = 0.25  # Maximum value to allow mean of the tiles to go.
 _NAME = os.path.basename(__file__).split('.')[0]
 _NAME_READ_NOISE = os.path.join(tempfile.gettempdir(), 'CameraITS/ReadNoise')
 _NAME_READ_NOISE_FILE = 'read_noise_results.pkl'
 _STATS_FILE_NAME = 'stats.pkl'
 _OUTLIER_MEDIAN_ABS_DEVS = 10  # Defines the number of Median Absolute
                                # Deviations that constitutes acceptable data
 _READ_NOISE_STEPS_PER_STOP = 12  # Sensitivities per stop to sample for read
                                  # noise
 _REMOVE_VAR_OUTLIERS = False  # When True, filters the variance to remove
                               # outliers
 _STEPS_PER_STOP = 3  # How many sensitivities per stop to sample.
 _ISO_MULTIPLIER = math.pow(2, 1.0 / _STEPS_PER_STOP)
 _TILE_CROP_N = 0  # Number of tiles to crop from edge of image. Usually 0.
 _TWO_STAGE_MODEL = False  # Require read noise data prior to running noise model
 _ZOOM_RATIO = 1  # Zoom target to be used while running the model
 _FIG_DPI = 100  # DPI for plotting noise model figures.
 _BAYER_COLORS_FOR_NOISE_PROFILE = [color.lower() for color in
                                    noise_model_utils.BAYER_COLORS]
 _QUAD_BAYER_COLORS_FOR_NOISE_PROFILE = [color.lower() for color in
                                         noise_model_utils.QUAD_BAYER_COLORS]


 class DngNoiseModel(its_base_test.ItsBaseTest):
   """Create DNG noise model.

   Captures RAW images with increasing analog gains to create the model.
   """

   def _create_noise_model_code(self, noise_model, sens_min, sens_max,
                                sens_max_analog, file_path):
     """Creates the C file for the noise model.

     Args:
       noise_model: Noise model parameters.
       sens_min: The minimum sensitivity value.
       sens_max: The maximum sensitivity value.
       sens_max_analog: The maximum analog sensitivity value.
       file_path: The path to the noise model file.
     """
     # Generate individual noise model components.
     scale_a, scale_b, offset_a, offset_b = zip(*noise_model)
     digital_gain_cdef = (
         f'(sens / {sens_max_analog:.1f}) < 1.0 ? '
         f'1.0 : (sens / {sens_max_analog:.1f})'
     )

     with open(file_path, 'w') as text_file:
       scale_a_str = ','.join([str(i) for i in scale_a])
       scale_b_str = ','.join([str(i) for i in scale_b])
       offset_a_str = ','.join([str(i) for i in offset_a])
       offset_b_str = ','.join([str(i) for i in offset_b])
       # pylint: disable=line-too-long
       code = textwrap.dedent(f"""\
               /* Generated test code to dump a table of data for external validation
               * of the noise model parameters.
               */
               #include <stdio.h>
               #include <assert.h>
               double compute_noise_model_entry_S(int plane, int sens);
               double compute_noise_model_entry_O(int plane, int sens);
               int main(void) {{
                   for (int plane = 0; plane < {len(scale_a)}; plane++) {{
                       for (int sens = {sens_min}; sens <= {sens_max}; sens += 100) {{
                           double o = compute_noise_model_entry_O(plane, sens);
                           double s = compute_noise_model_entry_S(plane, sens);
                           printf("%d,%d,%lf,%lf\\n", plane, sens, o, s);
                       }}
                   }}
                   return 0;
               }}

               /* Generated functions to map a given sensitivity to the O and S noise
               * model parameters in the DNG noise model. The planes are in
               * R, Gr, Gb, B order.
               */
               double compute_noise_model_entry_S(int plane, int sens) {{
                   static double noise_model_A[] = {{ {scale_a_str:s} }};
                   static double noise_model_B[] = {{ {scale_b_str:s} }};
                   double A = noise_model_A[plane];
                   double B = noise_model_B[plane];
                   double s = A * sens + B;
                   return s < 0.0 ? 0.0 : s;
               }}

               double compute_noise_model_entry_O(int plane, int sens) {{
                   static double noise_model_C[] = {{ {offset_a_str:s} }};
                   static double noise_model_D[] = {{ {offset_b_str:s} }};
                   double digital_gain = {digital_gain_cdef:s};
                   double C = noise_model_C[plane];
                   double D = noise_model_D[plane];
                   double o = C * sens * sens + D * digital_gain * digital_gain;
                   return o < 0.0 ? 0.0 : o;
               }}
               """)

       text_file.write(code)

   def _create_noise_profile_code(self, noise_model, color_channels, file_path):
     """Creates the noise profile C++ file.

     Args:
       noise_model: Noise model parameters.
       color_channels: Color channels in canonical order.
       file_path: The path to the noise profile C++ file.
     """
     # Generate individual noise model components.
     scale_a, scale_b, offset_a, offset_b = zip(*noise_model)
     num_channels = noise_model.shape[0]
     params = []
     for ch, color in enumerate(color_channels):
       prefix = f'.noise_coefficients_{color} = {{'
       spaces = ' ' * len(prefix)
       suffix = '},' if ch != num_channels - 1 else '}'
       params.append(textwrap.dedent(f"""
         {prefix}.gradient_slope = {scale_a[ch]},
         {spaces}.offset_slope = {scale_b[ch]},
         {spaces}.gradient_intercept = {offset_a[ch]},
         {spaces}.offset_intercept = {offset_b[ch]}{suffix}"""))

     with open(file_path, 'w') as text_file:
       # pylint: disable=line-too-long
       code_comment = textwrap.dedent("""\
               /* noise_profile.cc
                 Note: gradient_slope --> gradient of API s_measured parameter
                       offset_slope --> o_model of API s_measured parameter
                       gradient_intercept--> gradient of API o_measured parameter
                       offset_intercept --> o_model of API o_measured parameter
                 Note: SENSOR_NOISE_PROFILE in Android Developers doc uses
                       N(x) = sqrt(Sx + O), where 'S' is 's_measured' & 'O' is 'o_measured'
               */
       """)
       params_str = textwrap.indent(''.join(params), ' ' * 4)
       code_params = '.noise_profile = {' + params_str + '},'
       code = code_comment + code_params
       text_file.write(code)

   def _create_noise_model_and_profile_code(self, noise_model, sens_min,
                                            sens_max, sens_max_analog, log_path):
     """Creates the code file with noise model parameters.

     Args:
       noise_model: Noise model parameters.
       sens_min: The minimum sensitivity value.
       sens_max: The maximum sensitivity value.
       sens_max_analog: The maximum analog sensitivity value.
       log_path: The path to the log file.
     """
     noise_model_utils.check_noise_model_shape(noise_model)
     # Create noise model code with noise model parameters.
     self._create_noise_model_code(
         noise_model,
         sens_min,
         sens_max,
         sens_max_analog,
         os.path.join(log_path, 'noise_model.c'),
     )

     num_channels = noise_model.shape[0]
     is_quad_bayer = num_channels == noise_model_utils.NUM_QUAD_BAYER_CHANNELS
     if is_quad_bayer:
       # Average noise model parameters of every four channels.
       avg_noise_model = noise_model.reshape(-1, 4, noise_model.shape[1]).mean(
           axis=1
       )
       # Create noise model code with average noise model parameters.
       self._create_noise_model_code(
           avg_noise_model,
           sens_min,
           sens_max,
           sens_max_analog,
           os.path.join(log_path, 'noise_model_avg.c'),
       )
       # Create noise profile code with average noise model parameters.
       self._create_noise_profile_code(
           avg_noise_model,
           _BAYER_COLORS_FOR_NOISE_PROFILE,
           os.path.join(log_path, 'noise_profile_avg.cc'),
       )
       # Create noise profile code with noise model parameters.
       self._create_noise_profile_code(
           noise_model,
           _QUAD_BAYER_COLORS_FOR_NOISE_PROFILE,
           os.path.join(log_path, 'noise_profile.cc'),
       )

     else:
       # Create noise profile code with noise model parameters.
       self._create_noise_profile_code(
           noise_model,
           _BAYER_COLORS_FOR_NOISE_PROFILE,
           os.path.join(log_path, 'noise_profile.cc'),
       )

   def _plot_stats_and_noise_model_fittings(
       self, iso_to_stats_dict, measured_models, noise_model, sens_max_analog,
       folder_path_prefix):
     """Plots the stats (means, vars_) and noise models fittings.

     Args:
       iso_to_stats_dict: A dictionary mapping ISO to a list of tuples of
         exposure time in milliseconds, mean values, and variance values.
       measured_models: A list of measured noise models for each ISO value.
       noise_model: A numpy array of global noise model parameters for all ISO
         values.
       sens_max_analog: The maximum analog sensitivity value.
       folder_path_prefix: The prefix of path to save figures.

     Raises:
       ValueError: If the noise model shape is invalid.
     """
     noise_model_utils.check_noise_model_shape(noise_model)
     # Separate individual noise model components.
     scale_a, scale_b, offset_a, offset_b = zip(*noise_model)

     iso_pidx_to_measured_model_dict = {}
     num_channels = noise_model.shape[0]
     for pidx in range(num_channels):
       for iso, s_measured, o_measured in measured_models[pidx]:
         iso_pidx_to_measured_model_dict[(iso, pidx)] = (s_measured, o_measured)

     isos = np.asarray(sorted(iso_to_stats_dict.keys()))
     digital_gains = noise_model_utils.compute_digital_gains(
         isos, sens_max_analog
     )

     x_range = [0, _MAX_SIGNAL_VALUE]
     for iso, digital_gain in zip(isos, digital_gains):
       logging.info('Plotting stats and noise model for ISO %d.', iso)
       fig, subplots = noise_model_utils.create_stats_figure(
           iso, _COLOR_CHANNEL_NAMES
       )

       xmax = 0
       stats_per_plane = [[] for _ in range(num_channels)]
       for exposure_ms, means, vars_ in iso_to_stats_dict[iso]:
         exposure_norm = noise_model_utils.COLOR_NORM(np.log2(exposure_ms))
         exposure_color = noise_model_utils.RAINBOW_CMAP(exposure_norm)
         for pidx in range(num_channels):
           means_p = means[pidx]
           vars_p = vars_[pidx]

           if means_p.size > 0 and vars_p.size > 0:
             subplots[pidx].plot(
                 means_p,
                 vars_p,
                 color=exposure_color,
                 marker='.',
                 markeredgecolor=exposure_color,
                 markersize=1,
                 linestyle='None',
                 alpha=0.5,
             )

             stats_per_plane[pidx].extend(list(zip(means_p, vars_p)))
             xmax = max(xmax, max(means_p))

       iso_sq = iso ** 2
       digital_gain_sq = digital_gain ** 2
       for pidx in range(num_channels):
         # Add the final noise model to subplots.
         s_model = scale_a[pidx] * iso * digital_gain + scale_b[pidx]
         o_model = (offset_a[pidx] * iso_sq + offset_b[pidx]) * digital_gain_sq

         plot_color = _PLOT_COLORS[pidx]
         subplots[pidx].plot(
             x_range,
             [o_model, s_model * _MAX_SIGNAL_VALUE + o_model],
             color=plot_color,
             linestyle='-',
             label='Model',
             alpha=0.5,
         )

         # Add the noise model measured by captures with current iso to subplots.
         if (iso, pidx) not in iso_pidx_to_measured_model_dict:
           continue

         s_measured, o_measured = iso_pidx_to_measured_model_dict[(iso, pidx)]

         subplots[pidx].plot(
             x_range,
             [o_measured, s_measured * _MAX_SIGNAL_VALUE + o_measured],
             color=plot_color,
             linestyle='--',
             label='Linear fit',
         )

         ymax = (o_measured + s_measured * xmax) * _MAX_SCALE_FUDGE
         subplots[pidx].set_xlim(xmin=0, xmax=xmax)
         subplots[pidx].set_ylim(ymin=0, ymax=ymax)
         subplots[pidx].legend()

       fig.savefig(
           f'{folder_path_prefix}_samples_iso{iso:04d}.png', dpi=_FIG_DPI
       )

   def _plot_noise_model_single_plane(
       self, pidx, plot, sens, measured_params, modeled_params):
     """Plots the noise model for one color plane specified by pidx.

     Args:
       pidx: The index of the color plane in Bayer pattern.
       plot: The ax to plot on.
       sens: The sensitivity of the sensor.
       measured_params:  The measured parameters.
       modeled_params: The modeled parameters.
     """
     color_channel = _COLOR_CHANNEL_NAMES[pidx]
     measured_label = f'{color_channel}-Measured'
     model_label = f'{color_channel}-Model'

     plot_color = _PLOT_COLORS[pidx]
     # Plot the measured parameters.
     plot.loglog(
         sens,
         measured_params,
         color=plot_color,
         marker='+',
         markeredgecolor=plot_color,
         linestyle='None',
         base=10,
         label=measured_label,
     )
     # Plot the modeled parameters.
     plot.loglog(
         sens,
         modeled_params,
         color=plot_color,
         marker='o',
         markeredgecolor=plot_color,
         linestyle='None',
         base=10,
         label=model_label,
         alpha=0.3,
     )

   def _plot_noise_model(self, isos, measured_models, noise_model,
                         sens_max_analog, name_with_log_path):
     """Plot the noise model for a given set of ISO values.

     The read noise model is defined by the following equation:
       f(x) = s_model * x + o_model
     where we have:
     s_model = scale_a * analog_gain * digital_gain + scale_b is the
     multiplicative factor,
     o_model = (offset_a * analog_gain^2 + offset_b) * digital_gain^2
     is the offset term.

     Args:
       isos: A list of ISO values.
       measured_models: A list of measured models, each of which is a tuple of
         (sens, s_measured, o_measured).
       noise_model: Noise model parameters of each plane, each of which is a
         tuple of (scale_a, scale_b, offset_a, offset_b).
       sens_max_analog: The maximum analog gain.
       name_with_log_path: The name of the file to save the logs to.
     """
     noise_model_utils.check_noise_model_shape(noise_model)

     # Plot noise model parameters.
     fig, axes = plt.subplots(4, 2, figsize=(22, 17))
     s_plots, o_plots = axes[:, 0], axes[:, 1]
     num_channels = noise_model.shape[0]
     is_quad_bayer = num_channels == noise_model_utils.NUM_QUAD_BAYER_CHANNELS
     for pidx, measured_model in enumerate(measured_models):
       # Grab the sensitivities and line parameters of each sensitivity.
       sens, s_measured, o_measured = zip(*measured_model)
       sens = np.asarray(sens)
       sens_sq = np.square(sens)
       scale_a, scale_b, offset_a, offset_b = noise_model[pidx]
       # Plot noise model components with the values predicted by the model.
       digital_gains = noise_model_utils.compute_digital_gains(
           sens, sens_max_analog
       )

       # s_model = scale_a * analog_gain * digital_gain + scale_b,
       # o_model = (offset_a * analog_gain^2 + offset_b) * digital_gain^2.
       s_model = scale_a * sens * digital_gains + scale_b
       o_model = (offset_a * sens_sq + offset_b) * np.square(digital_gains)
       if is_quad_bayer:
         s_plot, o_plot = s_plots[pidx // 4], o_plots[pidx // 4]
       else:
         s_plot, o_plot = s_plots[pidx], o_plots[pidx]

       self._plot_noise_model_single_plane(
           pidx, s_plot, sens, s_measured, s_model)
       self._plot_noise_model_single_plane(
           pidx, o_plot, sens, o_measured, o_model)

     # Set figure attributes after plotting noise model parameters.
     for s_plot, o_plot in zip(s_plots, o_plots):
       s_plot.set_xlabel('ISO')
       s_plot.set_ylabel('S')

       o_plot.set_xlabel('ISO')
       o_plot.set_ylabel('O')

       for sub_plot in (s_plot, o_plot):
         sub_plot.set_xticks(isos)
         # No minor ticks.
         sub_plot.xaxis.set_minor_locator(matplotlib.ticker.NullLocator())
         sub_plot.xaxis.set_major_formatter(matplotlib.ticker.ScalarFormatter())
         sub_plot.legend()

     fig.suptitle('Noise model: N(x) = sqrt(Sx + O)', x=0.54, y=0.99)
     pylab.tight_layout()
     fig.savefig(f'{name_with_log_path}.png', dpi=_FIG_DPI)

   def test_dng_noise_model_generation(self):
     """Calibrates standard Bayer or quad Bayer noise model.

     def requires 'test' in name to actually run.
     This function:
     * Calibrates read noise (optional).
     * Captures stats images of different ISO values and exposure times.
     * Measures linear fittings for each ISO value.
     * Computes and validates overall noise model parameters.
     * Plots noise model parameters figures.
     * Plots stats samples, linear fittings and model fittings.
     * Saves the read noise plot and csv data (optional).
     * Generates noise model and noise profile code.
     """
     read_noise_file_path = noise_model_utils.calibrate_read_noise(
         self.dut.serial,
         self.camera_id,
         self.hidden_physical_id,
         _NAME_READ_NOISE,
         _NAME_READ_NOISE_FILE,
         _READ_NOISE_STEPS_PER_STOP,
         raw_format=_READ_NOISE_RAW_FORMAT,
         is_two_stage_model=_TWO_STAGE_MODEL,
     )

     # Begin DNG Noise Model Calibration
     with its_session_utils.ItsSession(
         device_id=self.dut.serial,
         camera_id=self.camera_id,
         hidden_physical_id=self.hidden_physical_id) as cam:
       props = cam.get_camera_properties()
       props = cam.override_with_hidden_physical_camera_props(props)
       log_path = self.log_path
       name_with_log_path = os.path.join(log_path, _NAME)
       logging.info('Starting %s for camera %s', _NAME, cam.get_camera_name())

       # Get basic properties we need.
       sens_min, sens_max = props['android.sensor.info.sensitivityRange']
       sens_max_analog = props['android.sensor.maxAnalogSensitivity']
       # Maximum sensitivity for measuring noise model.
       sens_max_meas = sens_max_analog

       # Change the ISO min and/or max values if specified
       if _ISO_MIN_VALUE is not None:
         sens_min = _ISO_MIN_VALUE
       if _ISO_MAX_VALUE is not None:
         sens_max_meas = _ISO_MAX_VALUE

       logging.info('Sensitivity range: [%d, %d]', sens_min, sens_max)
       logging.info('Max analog sensitivity: %d', sens_max_analog)
       logging.info(
           'Sensitivity range for measurement: [%d, %d]',
           sens_min, sens_max_meas,
       )

       offset_a, offset_b = None, None
       read_noise_data = None
       if _TWO_STAGE_MODEL:
         # Check if read noise results exist for this device and camera
         if not os.path.exists(read_noise_file_path):
           raise AssertionError(
               'Read noise results file does not exist for this device. Run'
               ' capture_read_noise_file_path script to gather read noise data'
               ' for current sensor'
           )

         with open(read_noise_file_path, 'rb') as f:
           read_noise_data = pickle.load(f)

         offset_a, offset_b = (
             capture_read_noise_utils.get_read_noise_coefficients(
                 read_noise_data,
                 sens_min,
                 sens_max_meas,
             )
         )

       iso_to_stats_dict = noise_model_utils.capture_stats_images(
           cam,
           props,
           _STATS_CONFIG,
           sens_min,
           sens_max_meas,
           _ZOOM_RATIO,
           _TILE_CROP_N,
           _MAX_SIGNAL_VALUE,
           _ISO_MULTIPLIER,
           _BRACKET_MAX,
           _BRACKET_FACTOR,
           self.log_path,
           stats_file_name=_STATS_FILE_NAME,
           is_remove_var_outliers=_REMOVE_VAR_OUTLIERS,
           outlier_median_abs_deviations=_OUTLIER_MEDIAN_ABS_DEVS,
           is_debug_mode=self.debug_mode,
       )

     measured_models, samples = noise_model_utils.measure_linear_noise_models(
         iso_to_stats_dict,
         _COLOR_CHANNEL_NAMES,
     )

     noise_model = noise_model_utils.compute_noise_model(
         samples,
         sens_max_analog,
         offset_a,
         offset_b,
         _TWO_STAGE_MODEL,
     )

     noise_model_utils.validate_noise_model(
         noise_model,
         _COLOR_CHANNEL_NAMES,
         sens_min,
     )

     self._plot_noise_model(
         sorted(iso_to_stats_dict.keys()),
         measured_models,
         noise_model,
         sens_max_analog,
         name_with_log_path,
     )

     self._plot_stats_and_noise_model_fittings(
         iso_to_stats_dict,
         measured_models,
         noise_model,
         sens_max_analog,
         name_with_log_path,
     )

     # If 2-Stage model is enabled, save the read noise graph and csv data
     if _TWO_STAGE_MODEL:
       # Save the linear plot of the read noise data
       filename = f'{pathlib.Path(_NAME_READ_NOISE_FILE).stem}.png'
       file_path = os.path.join(log_path, filename)
       capture_read_noise_utils.plot_read_noise_data(
           read_noise_data,
           sens_min,
           sens_max_meas,
           file_path,
           _COLOR_CHANNEL_NAMES,
           _PLOT_COLORS,
       )

       # Save the data as a csv file
       filename = f'{pathlib.Path(_NAME_READ_NOISE_FILE).stem}.csv'
       file_path = os.path.join(log_path, filename)
       capture_read_noise_utils.save_read_noise_data_as_csv(
           read_noise_data,
           sens_min,
           sens_max_meas,
           file_path,
           _COLOR_CHANNEL_NAMES,
       )

     # Generate the noise model file.
     self._create_noise_model_and_profile_code(
         noise_model,
         sens_min,
         sens_max,
         sens_max_analog,
         log_path,
     )


 if __name__ == '__main__':
   test_runner.main()
	# Copyright 2014 The Android Open Source Project
	#
	# Licensed under the Apache License, Version 2.0 (the "License");
	# you may not use this file except in compliance with the License.
	# You may obtain a copy of the License at
	#
	# http://www.apache.org/licenses/LICENSE-2.0
	#
	# Unless required by applicable law or agreed to in writing, software
	# distributed under the License is distributed on an "AS IS" BASIS,
	# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	# See the License for the specific language governing permissions and
	# limitations under the License.
	"""CameraITS script to generate noise models."""

	import logging
	import math
	import os.path
	import pathlib
	import pickle
	import tempfile
	import textwrap

	import capture_read_noise_utils
	import its_base_test
	import its_session_utils
	from matplotlib import pylab
	import matplotlib.pyplot as plt
	import matplotlib.ticker
	from mobly import test_runner
	import noise_model_utils
	import numpy as np

	_IS_QUAD_BAYER = False # A manual flag to choose standard or quad Bayer noise
	# model generation.
	if _IS_QUAD_BAYER:
	_COLOR_CHANNEL_NAMES = noise_model_utils.QUAD_BAYER_COLORS
	_PLOT_COLORS = noise_model_utils.QUAD_BAYER_PLOT_COLORS
	_TILE_SIZE = 64 # Tile size to compute mean/variance. Large tiles may have
	# their variance corrupted by low freq image changes.
	_STATS_FORMAT = 'raw10QuadBayerStats' # rawQuadBayerStats\|raw10QuadBayerStats
	_READ_NOISE_RAW_FORMAT = 'raw10QuadBayer' # rawQuadBayer\|raw10QuadBayer
	else:
	_COLOR_CHANNEL_NAMES = noise_model_utils.BAYER_COLORS
	_PLOT_COLORS = noise_model_utils.BAYER_PLOT_COLORS
	_TILE_SIZE = 32 # Tile size to compute mean/variance. Large tiles may have
	# their variance corrupted by low freq image changes.
	_STATS_FORMAT = 'rawStats' # rawStats\|raw10Stats
	_READ_NOISE_RAW_FORMAT = 'raw' # raw\|raw10

	_STATS_CONFIG = {
	'format': _STATS_FORMAT,
	'gridWidth': _TILE_SIZE,
	'gridHeight': _TILE_SIZE,
	}
	_BRACKET_MAX = 8 # Exposure bracketing range in stops
	_BRACKET_FACTOR = math.pow(2, _BRACKET_MAX)
	_ISO_MAX_VALUE = None # ISO range max value, uses sensor max if None
	_ISO_MIN_VALUE = None # ISO range min value, uses sensor min if None
	_MAX_SCALE_FUDGE = 1.1
	_MAX_SIGNAL_VALUE = 0.25 # Maximum value to allow mean of the tiles to go.
	_NAME = os.path.basename(__file__).split('.')[0]
	_NAME_READ_NOISE = os.path.join(tempfile.gettempdir(), 'CameraITS/ReadNoise')
	_NAME_READ_NOISE_FILE = 'read_noise_results.pkl'
	_STATS_FILE_NAME = 'stats.pkl'
	_OUTLIER_MEDIAN_ABS_DEVS = 10 # Defines the number of Median Absolute
	# Deviations that constitutes acceptable data
	_READ_NOISE_STEPS_PER_STOP = 12 # Sensitivities per stop to sample for read
	# noise
	_REMOVE_VAR_OUTLIERS = False # When True, filters the variance to remove
	# outliers
	_STEPS_PER_STOP = 3 # How many sensitivities per stop to sample.
	_ISO_MULTIPLIER = math.pow(2, 1.0 / _STEPS_PER_STOP)
	_TILE_CROP_N = 0 # Number of tiles to crop from edge of image. Usually 0.
	_TWO_STAGE_MODEL = False # Require read noise data prior to running noise model
	_ZOOM_RATIO = 1 # Zoom target to be used while running the model
	_FIG_DPI = 100 # DPI for plotting noise model figures.
	_BAYER_COLORS_FOR_NOISE_PROFILE = [color.lower() for color in
	noise_model_utils.BAYER_COLORS]
	_QUAD_BAYER_COLORS_FOR_NOISE_PROFILE = [color.lower() for color in
	noise_model_utils.QUAD_BAYER_COLORS]


	class DngNoiseModel(its_base_test.ItsBaseTest):
	"""Create DNG noise model.

	Captures RAW images with increasing analog gains to create the model.
	"""

	def _create_noise_model_code(self, noise_model, sens_min, sens_max,
	sens_max_analog, file_path):
	"""Creates the C file for the noise model.

	Args:
	noise_model: Noise model parameters.
	sens_min: The minimum sensitivity value.
	sens_max: The maximum sensitivity value.
	sens_max_analog: The maximum analog sensitivity value.
	file_path: The path to the noise model file.
	"""
	# Generate individual noise model components.
	scale_a, scale_b, offset_a, offset_b = zip(*noise_model)
	digital_gain_cdef = (
	f'(sens / {sens_max_analog:.1f}) < 1.0 ? '
	f'1.0 : (sens / {sens_max_analog:.1f})'
	)

	with open(file_path, 'w') as text_file:
	scale_a_str = ','.join([str(i) for i in scale_a])
	scale_b_str = ','.join([str(i) for i in scale_b])
	offset_a_str = ','.join([str(i) for i in offset_a])
	offset_b_str = ','.join([str(i) for i in offset_b])
	# pylint: disable=line-too-long
	code = textwrap.dedent(f"""\
	/* Generated test code to dump a table of data for external validation
	* of the noise model parameters.
	*/
	#include <stdio.h>
	#include <assert.h>
	double compute_noise_model_entry_S(int plane, int sens);
	double compute_noise_model_entry_O(int plane, int sens);
	int main(void) {{
	for (int plane = 0; plane < {len(scale_a)}; plane++) {{
	for (int sens = {sens_min}; sens <= {sens_max}; sens += 100) {{
	double o = compute_noise_model_entry_O(plane, sens);
	double s = compute_noise_model_entry_S(plane, sens);
	printf("%d,%d,%lf,%lf\\n", plane, sens, o, s);
	}}
	}}
	return 0;
	}}

	/* Generated functions to map a given sensitivity to the O and S noise
	* model parameters in the DNG noise model. The planes are in
	* R, Gr, Gb, B order.
	*/
	double compute_noise_model_entry_S(int plane, int sens) {{
	static double noise_model_A[] = {{ {scale_a_str:s} }};
	static double noise_model_B[] = {{ {scale_b_str:s} }};
	double A = noise_model_A[plane];
	double B = noise_model_B[plane];
	double s = A * sens + B;
	return s < 0.0 ? 0.0 : s;
	}}

	double compute_noise_model_entry_O(int plane, int sens) {{
	static double noise_model_C[] = {{ {offset_a_str:s} }};
	static double noise_model_D[] = {{ {offset_b_str:s} }};
	double digital_gain = {digital_gain_cdef:s};
	double C = noise_model_C[plane];
	double D = noise_model_D[plane];
	double o = C * sens * sens + D * digital_gain * digital_gain;
	return o < 0.0 ? 0.0 : o;
	}}
	""")

	text_file.write(code)

	def _create_noise_profile_code(self, noise_model, color_channels, file_path):
	"""Creates the noise profile C++ file.

	Args:
	noise_model: Noise model parameters.
	color_channels: Color channels in canonical order.
	file_path: The path to the noise profile C++ file.
	"""
	# Generate individual noise model components.
	scale_a, scale_b, offset_a, offset_b = zip(*noise_model)
	num_channels = noise_model.shape[0]
	params = []
	for ch, color in enumerate(color_channels):
	prefix = f'.noise_coefficients_{color} = {{'
	spaces = ' ' * len(prefix)
	suffix = '},' if ch != num_channels - 1 else '}'
	params.append(textwrap.dedent(f"""
	{prefix}.gradient_slope = {scale_a[ch]},
	{spaces}.offset_slope = {scale_b[ch]},
	{spaces}.gradient_intercept = {offset_a[ch]},
	{spaces}.offset_intercept = {offset_b[ch]}{suffix}"""))

	with open(file_path, 'w') as text_file:
	# pylint: disable=line-too-long
	code_comment = textwrap.dedent("""\
	/* noise_profile.cc
	Note: gradient_slope --> gradient of API s_measured parameter
	offset_slope --> o_model of API s_measured parameter
	gradient_intercept--> gradient of API o_measured parameter
	offset_intercept --> o_model of API o_measured parameter
	Note: SENSOR_NOISE_PROFILE in Android Developers doc uses
	N(x) = sqrt(Sx + O), where 'S' is 's_measured' & 'O' is 'o_measured'
	*/
	""")
	params_str = textwrap.indent(''.join(params), ' ' * 4)
	code_params = '.noise_profile = {' + params_str + '},'
	code = code_comment + code_params
	text_file.write(code)

	def _create_noise_model_and_profile_code(self, noise_model, sens_min,
	sens_max, sens_max_analog, log_path):
	"""Creates the code file with noise model parameters.

	Args:
	noise_model: Noise model parameters.
	sens_min: The minimum sensitivity value.
	sens_max: The maximum sensitivity value.
	sens_max_analog: The maximum analog sensitivity value.
	log_path: The path to the log file.
	"""
	noise_model_utils.check_noise_model_shape(noise_model)
	# Create noise model code with noise model parameters.
	self._create_noise_model_code(
	noise_model,
	sens_min,
	sens_max,
	sens_max_analog,
	os.path.join(log_path, 'noise_model.c'),
	)

	num_channels = noise_model.shape[0]
	is_quad_bayer = num_channels == noise_model_utils.NUM_QUAD_BAYER_CHANNELS
	if is_quad_bayer:
	# Average noise model parameters of every four channels.
	avg_noise_model = noise_model.reshape(-1, 4, noise_model.shape[1]).mean(
	axis=1
	)
	# Create noise model code with average noise model parameters.
	self._create_noise_model_code(
	avg_noise_model,
	sens_min,
	sens_max,
	sens_max_analog,
	os.path.join(log_path, 'noise_model_avg.c'),
	)
	# Create noise profile code with average noise model parameters.
	self._create_noise_profile_code(
	avg_noise_model,
	_BAYER_COLORS_FOR_NOISE_PROFILE,
	os.path.join(log_path, 'noise_profile_avg.cc'),
	)
	# Create noise profile code with noise model parameters.
	self._create_noise_profile_code(
	noise_model,
	_QUAD_BAYER_COLORS_FOR_NOISE_PROFILE,
	os.path.join(log_path, 'noise_profile.cc'),
	)

	else:
	# Create noise profile code with noise model parameters.
	self._create_noise_profile_code(
	noise_model,
	_BAYER_COLORS_FOR_NOISE_PROFILE,
	os.path.join(log_path, 'noise_profile.cc'),
	)

	def _plot_stats_and_noise_model_fittings(
	self, iso_to_stats_dict, measured_models, noise_model, sens_max_analog,
	folder_path_prefix):
	"""Plots the stats (means, vars_) and noise models fittings.

	Args:
	iso_to_stats_dict: A dictionary mapping ISO to a list of tuples of
	exposure time in milliseconds, mean values, and variance values.
	measured_models: A list of measured noise models for each ISO value.
	noise_model: A numpy array of global noise model parameters for all ISO
	values.
	sens_max_analog: The maximum analog sensitivity value.
	folder_path_prefix: The prefix of path to save figures.

	Raises:
	ValueError: If the noise model shape is invalid.
	"""
	noise_model_utils.check_noise_model_shape(noise_model)
	# Separate individual noise model components.
	scale_a, scale_b, offset_a, offset_b = zip(*noise_model)

	iso_pidx_to_measured_model_dict = {}
	num_channels = noise_model.shape[0]
	for pidx in range(num_channels):
	for iso, s_measured, o_measured in measured_models[pidx]:
	iso_pidx_to_measured_model_dict[(iso, pidx)] = (s_measured, o_measured)

	isos = np.asarray(sorted(iso_to_stats_dict.keys()))
	digital_gains = noise_model_utils.compute_digital_gains(
	isos, sens_max_analog
	)

	x_range = [0, _MAX_SIGNAL_VALUE]
	for iso, digital_gain in zip(isos, digital_gains):
	logging.info('Plotting stats and noise model for ISO %d.', iso)
	fig, subplots = noise_model_utils.create_stats_figure(
	iso, _COLOR_CHANNEL_NAMES
	)

	xmax = 0
	stats_per_plane = [[] for _ in range(num_channels)]
	for exposure_ms, means, vars_ in iso_to_stats_dict[iso]:
	exposure_norm = noise_model_utils.COLOR_NORM(np.log2(exposure_ms))
	exposure_color = noise_model_utils.RAINBOW_CMAP(exposure_norm)
	for pidx in range(num_channels):
	means_p = means[pidx]
	vars_p = vars_[pidx]

	if means_p.size > 0 and vars_p.size > 0:
	subplots[pidx].plot(
	means_p,
	vars_p,
	color=exposure_color,
	marker='.',
	markeredgecolor=exposure_color,
	markersize=1,
	linestyle='None',
	alpha=0.5,
	)

	stats_per_plane[pidx].extend(list(zip(means_p, vars_p)))
	xmax = max(xmax, max(means_p))

	iso_sq = iso ** 2
	digital_gain_sq = digital_gain ** 2
	for pidx in range(num_channels):
	# Add the final noise model to subplots.
	s_model = scale_a[pidx] * iso * digital_gain + scale_b[pidx]
	o_model = (offset_a[pidx] * iso_sq + offset_b[pidx]) * digital_gain_sq

	plot_color = _PLOT_COLORS[pidx]
	subplots[pidx].plot(
	x_range,
	[o_model, s_model * _MAX_SIGNAL_VALUE + o_model],
	color=plot_color,
	linestyle='-',
	label='Model',
	alpha=0.5,
	)

	# Add the noise model measured by captures with current iso to subplots.
	if (iso, pidx) not in iso_pidx_to_measured_model_dict:
	continue

	s_measured, o_measured = iso_pidx_to_measured_model_dict[(iso, pidx)]

	subplots[pidx].plot(
	x_range,
	[o_measured, s_measured * _MAX_SIGNAL_VALUE + o_measured],
	color=plot_color,
	linestyle='--',
	label='Linear fit',
	)

	ymax = (o_measured + s_measured * xmax) * _MAX_SCALE_FUDGE
	subplots[pidx].set_xlim(xmin=0, xmax=xmax)
	subplots[pidx].set_ylim(ymin=0, ymax=ymax)
	subplots[pidx].legend()

	fig.savefig(
	f'{folder_path_prefix}_samples_iso{iso:04d}.png', dpi=_FIG_DPI
	)

	def _plot_noise_model_single_plane(
	self, pidx, plot, sens, measured_params, modeled_params):
	"""Plots the noise model for one color plane specified by pidx.

	Args:
	pidx: The index of the color plane in Bayer pattern.
	plot: The ax to plot on.
	sens: The sensitivity of the sensor.
	measured_params: The measured parameters.
	modeled_params: The modeled parameters.
	"""
	color_channel = _COLOR_CHANNEL_NAMES[pidx]
	measured_label = f'{color_channel}-Measured'
	model_label = f'{color_channel}-Model'

	plot_color = _PLOT_COLORS[pidx]
	# Plot the measured parameters.
	plot.loglog(
	sens,
	measured_params,
	color=plot_color,
	marker='+',
	markeredgecolor=plot_color,
	linestyle='None',
	base=10,
	label=measured_label,
	)
	# Plot the modeled parameters.
	plot.loglog(
	sens,
	modeled_params,
	color=plot_color,
	marker='o',
	markeredgecolor=plot_color,
	linestyle='None',
	base=10,
	label=model_label,
	alpha=0.3,
	)

	def _plot_noise_model(self, isos, measured_models, noise_model,
	sens_max_analog, name_with_log_path):
	"""Plot the noise model for a given set of ISO values.

	The read noise model is defined by the following equation:
	f(x) = s_model * x + o_model
	where we have:
	s_model = scale_a * analog_gain * digital_gain + scale_b is the
	multiplicative factor,
	o_model = (offset_a * analog_gain^2 + offset_b) * digital_gain^2
	is the offset term.

	Args:
	isos: A list of ISO values.
	measured_models: A list of measured models, each of which is a tuple of
	(sens, s_measured, o_measured).
	noise_model: Noise model parameters of each plane, each of which is a
	tuple of (scale_a, scale_b, offset_a, offset_b).
	sens_max_analog: The maximum analog gain.
	name_with_log_path: The name of the file to save the logs to.
	"""
	noise_model_utils.check_noise_model_shape(noise_model)

	# Plot noise model parameters.
	fig, axes = plt.subplots(4, 2, figsize=(22, 17))
	s_plots, o_plots = axes[:, 0], axes[:, 1]
	num_channels = noise_model.shape[0]
	is_quad_bayer = num_channels == noise_model_utils.NUM_QUAD_BAYER_CHANNELS
	for pidx, measured_model in enumerate(measured_models):
	# Grab the sensitivities and line parameters of each sensitivity.
	sens, s_measured, o_measured = zip(*measured_model)
	sens = np.asarray(sens)
	sens_sq = np.square(sens)
	scale_a, scale_b, offset_a, offset_b = noise_model[pidx]
	# Plot noise model components with the values predicted by the model.
	digital_gains = noise_model_utils.compute_digital_gains(
	sens, sens_max_analog
	)

	# s_model = scale_a * analog_gain * digital_gain + scale_b,
	# o_model = (offset_a * analog_gain^2 + offset_b) * digital_gain^2.
	s_model = scale_a * sens * digital_gains + scale_b
	o_model = (offset_a * sens_sq + offset_b) * np.square(digital_gains)
	if is_quad_bayer:
	s_plot, o_plot = s_plots[pidx // 4], o_plots[pidx // 4]
	else:
	s_plot, o_plot = s_plots[pidx], o_plots[pidx]

	self._plot_noise_model_single_plane(
	pidx, s_plot, sens, s_measured, s_model)
	self._plot_noise_model_single_plane(
	pidx, o_plot, sens, o_measured, o_model)

	# Set figure attributes after plotting noise model parameters.
	for s_plot, o_plot in zip(s_plots, o_plots):
	s_plot.set_xlabel('ISO')
	s_plot.set_ylabel('S')

	o_plot.set_xlabel('ISO')
	o_plot.set_ylabel('O')

	for sub_plot in (s_plot, o_plot):
	sub_plot.set_xticks(isos)
	# No minor ticks.
	sub_plot.xaxis.set_minor_locator(matplotlib.ticker.NullLocator())
	sub_plot.xaxis.set_major_formatter(matplotlib.ticker.ScalarFormatter())
	sub_plot.legend()

	fig.suptitle('Noise model: N(x) = sqrt(Sx + O)', x=0.54, y=0.99)
	pylab.tight_layout()
	fig.savefig(f'{name_with_log_path}.png', dpi=_FIG_DPI)

	def test_dng_noise_model_generation(self):
	"""Calibrates standard Bayer or quad Bayer noise model.

	def requires 'test' in name to actually run.
	This function:
	* Calibrates read noise (optional).
	* Captures stats images of different ISO values and exposure times.
	* Measures linear fittings for each ISO value.
	* Computes and validates overall noise model parameters.
	* Plots noise model parameters figures.
	* Plots stats samples, linear fittings and model fittings.
	* Saves the read noise plot and csv data (optional).
	* Generates noise model and noise profile code.
	"""
	read_noise_file_path = noise_model_utils.calibrate_read_noise(
	self.dut.serial,
	self.camera_id,
	self.hidden_physical_id,
	_NAME_READ_NOISE,
	_NAME_READ_NOISE_FILE,
	_READ_NOISE_STEPS_PER_STOP,
	raw_format=_READ_NOISE_RAW_FORMAT,
	is_two_stage_model=_TWO_STAGE_MODEL,
	)

	# Begin DNG Noise Model Calibration
	with its_session_utils.ItsSession(
	device_id=self.dut.serial,
	camera_id=self.camera_id,
	hidden_physical_id=self.hidden_physical_id) as cam:
	props = cam.get_camera_properties()
	props = cam.override_with_hidden_physical_camera_props(props)
	log_path = self.log_path
	name_with_log_path = os.path.join(log_path, _NAME)
	logging.info('Starting %s for camera %s', _NAME, cam.get_camera_name())

	# Get basic properties we need.
	sens_min, sens_max = props['android.sensor.info.sensitivityRange']
	sens_max_analog = props['android.sensor.maxAnalogSensitivity']
	# Maximum sensitivity for measuring noise model.
	sens_max_meas = sens_max_analog

	# Change the ISO min and/or max values if specified
	if _ISO_MIN_VALUE is not None:
	sens_min = _ISO_MIN_VALUE
	if _ISO_MAX_VALUE is not None:
	sens_max_meas = _ISO_MAX_VALUE

	logging.info('Sensitivity range: [%d, %d]', sens_min, sens_max)
	logging.info('Max analog sensitivity: %d', sens_max_analog)
	logging.info(
	'Sensitivity range for measurement: [%d, %d]',
	sens_min, sens_max_meas,
	)

	offset_a, offset_b = None, None
	read_noise_data = None
	if _TWO_STAGE_MODEL:
	# Check if read noise results exist for this device and camera
	if not os.path.exists(read_noise_file_path):
	raise AssertionError(
	'Read noise results file does not exist for this device. Run'
	' capture_read_noise_file_path script to gather read noise data'
	' for current sensor'
	)

	with open(read_noise_file_path, 'rb') as f:
	read_noise_data = pickle.load(f)

	offset_a, offset_b = (
	capture_read_noise_utils.get_read_noise_coefficients(
	read_noise_data,
	sens_min,
	sens_max_meas,
	)
	)

	iso_to_stats_dict = noise_model_utils.capture_stats_images(
	cam,
	props,
	_STATS_CONFIG,
	sens_min,
	sens_max_meas,
	_ZOOM_RATIO,
	_TILE_CROP_N,
	_MAX_SIGNAL_VALUE,
	_ISO_MULTIPLIER,
	_BRACKET_MAX,
	_BRACKET_FACTOR,
	self.log_path,
	stats_file_name=_STATS_FILE_NAME,
	is_remove_var_outliers=_REMOVE_VAR_OUTLIERS,
	outlier_median_abs_deviations=_OUTLIER_MEDIAN_ABS_DEVS,
	is_debug_mode=self.debug_mode,
	)

	measured_models, samples = noise_model_utils.measure_linear_noise_models(
	iso_to_stats_dict,
	_COLOR_CHANNEL_NAMES,
	)

	noise_model = noise_model_utils.compute_noise_model(
	samples,
	sens_max_analog,
	offset_a,
	offset_b,
	_TWO_STAGE_MODEL,
	)

	noise_model_utils.validate_noise_model(
	noise_model,
	_COLOR_CHANNEL_NAMES,
	sens_min,
	)

	self._plot_noise_model(
	sorted(iso_to_stats_dict.keys()),
	measured_models,
	noise_model,
	sens_max_analog,
	name_with_log_path,
	)

	self._plot_stats_and_noise_model_fittings(
	iso_to_stats_dict,
	measured_models,
	noise_model,
	sens_max_analog,
	name_with_log_path,
	)

	# If 2-Stage model is enabled, save the read noise graph and csv data
	if _TWO_STAGE_MODEL:
	# Save the linear plot of the read noise data
	filename = f'{pathlib.Path(_NAME_READ_NOISE_FILE).stem}.png'
	file_path = os.path.join(log_path, filename)
	capture_read_noise_utils.plot_read_noise_data(
	read_noise_data,
	sens_min,
	sens_max_meas,
	file_path,
	_COLOR_CHANNEL_NAMES,
	_PLOT_COLORS,
	)

	# Save the data as a csv file
	filename = f'{pathlib.Path(_NAME_READ_NOISE_FILE).stem}.csv'
	file_path = os.path.join(log_path, filename)
	capture_read_noise_utils.save_read_noise_data_as_csv(
	read_noise_data,
	sens_min,
	sens_max_meas,
	file_path,
	_COLOR_CHANNEL_NAMES,
	)

	# Generate the noise model file.
	self._create_noise_model_and_profile_code(
	noise_model,
	sens_min,
	sens_max,
	sens_max_analog,
	log_path,
	)


	if __name__ == '__main__':
	test_runner.main()