1# Copyright 2022 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"""Utility functions to enable capture read noise analysis.""" 15 16import csv 17import logging 18import math 19import os 20import pickle 21 22import matplotlib.pyplot as plt 23from matplotlib.ticker import NullLocator 24from matplotlib.ticker import ScalarFormatter 25import numpy as np 26 27import camera_properties_utils 28import capture_request_utils 29 30 31_BAYER_COLOR_PLANE = ('red', 'green_r', 'blue', 'green_b') 32_LINEAR_FIT_NUM_SAMPLES = 100 # Number of samples to plot for the linear fit 33_PLOT_AXIS_TICKS = 5 # Number of ticks to display on the plot axis 34 35 36def create_and_save_csv_from_results(rn_data, iso_low, iso_high, cmap, file): 37 """Creates a .csv file for the read noise results. 38 39 Args: 40 rn_data: Read noise data object from capture_read_noise_for_iso_range 41 iso_low: int; minimum iso value to include 42 iso_high: int; maximum iso value to include 43 cmap: str; string containing each color symbol 44 file: str; path to csv where this will be created 45 """ 46 with open(file, 'w+') as f: 47 writer = csv.writer(f) 48 49 results = list(filter( 50 lambda x: x[0]['iso'] >= iso_low and x[0]['iso'] <= iso_high, rn_data 51 )) 52 53 color_channels = range(len(cmap)) 54 55 # Create headers for csv file 56 headers = ['iso', 'iso^2'] 57 headers.extend([f'mean_{cmap[i]}' for i in color_channels]) 58 headers.extend([f'var_{cmap[i]}' for i in color_channels]) 59 headers.extend([f'norm_var_{cmap[i]}' for i in color_channels]) 60 61 writer.writerow(headers) 62 63 # Create data rows 64 for data_row in results: 65 row = [data_row[0]['iso']] 66 row.append(data_row[0]['iso']**2) 67 row.extend([data_row[i]['mean'] for i in color_channels]) 68 row.extend([data_row[i]['var'] for i in color_channels]) 69 row.extend([data_row[i]['norm_var'] for i in color_channels]) 70 71 writer.writerow(row) 72 73 writer.writerow([]) # divider line 74 75 # Create row containing the offset coefficients calculated by np.polyfit 76 coeff_headers = ['', 'offset_coefficient_a', 'offset_coefficient_b'] 77 writer.writerow(coeff_headers) 78 79 coeff_a, coeff_b = get_read_noise_coefficients(results) 80 for i in range(len(cmap)): 81 writer.writerow([cmap[i], coeff_a[i], coeff_b[i]]) 82 83 84def create_read_noise_plots_from_results(rn_data, iso_low, iso_high, cmap, 85 file): 86 """Plot the read noise data for the given ISO range. 87 88 Args: 89 rn_data: Read noise data object from capture_read_noise_for_iso_range 90 iso_low: int; minimum iso value to include 91 iso_high: int; maximum iso value to include 92 cmap: str; string containing the Bayer format 93 file: str; file path for the plot image 94 """ 95 # Get a list of color names and plot color arrangements for the given cmap. 96 # This will be used for chart labels and color schemes 97 bayer_color_list = [] 98 plot_colors = '' 99 if cmap.lower() == 'grbg': 100 bayer_color_list = ['GR', 'R', 'B', 'GB'] 101 plot_colors = 'grby' 102 elif cmap.lower() == 'rggb': 103 bayer_color_list = ['R', 'GR', 'GB', 'B'] 104 plot_colors = 'rgyb' 105 elif cmap.lower() == 'bggr': 106 bayer_color_list = ['B', 'GB', 'GR', 'R'] 107 plot_colors = 'bygr' 108 elif cmap.lower() == 'gbrg': 109 bayer_color_list = ['GB', 'B', 'R', 'GR'] 110 plot_colors = 'ybrg' 111 else: 112 raise AssertionError('cmap parameter does not match any known Bayer format') 113 114 # Create the figure for plotting the read noise to ISO^2 curve 115 fig = plt.figure(figsize=(11, 11)) 116 fig.suptitle('Read Noise to ISO^2', x=0.54, y=0.99) 117 118 iso_range = fig.add_subplot(111) 119 120 iso_range.set_xlabel('ISO^2') 121 iso_range.set_ylabel('Read Noise') 122 123 # Get the ISO values for the current range 124 current_range = list(filter( 125 lambda x: x[0]['iso'] >= iso_low and x[0]['iso'] <= iso_high, rn_data 126 )) 127 128 # Get X-axis values (ISO^2) for current_range 129 iso_sq = [data[0]['iso']**2 for data in current_range] 130 131 # Get X-axis values for the calculated linear fit for the read noise 132 iso_sq_values = np.linspace(iso_low**2, iso_high**2, _LINEAR_FIT_NUM_SAMPLES) 133 134 # Get the line fit coeff for plotting the linear fit of read noise to iso^2 135 coeff_a, coeff_b = get_read_noise_coefficients(current_range) 136 137 # Plot the read noise to iso^2 data 138 for pidx in range(len(bayer_color_list)): 139 norm_vars = [data[pidx]['norm_var'] for data in current_range] 140 141 # Plot the measured read noise to ISO^2 values 142 iso_range.plot(iso_sq, norm_vars, plot_colors[pidx]+'o', 143 label=f'{bayer_color_list[pidx]}', alpha=0.3) 144 145 # Plot the line fit calculated from the read noise values 146 iso_range.plot(iso_sq_values, coeff_a[pidx]*iso_sq_values + coeff_b[pidx], 147 color=plot_colors[pidx]) 148 149 # Create a numpy array containing all normalized variance values for the 150 # current range, this will be used for labelling the X-axis 151 y_values = np.array( 152 [[color['norm_var'] for color in x] for x in current_range]) 153 154 x_ticks = np.linspace(iso_low**2, iso_high**2, _PLOT_AXIS_TICKS) 155 y_ticks = np.linspace(np.min(y_values), np.max(y_values), _PLOT_AXIS_TICKS) 156 157 iso_range.set_xticks(x_ticks) 158 iso_range.xaxis.set_minor_locator(NullLocator()) 159 iso_range.xaxis.set_major_formatter(ScalarFormatter()) 160 161 iso_range.set_yticks(y_ticks) 162 iso_range.yaxis.set_minor_locator(NullLocator()) 163 iso_range.yaxis.set_major_formatter(ScalarFormatter()) 164 165 iso_range.legend() 166 167 fig.savefig(file) 168 169 170def _generate_image_data_bayer(img, iso, white_level, cmap): 171 """Generates read noise data for a given image. 172 173 Each element in the list corresponds to each color channel, and each dict 174 contains information relevant to the read noise calculation. 175 176 Args: 177 img: np.array; image for the given iso 178 iso: float; iso value which the 179 white_level: int; white level value for the sensor 180 cmap: str; color map of the sensor 181 Returns: 182 list(dict) list containing information for each color channel 183 """ 184 result = [] 185 186 color_channel_img = np.empty((len(_BAYER_COLOR_PLANE), 187 int(img.shape[0]/2), 188 int(img.shape[1]/2))) 189 190 # Create a dict of read noise values for each color channel in the image 191 for i, color_plane in enumerate(_BAYER_COLOR_PLANE): 192 color_channel_img[i] = _subsample(img, color_plane, cmap) 193 var = np.var(color_channel_img[i]) 194 mean = np.mean(color_channel_img[i]) 195 norm_var = var / ((white_level - mean)**2) 196 result.append({ 197 'iso': iso, 198 'mean': mean, 199 'var': var, 200 'norm_var': norm_var 201 }) 202 203 return result 204 205 206def _subsample(img, color_plane, cmap): 207 """Subsample image array based on color_plane. 208 209 Args: 210 img: 2-D numpy array of image 211 color_plane: string; color to extract 212 cmap: list; color map of the sensor 213 Returns: 214 img_subsample: 2-D numpy subarray of image with only color plane 215 """ 216 subsample_img_2x = lambda img, x, h, v: img[int(x / 2):v:2, x % 2:h:2] 217 size_h = img.shape[1] 218 size_v = img.shape[0] 219 if color_plane == 'red': 220 cmap_index = cmap.index('R') 221 elif color_plane == 'blue': 222 cmap_index = cmap.index('B') 223 elif color_plane == 'green_r': 224 color_plane_map_index = { 225 'GRBG': 0, 226 'RGGB': 1, 227 'BGGR': 2, 228 'GBRG': 3 229 } 230 cmap_index = color_plane_map_index[cmap] 231 elif color_plane == 'green_b': 232 color_plane_map_index = { 233 'GBRG': 0, 234 'BGGR': 1, 235 'RGGB': 2, 236 'GRBG': 3 237 } 238 cmap_index = color_plane_map_index[cmap] 239 else: 240 logging.error('Wrong color_plane entered!') 241 return None 242 243 return subsample_img_2x(img, cmap_index, size_h, size_v) 244 245 246def get_read_noise_coefficients(rn_data, iso_low=0, iso_high=1000000): 247 """Calculate the read noise coefficients from the read noise data. 248 249 Args: 250 rn_data: Read noise data object from capture_read_noise_for_iso_range 251 iso_low: int; minimum iso value to include 252 iso_high: int; maximum iso value to include 253 Returns: 254 (list, list) Offset coefficients for the linear fit to read noise data 255 """ 256 # Filter the values by the given ISO range 257 iso_range = list(filter( 258 lambda x: x[0]['iso'] >= iso_low and x[0]['iso'] <= iso_high, rn_data 259 )) 260 261 read_noise_coefficients_a = [] 262 read_noise_coefficients_b = [] 263 264 # Get ISO^2 values used for X-axis in polyfit 265 iso_sq = [data[0]['iso']**2 for data in iso_range] 266 267 # Find the linear equation coefficients for each color channel 268 for i in range(len(iso_range[0])): 269 norm_var = [data[i]['norm_var'] for data in iso_range] 270 271 coeffs = np.polyfit(iso_sq, norm_var, 1) 272 273 read_noise_coefficients_a.append(coeffs[0]) 274 read_noise_coefficients_b.append(coeffs[1]) 275 276 return read_noise_coefficients_a, read_noise_coefficients_b 277 278 279def capture_read_noise_for_iso_range(cam, low_iso, high_iso, steps, cmap, 280 dest_file): 281 """Captures read noise data at the lowest advertised exposure value. 282 283 Args: 284 cam: ItsSession; camera for the current ItsSession 285 low_iso: int; lowest iso value in range 286 high_iso: int; highest iso value in range 287 steps: int; steps to take per stop 288 cmap: str; color map of the sensor 289 dest_file: str; path where the results should be saved 290 Returns: 291 list(list(dict)) Read noise results for each frame 292 """ 293 props = cam.get_camera_properties() 294 props = cam.override_with_hidden_physical_camera_props(props) 295 camera_properties_utils.skip_unless( 296 camera_properties_utils.raw16(props) and 297 camera_properties_utils.manual_sensor(props) and 298 camera_properties_utils.read_3a(props) and 299 camera_properties_utils.per_frame_control(props)) 300 301 min_exposure_ns, _ = props['android.sensor.info.exposureTimeRange'] 302 min_fd = props['android.lens.info.minimumFocusDistance'] 303 white_level = props['android.sensor.info.whiteLevel'] 304 305 iso = low_iso 306 307 results = [] 308 309 # This operation can last a very long time, if it happens to fail halfway 310 # through, this section of code will allow us to pick up where we left off 311 if os.path.exists(dest_file): 312 # If there already exists a results file, retrieve them 313 with open(dest_file, 'rb') as f: 314 results = pickle.load(f) 315 # Set the starting iso to the last iso of results 316 iso = results[-1][0]['iso'] 317 iso *= math.pow(2, 1.0/steps) 318 319 while int(round(iso)) <= high_iso: 320 iso_int = int(iso) 321 req = capture_request_utils.manual_capture_request(iso_int, min_exposure_ns) 322 req['android.lens.focusDistance'] = min_fd 323 fmt = {'format': 'raw'} 324 cap = cam.do_capture(req, fmt) 325 w = cap['width'] 326 h = cap['height'] 327 328 img = np.ndarray(shape=(h*w,), dtype='<u2', buffer=cap['data'][0:w*h*2]) 329 img = img.astype(dtype=np.uint16).reshape(h, w) 330 331 # Add values to results, organized as a dictionary 332 results.append(_generate_image_data_bayer(img, iso, white_level, cmap)) 333 334 logging.info('iso: %.2f, mean: %.2f, var: %.2f, min: %d, max: %d', iso, 335 np.mean(img), np.var(img), np.min(img), np.max(img)) 336 337 with open(dest_file, 'wb+') as f: 338 pickle.dump(results, f) 339 340 iso *= math.pow(2, 1.0/steps) 341 342 logging.info('Results pickled into file %s', dest_file) 343 344 return results 345