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 its.device
16import its.caps
17import its.objects
18import its.image
19import os.path
20from matplotlib import pylab
21import matplotlib
22import matplotlib.pyplot as plt
23import math
24import textwrap
25import time
26import numpy as np
27import scipy.stats
28import scipy.signal
29
30
31# Convert a 2D array a to a 4D array with dimensions [tile_size,
32# tile_size, row, col] where row, col are tile indices.
33def tile(a, tile_size):
34    tile_rows, tile_cols = a.shape[0]/tile_size, a.shape[1]/tile_size
35    a = a.reshape([tile_rows, tile_size, tile_cols, tile_size])
36    a = a.transpose([1, 3, 0, 2])
37    return a
38
39
40def main():
41    """Capture a set of raw images with increasing gains and measure the noise.
42    """
43    NAME = os.path.basename(__file__).split(".")[0]
44    BAYER_LIST = ['R', 'GR', 'GB', 'B']
45
46    # How many sensitivities per stop to sample.
47    steps_per_stop = 2
48    # How large of tiles to use to compute mean/variance.
49    tile_size = 64
50    # Exposure bracketing range in stops
51    bracket_stops = 4
52    # How high to allow the mean of the tiles to go.
53    max_signal_level = 0.5
54    # Colors used for plotting the data for each exposure.
55    colors = 'rygcbm'
56
57    # Define a first order high pass filter to eliminate low frequency
58    # signal content when computing variance.
59    f = np.array([-1, 1]).astype('float32')
60    # Make it a higher order filter by convolving the first order
61    # filter with itself a few times.
62    f = np.convolve(f, f)
63    f = np.convolve(f, f)
64
65    # Compute the normalization of the filter to preserve noise
66    # power. Let a be the normalization factor we're looking for, and
67    # Let X and X' be the random variables representing the noise
68    # before and after filtering, respectively. First, compute
69    # Var[a*X']:
70    #
71    #   Var[a*X'] = a^2*Var[X*f_0 + X*f_1 + ... + X*f_N-1]
72    #             = a^2*(f_0^2*Var[X] + f_1^2*Var[X] + ... + (f_N-1)^2*Var[X])
73    #             = sum(f_i^2)*a^2*Var[X]
74    #
75    # We want Var[a*X'] to be equal to Var[X]:
76    #
77    #    sum(f_i^2)*a^2*Var[X] = Var[X] -> a = sqrt(1/sum(f_i^2))
78    #
79    # We can just bake this normalization factor into the high pass
80    # filter kernel.
81    f /= math.sqrt(np.dot(f, f))
82
83    bracket_factor = math.pow(2, bracket_stops)
84
85    with its.device.ItsSession() as cam:
86        props = cam.get_camera_properties()
87
88        # Get basic properties we need.
89        sens_min, sens_max = props['android.sensor.info.sensitivityRange']
90        sens_max_analog = props['android.sensor.maxAnalogSensitivity']
91        white_level = props['android.sensor.info.whiteLevel']
92
93        print "Sensitivity range: [%f, %f]" % (sens_min, sens_max)
94        print "Max analog sensitivity: %f" % (sens_max_analog)
95
96        # Do AE to get a rough idea of where we are.
97        s_ae, e_ae, _, _, _  = \
98            cam.do_3a(get_results=True, do_awb=False, do_af=False)
99        # Underexpose to get more data for low signal levels.
100        auto_e = s_ae*e_ae/bracket_factor
101        # Focus at zero to intentionally blur the scene as much as possible.
102        f_dist = 0.0
103
104        # If the auto-exposure result is too bright for the highest
105        # sensitivity or too dark for the lowest sensitivity, report
106        # an error.
107        min_exposure_ns, max_exposure_ns = \
108            props['android.sensor.info.exposureTimeRange']
109        if auto_e < min_exposure_ns*sens_max:
110            raise its.error.Error("Scene is too bright to properly expose \
111                                  at the highest sensitivity")
112        if auto_e*bracket_factor > max_exposure_ns*sens_min:
113            raise its.error.Error("Scene is too dark to properly expose \
114                                  at the lowest sensitivity")
115
116        # Start the sensitivities at the minimum.
117        s = sens_min
118
119        samples = [[], [], [], []]
120        plots = []
121        measured_models = [[], [], [], []]
122        while s <= sens_max + 1:
123            print "ISO %d" % round(s)
124            fig = plt.figure()
125            plt_s = fig.gca()
126            plt_s.set_title("ISO %d" % round(s))
127            plt_s.set_xlabel("Mean signal level")
128            plt_s.set_ylabel("Variance")
129
130            samples_s = [[], [], [], []]
131            for b in range(0, bracket_stops + 1):
132                # Get the exposure for this sensitivity and exposure time.
133                e = int(math.pow(2, b)*auto_e/float(s))
134                req = its.objects.manual_capture_request(round(s), e, f_dist)
135                cap = cam.do_capture(req, cam.CAP_RAW)
136                planes = its.image.convert_capture_to_planes(cap, props)
137
138                for (pidx, p) in enumerate(planes):
139                    p = p.squeeze()
140
141                    # Crop the plane to be a multiple of the tile size.
142                    p = p[0:p.shape[0] - p.shape[0]%tile_size,
143                          0:p.shape[1] - p.shape[1]%tile_size]
144
145                    # convert_capture_to_planes normalizes the range
146                    # to [0, 1], but without subtracting the black
147                    # level.
148                    black_level = its.image.get_black_level(
149                        pidx, props, cap["metadata"])
150                    p *= white_level
151                    p = (p - black_level)/(white_level - black_level)
152
153                    # Use our high pass filter to filter this plane.
154                    hp = scipy.signal.sepfir2d(p, f, f).astype('float32')
155
156                    means_tiled = \
157                        np.mean(tile(p, tile_size), axis=(0, 1)).flatten()
158                    vars_tiled = \
159                        np.var(tile(hp, tile_size), axis=(0, 1)).flatten()
160
161                    samples_e = []
162                    for (mean, var) in zip(means_tiled, vars_tiled):
163                        # Don't include the tile if it has samples that might
164                        # be clipped.
165                        if mean + 2*math.sqrt(var) < max_signal_level:
166                            samples_e.append([mean, var])
167
168                    if len(samples_e) > 0:
169                        means_e, vars_e = zip(*samples_e)
170                        plt_s.plot(means_e, vars_e, colors[b%len(colors)] + ',')
171
172                        samples_s[pidx].extend(samples_e)
173
174            for (pidx, p) in enumerate(samples_s):
175                [S, O, R, p, stderr] = scipy.stats.linregress(samples_s[pidx])
176                measured_models[pidx].append([round(s), S, O])
177                print "Sensitivity %d: %e*y + %e (R=%f)" % (round(s), S, O, R)
178
179                # Add the samples for this sensitivity to the global samples list.
180                samples[pidx].extend([(round(s), mean, var) for (mean, var) in samples_s[pidx]])
181
182                # Add the linear fit to the plot for this sensitivity.
183                plt_s.plot([0, max_signal_level], [O, O + S*max_signal_level], 'rgkb'[pidx]+'--',
184                           label="Linear fit")
185
186            xmax = max([max([x for (x, _) in p]) for p in samples_s])*1.25
187            ymax = max([max([y for (_, y) in p]) for p in samples_s])*1.25
188            plt_s.set_xlim(xmin=0, xmax=xmax)
189            plt_s.set_ylim(ymin=0, ymax=ymax)
190
191            fig.savefig("%s_samples_iso%04d.png" % (NAME, round(s)))
192            plots.append([round(s), fig])
193
194            # Move to the next sensitivity.
195            s *= math.pow(2, 1.0/steps_per_stop)
196
197        (fig, (plt_S, plt_O)) = plt.subplots(2, 1)
198        plt_S.set_title("Noise model")
199        plt_S.set_ylabel("S")
200        plt_S.legend(loc=2)
201        plt_O.set_xlabel("ISO")
202        plt_O.set_ylabel("O")
203
204        A = []
205        B = []
206        C = []
207        D = []
208        for (pidx, p) in enumerate(measured_models):
209            # Grab the sensitivities and line parameters from each sensitivity.
210            S_measured = [e[1] for e in measured_models[pidx]]
211            O_measured = [e[2] for e in measured_models[pidx]]
212            sens = np.asarray([e[0] for e in measured_models[pidx]])
213            sens_sq = np.square(sens)
214
215            # Use a global linear optimization to fit the noise model.
216            gains = np.asarray([s[0] for s in samples[pidx]])
217            means = np.asarray([s[1] for s in samples[pidx]])
218            vars_ = np.asarray([s[2] for s in samples[pidx]])
219
220            # Define digital gain as the gain above the max analog gain
221            # per the Camera2 spec. Also, define a corresponding C
222            # expression snippet to use in the generated model code.
223            digital_gains = np.maximum(gains/sens_max_analog, 1)
224            digital_gain_cdef = "(sens / %d.0) < 1.0 ? 1.0 : (sens / %d.0)" % \
225                (sens_max_analog, sens_max_analog)
226
227            # Find the noise model parameters via least squares fit.
228            ad = gains*means
229            bd = means
230            cd = gains*gains
231            dd = digital_gains*digital_gains
232            a = np.asarray([ad, bd, cd, dd]).T
233            b = vars_
234
235            # To avoid overfitting to high ISOs (high variances), divide the system
236            # by the gains.
237            a /= (np.tile(gains, (a.shape[1], 1)).T)
238            b /= gains
239
240            [A_p, B_p, C_p, D_p], _, _, _ = np.linalg.lstsq(a, b)
241            A.append(A_p)
242            B.append(B_p)
243            C.append(C_p)
244            D.append(D_p)
245
246            # Plot the noise model components with the values predicted by the
247            # noise model.
248            S_model = A_p*sens + B_p
249            O_model = \
250                C_p*sens_sq + D_p*np.square(np.maximum(sens/sens_max_analog, 1))
251
252            plt_S.loglog(sens, S_measured, 'rgkb'[pidx]+'+', basex=10, basey=10,
253                         label="Measured")
254            plt_S.loglog(sens, S_model, 'rgkb'[pidx]+'x', basex=10, basey=10, label="Model")
255
256            plt_O.loglog(sens, O_measured, 'rgkb'[pidx]+'+', basex=10, basey=10,
257                         label="Measured")
258            plt_O.loglog(sens, O_model, 'rgkb'[pidx]+'x', basex=10, basey=10, label="Model")
259
260        fig.savefig("%s.png" % (NAME))
261
262        for [s, fig] in plots:
263            plt_s = fig.gca()
264
265            dg = max(s/sens_max_analog, 1)
266            for (pidx, p) in enumerate(measured_models):
267                S = A[pidx]*s + B[pidx]
268                O = C[pidx]*s*s + D[pidx]*dg*dg
269                plt_s.plot([0, max_signal_level], [O, O + S*max_signal_level], 'rgkb'[pidx]+'-',
270                           label="Model")
271
272            plt_s.legend(loc=2)
273            plt.figure(fig.number)
274
275            # Re-save the plot with the global model.
276            fig.savefig("%s_samples_iso%04d.png" % (NAME, round(s)))
277
278          # Generate the noise model implementation.
279        A_array = ",".join([str(i) for i in A])
280        B_array = ",".join([str(i) for i in B])
281        C_array = ",".join([str(i) for i in C])
282        D_array = ",".join([str(i) for i in D])
283        noise_model_code = textwrap.dedent("""\
284            /* Generated test code to dump a table of data for external validation
285             * of the noise model parameters.
286             */
287            #include <stdio.h>
288            #include <assert.h>
289            double compute_noise_model_entry_S(int plane, int sens);
290            double compute_noise_model_entry_O(int plane, int sens);
291            int main(void) {
292                for (int plane = 0; plane < %d; plane++) {
293                    for (int sens = %d; sens <= %d; sens += 100) {
294                        double o = compute_noise_model_entry_O(plane, sens);
295                        double s = compute_noise_model_entry_S(plane, sens);
296                        printf("%%d,%%d,%%lf,%%lf\\n", plane, sens, o, s);
297                    }
298                }
299                return 0;
300            }
301
302            /* Generated functions to map a given sensitivity to the O and S noise
303             * model parameters in the DNG noise model. The planes are in
304             * R, Gr, Gb, B order.
305             */
306            double compute_noise_model_entry_S(int plane, int sens) {
307                static double noise_model_A[] = { %s };
308                static double noise_model_B[] = { %s };
309                double A = noise_model_A[plane];
310                double B = noise_model_B[plane];
311                double s = A * sens + B;
312                return s < 0.0 ? 0.0 : s;
313            }
314
315            double compute_noise_model_entry_O(int plane, int sens) {
316                static double noise_model_C[] = { %s };
317                static double noise_model_D[] = { %s };
318                double digital_gain = %s;
319                double C = noise_model_C[plane];
320                double D = noise_model_D[plane];
321                double o = C * sens * sens + D * digital_gain * digital_gain;
322                return o < 0.0 ? 0.0 : o;
323            }
324            """ % (len(A), sens_min, sens_max, A_array, B_array, C_array, D_array, digital_gain_cdef))
325        print noise_model_code
326        text_file = open("noise_model.c", "w")
327        text_file.write("%s" % noise_model_code)
328        text_file.close()
329
330if __name__ == '__main__':
331    main()
332