• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# Copyright 2014 The Android Open Source Project
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7#      http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14
15import logging
16import math
17import os.path
18from pathlib import Path
19import pickle
20import tempfile
21import textwrap
22
23import capture_read_noise_utils
24import capture_request_utils
25import image_processing_utils
26import its_base_test
27import its_session_utils
28from matplotlib import pylab
29import matplotlib.pyplot as plt
30from matplotlib.ticker import NullLocator, ScalarFormatter
31from mobly import test_runner
32import numpy as np
33import scipy.signal
34import scipy.stats
35
36
37_BAYER_LIST = ('R', 'GR', 'GB', 'B')  # List of Bayer colors
38_BAYER_COLOR_FILTERS = [  # Bayer filters (SENSOR_INFO_COLOR_FILTER_ARRANGEMENT)
39    'RGGB',  # 0
40    'GRBG',  # 1
41    'GBRG',  # 2
42    'BGGR',  # 3
43    'RBG',   # 4
44    'MONO',  # 5
45    'NIR'    # 6
46]
47_BRACKET_MAX = 8  # Exposure bracketing range in stops
48_BRACKET_FACTOR = math.pow(2, _BRACKET_MAX)
49_ISO_MAX_VALUE = None  # ISO range max value, uses sensor max if None
50_ISO_MIN_VALUE = None  # ISO range min value, uses sensor min if None
51_PLOT_COLORS = 'rygcbm'  # Colors used for plotting the data for each exposure.
52_MAX_SCALE_FUDGE = 1.1
53_MAX_SIGNAL_VALUE = 0.25  # Maximum value to allow mean of the tiles to go.
54_NAME = os.path.basename(__file__).split('.')[0]
55_NAME_READ_NOISE = os.path.join(tempfile.gettempdir(), 'CameraITS/ReadNoise')
56_NAME_READ_NOISE_FILE = 'read_noise_results.pkl'
57_OUTLIE_MEDIAN_ABS_DEVS = 10  # Defines the number of Median Absolute Deviations
58                              # that consitutes acceptable data
59_READ_NOISE_STEPS_PER_STOP = 12  # Sensitivities per stop to sample for read
60                                 # noise
61_REMOVE_OUTLIERS = False  # When True, filters the variance to remove outliers
62_STEPS_PER_STOP = 3  # How many sensitivities per stop to sample.
63_TILE_SIZE = 32  # Tile size to compute mean/variance. Large tiles may have
64                 # their variance corrupted by low freq image changes.
65_TILE_CROP_N = 0  # Number of tiles to crop from edge of image. Usually 0.
66_TWO_STAGE_MODEL = False  # Require read noise data prior to running noise model
67_ZOOM_RATIO = 1  # Zoom target to be used while running the model
68
69
70def check_auto_exposure_targets(auto_e, sens_min, sens_max, props):
71  """Checks if AE too bright for highest gain & too dark for lowest gain."""
72
73  min_exposure_ns, max_exposure_ns = props[
74      'android.sensor.info.exposureTimeRange']
75  if auto_e < min_exposure_ns*sens_max:
76    raise AssertionError('Scene is too bright to properly expose at highest '
77                         f'sensitivity: {sens_max}')
78  if auto_e*_BRACKET_FACTOR > max_exposure_ns*sens_min:
79    raise AssertionError('Scene is too dark to properly expose at lowest '
80                         f'sensitivity: {sens_min}')
81
82
83def create_noise_model_code(noise_model_a, noise_model_b,
84                            noise_model_c, noise_model_d,
85                            sens_min, sens_max, digital_gain_cdef, log_path):
86  """Creates the c file for the noise model."""
87
88  noise_model_a_array = ','.join([str(i) for i in noise_model_a])
89  noise_model_b_array = ','.join([str(i) for i in noise_model_b])
90  noise_model_c_array = ','.join([str(i) for i in noise_model_c])
91  noise_model_d_array = ','.join([str(i) for i in noise_model_d])
92  code = textwrap.dedent(f"""\
93          /* Generated test code to dump a table of data for external validation
94           * of the noise model parameters.
95           */
96          #include <stdio.h>
97          #include <assert.h>
98          double compute_noise_model_entry_S(int plane, int sens);
99          double compute_noise_model_entry_O(int plane, int sens);
100          int main(void) {{
101              for (int plane = 0; plane < {len(noise_model_a)}; plane++) {{
102                  for (int sens = {sens_min}; sens <= {sens_max}; sens += 100) {{
103                      double o = compute_noise_model_entry_O(plane, sens);
104                      double s = compute_noise_model_entry_S(plane, sens);
105                      printf("%d,%d,%lf,%lf\\n", plane, sens, o, s);
106                  }}
107              }}
108              return 0;
109          }}
110
111          /* Generated functions to map a given sensitivity to the O and S noise
112           * model parameters in the DNG noise model. The planes are in
113           * R, Gr, Gb, B order.
114           */
115          double compute_noise_model_entry_S(int plane, int sens) {{
116              static double noise_model_A[] = {{ {noise_model_a_array:s} }};
117              static double noise_model_B[] = {{ {noise_model_b_array:s} }};
118              double A = noise_model_A[plane];
119              double B = noise_model_B[plane];
120              double s = A * sens + B;
121              return s < 0.0 ? 0.0 : s;
122          }}
123
124          double compute_noise_model_entry_O(int plane, int sens) {{
125              static double noise_model_C[] = {{ {noise_model_c_array:s} }};
126              static double noise_model_D[] = {{ {noise_model_d_array:s} }};
127              double digital_gain = {digital_gain_cdef:s};
128              double C = noise_model_C[plane];
129              double D = noise_model_D[plane];
130              double o = C * sens * sens + D * digital_gain * digital_gain;
131              return o < 0.0 ? 0.0 : o;
132          }}
133          """)
134  text_file = open(os.path.join(log_path, 'noise_model.c'), 'w')
135  text_file.write(code)
136  text_file.close()
137
138  # Creates the noise profile C++ file
139  code = textwrap.dedent(f"""\
140          /* noise_profile.cc
141             Note: gradient_slope --> gradient of API slope parameter
142                   offset_slope --> offset of API slope parameter
143                   gradient_intercept--> gradient of API intercept parameter
144                   offset_intercept --> offset of API intercept parameter
145             Note: SENSOR_NOISE_PROFILE in Android Developers doc uses
146                   N(x) = sqrt(Sx + O), where 'S' is 'slope' & 'O' is 'intercept'
147          */
148          .noise_profile =
149              {{.noise_coefficients_r = {{.gradient_slope = {noise_model_a[0]},
150                                        .offset_slope = {noise_model_b[0]},
151                                        .gradient_intercept = {noise_model_c[0]},
152                                        .offset_intercept = {noise_model_d[0]}}},
153               .noise_coefficients_gr = {{.gradient_slope = {noise_model_a[1]},
154                                         .offset_slope = {noise_model_b[1]},
155                                         .gradient_intercept = {noise_model_c[1]},
156                                         .offset_intercept = {noise_model_d[1]}}},
157               .noise_coefficients_gb = {{.gradient_slope = {noise_model_a[2]},
158                                         .offset_slope = {noise_model_b[2]},
159                                         .gradient_intercept = {noise_model_c[2]},
160                                         .offset_intercept = {noise_model_d[2]}}},
161               .noise_coefficients_b = {{.gradient_slope = {noise_model_a[3]},
162                                        .offset_slope = {noise_model_b[3]},
163                                        .gradient_intercept = {noise_model_c[3]},
164                                        .offset_intercept = {noise_model_d[3]}}}}},
165          """)
166  text_file = open(os.path.join(log_path, 'noise_profile.cc'), 'w')
167  text_file.write(code)
168  text_file.close()
169
170
171def outlier_removed_indices(data, deviations=3):
172  """Removes outliers using median absolute deviation and returns indices kept.
173
174  Args:
175      data:             list to remove outliers from
176      deviations:       number of deviations from median to keep
177  Returns:
178      keep_indices:     The indices of data which should be kept
179  """
180  std_dev = scipy.stats.median_abs_deviation(data, axis=None, scale=1)
181  med = np.median(data)
182  keep_indices = np.where(
183      np.logical_and(data > med-deviations*std_dev,
184                     data < med+deviations*std_dev))
185  return keep_indices
186
187
188def reorganize_read_noise_coeff_to_rggb(data, cmap):
189  """Reorganize the list of read noise coeffs to RGGB Bayer form.
190
191  Args:
192    data:      list; List of color channel data
193    cmap:      str; Color map filter of the given data
194  Returns:
195    list       Returns data reformatted in the correct RGGB order
196  """
197  if cmap not in _BAYER_COLOR_FILTERS:
198    raise AssertionError(f'Unexpected color map {cmap}')
199
200  if cmap.lower() == 'rggb':
201    return data
202  if cmap.lower() == 'grbg':
203    return [data[1], data[0], data[3], data[2]]
204  if cmap.lower() == 'gbrg':
205    return [data[2], data[3], data[0], data[1]]
206  if cmap.lower() == 'bggr':
207    return [data[3], data[2], data[1], data[0]]
208  else:
209    raise AssertionError(
210        'Currently only 4-channel filters supported for 2-Stage model')
211
212
213class DngNoiseModel(its_base_test.ItsBaseTest):
214  """Create DNG noise model.
215
216  Captures RAW images with increasing analog gains to create the model.
217  def requires 'test' in name to actually run.
218  """
219
220  def test_dng_noise_model_generation(self):
221    read_noise_folder = ''
222    read_noise_data = ''
223
224    # If 2-Stage model is enabled, check/collect read noise data
225    if _TWO_STAGE_MODEL:
226      with its_session_utils.ItsSession(
227          device_id=self.dut.serial,
228          camera_id=self.camera_id,
229          hidden_physical_id=self.hidden_physical_id) as cam:
230        props = cam.get_camera_properties()
231        props = cam.override_with_hidden_physical_camera_props(props)
232
233        # Get sensor ISO range
234        sens_min, _ = props['android.sensor.info.sensitivityRange']
235        sens_max_analog = props['android.sensor.maxAnalogSensitivity']
236        color_map_index = props['android.sensor.info.colorFilterArrangement']
237        color_map = _BAYER_COLOR_FILTERS[color_map_index]
238
239        sens_max_meas = sens_max_analog
240
241        # Create the folder structure
242        if self.hidden_physical_id:
243          camera_name = f'{self.camera_id}.{self.hidden_physical_id}'
244        else:
245          camera_name = self.camera_id
246
247        read_noise_folder = os.path.join(_NAME_READ_NOISE,
248                                         self.dut.serial.replace(':', '_'),
249                                         camera_name)
250        read_noise_data = os.path.join(read_noise_folder,
251                                       _NAME_READ_NOISE_FILE)
252
253        if not os.path.exists(read_noise_folder):
254          os.makedirs(read_noise_folder)
255
256        logging.info('Read noise data folder: %s', read_noise_folder)
257        # Collect or retrieve read noise data
258        if not os.path.isfile(read_noise_data):
259          # Wait until camera is repositioned for read noise data collection
260          input(f'\nPress <ENTER> after concealing camera {self.camera_id} in'
261                ' complete darkness.\n')
262          logging.info('Collecting read noise data for %s', self.camera_id)
263          # Results file does not exist, collect read noise data
264          capture_read_noise_utils.capture_read_noise_for_iso_range(
265              cam, sens_min, sens_max_meas, _READ_NOISE_STEPS_PER_STOP,
266              color_map, read_noise_data)
267        else:  # If data exists, check if it covers the full range
268          with open(read_noise_data, 'rb') as f:
269            results = pickle.load(f)
270            # The +5 offset takes write to read error into account
271            if results[-1][0]['iso'] + 50 < sens_max_meas:
272              logging.info('\nNot enough ISO data points exist. '
273                           '\nMax ISO measured: %.2f'
274                           '\nMax ISO possible: %.2f',
275                           results[-1][0]['iso'],
276                           sens_max_meas)
277              # Wait until camera is repositioned for read noise data collection
278              input(f'\nPress <ENTER> after concealing camera {self.camera_id} '
279                    'in complete darkness.\n')
280              # Not all data points were captured, continue capture
281              capture_read_noise_utils.capture_read_noise_for_iso_range(
282                  cam, sens_min, sens_max_meas, _READ_NOISE_STEPS_PER_STOP,
283                  color_map, read_noise_data)
284
285    # Begin DNG Noise Model Calibration
286    with its_session_utils.ItsSession(
287        device_id=self.dut.serial,
288        camera_id=self.camera_id,
289        hidden_physical_id=self.hidden_physical_id) as cam:
290      props = cam.get_camera_properties()
291      props = cam.override_with_hidden_physical_camera_props(props)
292      log_path = self.log_path
293      name_with_log_path = os.path.join(log_path, _NAME)
294      if self.hidden_physical_id:
295        camera_name = f'{self.camera_id}.{self.hidden_physical_id}'
296      else:
297        camera_name = self.camera_id
298      logging.info('Starting %s for camera %s', _NAME, camera_name)
299
300      # Get basic properties we need.
301      sens_min, sens_max = props['android.sensor.info.sensitivityRange']
302      sens_max_analog = props['android.sensor.maxAnalogSensitivity']
303      sens_max_meas = sens_max_analog
304      white_level = props['android.sensor.info.whiteLevel']
305      min_exposure_ns, _ = props[
306          'android.sensor.info.exposureTimeRange']
307
308      # Change the ISO min and/or max values if specified
309      if _ISO_MIN_VALUE is not None:
310        sens_min = _ISO_MIN_VALUE
311      if _ISO_MAX_VALUE is not None:
312        sens_max_meas = _ISO_MAX_VALUE
313
314      # Wait until camera is repositioned for DNG noise model calibration
315      input(f'\nPress <ENTER> after covering camera lense {self.camera_id} with'
316            ' frosted glass diffuser, and facing lense at evenly illuminated'
317            ' surface.\n')
318
319      if _TWO_STAGE_MODEL:
320        # Check if read noise results exist for this device and camera
321        if not os.path.exists(read_noise_data):
322          raise AssertionError(
323              'Read noise results file does not exist for this device. Run '
324              'capture_read_noise_data script to gather read noise data for '
325              'current sensor')
326
327        color_map_index = props['android.sensor.info.colorFilterArrangement']
328        color_map = _BAYER_COLOR_FILTERS[color_map_index]
329
330        with open(read_noise_data, 'rb') as f:
331          read_noise_results = pickle.load(f)
332
333        coeff_a, coeff_b = capture_read_noise_utils.get_read_noise_coefficients(
334            read_noise_results, sens_min, sens_max_meas)
335
336        coeff_a = reorganize_read_noise_coeff_to_rggb(coeff_a, color_map)
337        coeff_b = reorganize_read_noise_coeff_to_rggb(coeff_b, color_map)
338
339      logging.info('Sensitivity range: [%d, %d]', sens_min, sens_max)
340      logging.info('Max analog sensitivity: %d', sens_max_analog)
341
342      # Do AE to get a rough idea of where we are.
343      iso_ae, exp_ae, _, _, _ = cam.do_3a(
344          get_results=True, do_awb=False, do_af=False)
345
346      # Underexpose to get more data for low signal levels.
347      auto_e = iso_ae * exp_ae / _BRACKET_FACTOR
348      check_auto_exposure_targets(auto_e, sens_min, sens_max_meas, props)
349
350      # Focus at zero to intentionally blur the scene as much as possible.
351      f_dist = 0.0
352
353      # Start the sensitivities at the minimum.
354      iso = sens_min
355      samples = [[], [], [], []]
356      plots = []
357      measured_models = [[], [], [], []]
358      color_plane_plots = {}
359      isos = []
360      fmt_raw = {'format': 'rawStats',
361                 'gridWidth': _TILE_SIZE,
362                 'gridHeight': _TILE_SIZE}
363
364      while int(round(iso)) <= sens_max_meas:
365        req = capture_request_utils.manual_capture_request(
366            int(round(iso)), min_exposure_ns, f_dist)
367        cap = cam.do_capture(req, fmt_raw)
368
369        # Instead of raising an error when the sensitivity readback != requested
370        # use the readback value for calculations instead
371        iso_cap = cap['metadata']['android.sensor.sensitivity']
372        isos.append(iso_cap)
373
374        logging.info('ISO %d', iso_cap)
375
376        fig, [[plt_r, plt_gr], [plt_gb, plt_b]] = plt.subplots(
377            2, 2, figsize=(11, 11))
378        fig.gca()
379        color_plane_plots[iso_cap] = [plt_r, plt_gr, plt_gb, plt_b]
380        fig.suptitle('ISO %d' % iso_cap, x=0.54, y=0.99)
381        for i, plot in enumerate(color_plane_plots[iso_cap]):
382          plot.set_title(_BAYER_LIST[i])
383          plot.set_xlabel('Mean signal level')
384          plot.set_ylabel('Variance')
385
386        samples_s = [[], [], [], []]
387        for b in range(_BRACKET_MAX):
388          # Get the exposure for this sensitivity and exposure time.
389          exposure = int(math.pow(2, b)*auto_e/iso)
390          logging.info('exp %.3fms', round(exposure*1.0E-6, 3))
391          req = capture_request_utils.manual_capture_request(iso_cap, exposure,
392                                                             f_dist)
393          req['android.control.zoomRatio'] = _ZOOM_RATIO
394          fmt_raw = {'format': 'rawStats',
395                     'gridWidth': _TILE_SIZE,
396                     'gridHeight': _TILE_SIZE}
397          cap = cam.do_capture(req, fmt_raw)
398          if self.debug_mode:
399            img = image_processing_utils.convert_capture_to_rgb_image(
400                cap, props=props)
401            image_processing_utils.write_image(
402                img, f'{name_with_log_path}_{iso_cap}_{exposure}ns.jpg', True)
403
404          mean_img, var_img = image_processing_utils.unpack_rawstats_capture(
405              cap)
406          idxs = image_processing_utils.get_canonical_cfa_order(props)
407          raw_stats_size = mean_img.shape
408          means = [mean_img[_TILE_CROP_N:raw_stats_size[0]-_TILE_CROP_N,
409                            _TILE_CROP_N:raw_stats_size[1]-_TILE_CROP_N, i]
410                   for i in idxs]
411          vars_ = [var_img[_TILE_CROP_N:raw_stats_size[0]-_TILE_CROP_N,
412                           _TILE_CROP_N:raw_stats_size[1]-_TILE_CROP_N, i]
413                   for i in idxs]
414          if self.debug_mode:
415            logging.info('rawStats image size: %s', str(raw_stats_size))
416            logging.info('R planes means image size: %s', str(means[0].shape))
417            logging.info('means min: %.3f, median: %.3f, max: %.3f',
418                         np.min(means), np.median(means), np.max(means))
419            logging.info('vars_ min: %.4f, median: %.4f, max: %.4f',
420                         np.min(vars_), np.median(vars_), np.max(vars_))
421
422          # If remove outliers is True, we will filter the variance data
423          if _REMOVE_OUTLIERS:
424            means_filtered = []
425            vars_filtered = []
426            for pidx in range(len(means)):
427              keep_indices = outlier_removed_indices(vars_[pidx],
428                                                     _OUTLIE_MEDIAN_ABS_DEVS)
429              means_filtered.append(means[pidx][keep_indices])
430              vars_filtered.append(vars_[pidx][keep_indices])
431
432            means = means_filtered
433            vars_ = vars_filtered
434
435          for pidx in range(len(means)):
436            plot = color_plane_plots[iso_cap][pidx]
437
438            # convert_capture_to_planes normalizes the range to [0, 1], but
439            # without subtracting the black level.
440            black_level = image_processing_utils.get_black_level(
441                pidx, props, cap['metadata'])
442            means_p = (means[pidx] - black_level)/(white_level - black_level)
443            vars_p = vars_[pidx]/((white_level - black_level)**2)
444
445            # TODO(dsharlet): It should be possible to account for low
446            # frequency variation by looking at neighboring means, but I
447            # have not been able to make this work.
448
449            means_p = np.asarray(means_p).flatten()
450            vars_p = np.asarray(vars_p).flatten()
451
452            samples_e = []
453            for (mean, var) in zip(means_p, vars_p):
454              # Don't include the tile if it has samples that might be clipped.
455              if mean + 2*math.sqrt(max(var, 0)) < _MAX_SIGNAL_VALUE:
456                samples_e.append([mean, var])
457
458            if samples_e:
459              means_e, vars_e = zip(*samples_e)
460              color_plane_plots[iso_cap][pidx].plot(
461                  means_e, vars_e, _PLOT_COLORS[b%len(_PLOT_COLORS)] + '.',
462                  alpha=0.5, markersize=1)
463              samples_s[pidx].extend(samples_e)
464
465        for (pidx, _) in enumerate(samples_s):
466          [slope, intercept, rvalue, _, _] = scipy.stats.linregress(
467              samples_s[pidx])
468
469          measured_models[pidx].append([iso_cap, slope, intercept])
470          logging.info('%s sensitivity %d: %e*y + %e (R=%f)',
471                       _BAYER_LIST[pidx], iso_cap, slope, intercept, rvalue)
472
473          # Add the samples for this sensitivity to the global samples list.
474          samples[pidx].extend(
475              [(iso_cap, mean, var) for (mean, var) in samples_s[pidx]])
476
477          # Add the linear fit to subplot for this sensitivity.
478          color_plane_plots[iso_cap][pidx].plot(
479              [0, _MAX_SIGNAL_VALUE],
480              [intercept, intercept + slope * _MAX_SIGNAL_VALUE],
481              'rgkb'[pidx] + '--',
482              label='Linear fit')
483
484          xmax = max([max([x for (x, _) in p]) for p in samples_s
485                     ]) * _MAX_SCALE_FUDGE
486          ymax = (intercept + slope * xmax) * _MAX_SCALE_FUDGE
487          color_plane_plots[iso_cap][pidx].set_xlim(xmin=0, xmax=xmax)
488          color_plane_plots[iso_cap][pidx].set_ylim(ymin=0, ymax=ymax)
489          color_plane_plots[iso_cap][pidx].legend()
490          pylab.tight_layout()
491
492        fig.savefig(f'{name_with_log_path}_samples_iso{iso_cap:04d}.png')
493        plots.append([iso_cap, fig])
494
495        # Move to the next sensitivity.
496        iso *= math.pow(2, 1.0/_STEPS_PER_STOP)
497
498    # Do model plots
499    (fig, (plt_s, plt_o)) = plt.subplots(2, 1, figsize=(11, 8.5))
500    plt_s.set_title('Noise model: N(x) = sqrt(Sx + O)')
501    plt_s.set_ylabel('S')
502    plt_o.set_xlabel('ISO')
503    plt_o.set_ylabel('O')
504
505    noise_model = []
506    for (pidx, _) in enumerate(measured_models):
507      # Grab the sensitivities and line parameters from each sensitivity.
508      s_measured = [e[1] for e in measured_models[pidx]]
509      o_measured = [e[2] for e in measured_models[pidx]]
510      sens = np.asarray([e[0] for e in measured_models[pidx]])
511      sens_sq = np.square(sens)
512
513      # Use a global linear optimization to fit the noise model.
514      gains = np.asarray([s[0] for s in samples[pidx]])
515      means = np.asarray([s[1] for s in samples[pidx]])
516      vars_ = np.asarray([s[2] for s in samples[pidx]])
517      gains = gains.flatten()
518      means = means.flatten()
519      vars_ = vars_.flatten()
520
521      # Define digital gain as the gain above the max analog gain
522      # per the Camera2 spec. Also, define a corresponding C
523      # expression snippet to use in the generated model code.
524      digital_gains = np.maximum(gains/sens_max_analog, 1)
525      if not np.all(digital_gains == 1):
526        raise AssertionError(f'Digital gain! gains: {gains}, '
527                             f'Max analog gain: {sens_max_analog}.')
528      digital_gain_cdef = '(sens / %d.0) < 1.0 ? 1.0 : (sens / %d.0)' % (
529          sens_max_analog, sens_max_analog)
530
531      # Noise model function:
532      # f(x) = scale * x + offset
533      # Where:
534      # scale = scale_a*analog_gain*digital_gain + scale_b
535      # offset = (offset_a*analog_gain^2 + offset_b)*digital_gain^2
536      #
537      # Function f will be used to train the scale and offset coefficients,
538      # scale_a, scale_b, offset_a, offset_b (a, b, c, d respectively)
539      # Divide the whole system by gains*means.
540      if _TWO_STAGE_MODEL:
541        # For the two-stage model, we want to use the line fit coefficients
542        # found from capturing read noise data (offset_a and offset_b) to
543        # train the scale coefficients
544        f = lambda x, a, b: (x[1]*a*x[0]+(x[1])*b+coeff_a[pidx]*(x[0]**2)+
545                             coeff_b[pidx])/(x[0])
546      else:
547        f = lambda x, a, b, c, d: (x[1]*a*x[0]+x[1]*b + c*(x[0]**2)+d)/(x[0])
548      result, _ = scipy.optimize.curve_fit(f, (gains, means), vars_/(gains))
549      # result[0:4] = s_gradient, s_offset, o_gradient, o_offset
550      # Note 'S' and 'O' are the API terms for the 2 model params.
551      # The noise_profile.cc uses 'slope' for 'S' and 'intercept' for 'O'.
552      # 'gradient' and 'offset' are used to describe the linear fit
553      # parameters for 'S' and 'O'.
554
555      # If using two-stage model, two of the coefficients calculated above are
556      # constant, so we need to append them to the result ndarray
557      if _TWO_STAGE_MODEL:
558        result = np.append(result, coeff_a[pidx])
559        result = np.append(result, coeff_b[pidx])
560
561      noise_model.append(result[0:4])
562
563      # Plot noise model components with the values predicted by the model.
564      s_model = result[0]*sens + result[1]
565      o_model = result[2]*sens_sq + result[3]*np.square(np.maximum(
566          sens/sens_max_analog, 1))
567
568      plt_s.loglog(sens, s_measured, 'rgkb'[pidx]+'+', base=10,
569                   label='Measured')
570      plt_s.loglog(sens, s_model, 'rgkb'[pidx]+'o', base=10,
571                   label='Model', alpha=0.3)
572      plt_o.loglog(sens, o_measured, 'rgkb'[pidx]+'+', base=10,
573                   label='Measured')
574      plt_o.loglog(sens, o_model, 'rgkb'[pidx]+'o', base=10,
575                   label='Model', alpha=0.3)
576
577    plt_s.set_xticks(isos)
578    plt_s.xaxis.set_minor_locator(NullLocator())  # no minor ticks
579    plt_s.xaxis.set_major_formatter(ScalarFormatter())
580    plt_s.legend()
581
582    plt_o.set_xticks(isos)
583    plt_o.xaxis.set_minor_locator(NullLocator())  # no minor ticks
584    plt_o.xaxis.set_major_formatter(ScalarFormatter())
585    plt_o.legend()
586    fig.savefig(f'{name_with_log_path}.png')
587
588    # Generate individual noise model components
589    noise_model_a, noise_model_b, noise_model_c, noise_model_d = zip(
590        *noise_model)
591
592    # Add models to subplots and re-save
593    for [iso, fig] in plots:  # re-step through figs...
594      dig_gain = max(iso/sens_max_analog, 1)
595      fig.gca()
596      for (pidx, _) in enumerate(measured_models):
597        s = noise_model_a[pidx]*iso + noise_model_b[pidx]
598        o = noise_model_c[pidx]*iso**2 + noise_model_d[pidx]*dig_gain**2
599        color_plane_plots[iso][pidx].plot(
600            [0, _MAX_SIGNAL_VALUE], [o, o+s*_MAX_SIGNAL_VALUE],
601            'rgkb'[pidx]+'-', label='Model', alpha=0.5)
602        color_plane_plots[iso][pidx].legend(loc='upper left')
603      fig.savefig(f'{name_with_log_path}_samples_iso{iso:04d}.png')
604
605    # Validity checks on model: read noise > 0, positive intercept gradient.
606    for i, _ in enumerate(_BAYER_LIST):
607      read_noise = noise_model_c[i] * sens_min * sens_min + noise_model_d[i]
608      if read_noise <= 0:
609        raise AssertionError(f'{_BAYER_LIST[i]} model min ISO noise < 0! '
610                             f'API intercept gradient: {noise_model_c[i]:.4e}, '
611                             f'API intercept offset: {noise_model_d[i]:.4e}, '
612                             f'read_noise: {read_noise:.4e}')
613      if noise_model_c[i] <= 0:
614        raise AssertionError(f'{_BAYER_LIST[i]} model API intercept gradient '
615                             f'is negative: {noise_model_c[i]:.4e}')
616
617    # If 2-Stage model is enabled, save the read noise graph and csv data
618    if _TWO_STAGE_MODEL:
619      # Save the linear plot of the read noise data
620      filename = f'{Path(_NAME_READ_NOISE_FILE).stem}.png'
621      file_path = os.path.join(log_path, filename)
622      capture_read_noise_utils.create_read_noise_plots_from_results(
623          read_noise_results,
624          sens_min,
625          sens_max_meas,
626          color_map,
627          file_path)
628
629      # Save the data as a csv file
630      filename = f'{Path(_NAME_READ_NOISE_FILE).stem}.csv'
631      file_path = os.path.join(log_path, filename)
632      capture_read_noise_utils.create_and_save_csv_from_results(
633          read_noise_results,
634          sens_min,
635          sens_max_meas,
636          color_map,
637          file_path)
638
639    # Generate the noise model file.
640    create_noise_model_code(
641        noise_model_a, noise_model_b, noise_model_c, noise_model_d,
642        sens_min, sens_max, digital_gain_cdef, log_path)
643
644if __name__ == '__main__':
645  test_runner.main()
646