1# Copyright 2013 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 matplotlib
16matplotlib.use('Agg')
17
18import its.error
19import sys
20from PIL import Image
21import numpy
22import math
23import unittest
24import cStringIO
25import copy
26import random
27
28DEFAULT_YUV_TO_RGB_CCM = numpy.matrix([
29                                [1.000,  0.000,  1.402],
30                                [1.000, -0.344, -0.714],
31                                [1.000,  1.772,  0.000]])
32
33DEFAULT_YUV_OFFSETS = numpy.array([0, 128, 128])
34
35DEFAULT_GAMMA_LUT = numpy.array(
36        [math.floor(65535 * math.pow(i/65535.0, 1/2.2) + 0.5)
37         for i in xrange(65536)])
38
39DEFAULT_INVGAMMA_LUT = numpy.array(
40        [math.floor(65535 * math.pow(i/65535.0, 2.2) + 0.5)
41         for i in xrange(65536)])
42
43MAX_LUT_SIZE = 65536
44
45NUM_TRYS = 2
46NUM_FRAMES = 4
47
48
49def convert_capture_to_rgb_image(cap,
50                                 ccm_yuv_to_rgb=DEFAULT_YUV_TO_RGB_CCM,
51                                 yuv_off=DEFAULT_YUV_OFFSETS,
52                                 props=None):
53    """Convert a captured image object to a RGB image.
54
55    Args:
56        cap: A capture object as returned by its.device.do_capture.
57        ccm_yuv_to_rgb: (Optional) the 3x3 CCM to convert from YUV to RGB.
58        yuv_off: (Optional) offsets to subtract from each of Y,U,V values.
59        props: (Optional) camera properties object (of static values);
60            required for processing raw images.
61
62    Returns:
63        RGB float-3 image array, with pixel values in [0.0, 1.0].
64    """
65    w = cap["width"]
66    h = cap["height"]
67    if cap["format"] == "raw10":
68        assert(props is not None)
69        cap = unpack_raw10_capture(cap, props)
70    if cap["format"] == "raw12":
71        assert(props is not None)
72        cap = unpack_raw12_capture(cap, props)
73    if cap["format"] == "yuv":
74        y = cap["data"][0:w*h]
75        u = cap["data"][w*h:w*h*5/4]
76        v = cap["data"][w*h*5/4:w*h*6/4]
77        return convert_yuv420_planar_to_rgb_image(y, u, v, w, h)
78    elif cap["format"] == "jpeg":
79        return decompress_jpeg_to_rgb_image(cap["data"])
80    elif cap["format"] == "raw" or cap["format"] == "rawStats":
81        assert(props is not None)
82        r,gr,gb,b = convert_capture_to_planes(cap, props)
83        return convert_raw_to_rgb_image(r,gr,gb,b, props, cap["metadata"])
84    elif cap["format"] == "y8":
85        y = cap["data"][0:w*h]
86        return convert_y8_to_rgb_image(y, w, h)
87    else:
88        raise its.error.Error('Invalid format %s' % (cap["format"]))
89
90
91def unpack_rawstats_capture(cap):
92    """Unpack a rawStats capture to the mean and variance images.
93
94    Args:
95        cap: A capture object as returned by its.device.do_capture.
96
97    Returns:
98        Tuple (mean_image var_image) of float-4 images, with non-normalized
99        pixel values computed from the RAW16 images on the device
100    """
101    assert(cap["format"] == "rawStats")
102    w = cap["width"]
103    h = cap["height"]
104    img = numpy.ndarray(shape=(2*h*w*4,), dtype='<f', buffer=cap["data"])
105    analysis_image = img.reshape(2,h,w,4)
106    mean_image = analysis_image[0,:,:,:].reshape(h,w,4)
107    var_image = analysis_image[1,:,:,:].reshape(h,w,4)
108    return mean_image, var_image
109
110
111def unpack_raw10_capture(cap, props):
112    """Unpack a raw-10 capture to a raw-16 capture.
113
114    Args:
115        cap: A raw-10 capture object.
116        props: Camera properties object.
117
118    Returns:
119        New capture object with raw-16 data.
120    """
121    # Data is packed as 4x10b pixels in 5 bytes, with the first 4 bytes holding
122    # the MSPs of the pixels, and the 5th byte holding 4x2b LSBs.
123    w,h = cap["width"], cap["height"]
124    if w % 4 != 0:
125        raise its.error.Error('Invalid raw-10 buffer width')
126    cap = copy.deepcopy(cap)
127    cap["data"] = unpack_raw10_image(cap["data"].reshape(h,w*5/4))
128    cap["format"] = "raw"
129    return cap
130
131
132def unpack_raw10_image(img):
133    """Unpack a raw-10 image to a raw-16 image.
134
135    Output image will have the 10 LSBs filled in each 16b word, and the 6 MSBs
136    will be set to zero.
137
138    Args:
139        img: A raw-10 image, as a uint8 numpy array.
140
141    Returns:
142        Image as a uint16 numpy array, with all row padding stripped.
143    """
144    if img.shape[1] % 5 != 0:
145        raise its.error.Error('Invalid raw-10 buffer width')
146    w = img.shape[1]*4/5
147    h = img.shape[0]
148    # Cut out the 4x8b MSBs and shift to bits [9:2] in 16b words.
149    msbs = numpy.delete(img, numpy.s_[4::5], 1)
150    msbs = msbs.astype(numpy.uint16)
151    msbs = numpy.left_shift(msbs, 2)
152    msbs = msbs.reshape(h,w)
153    # Cut out the 4x2b LSBs and put each in bits [1:0] of their own 8b words.
154    lsbs = img[::, 4::5].reshape(h,w/4)
155    lsbs = numpy.right_shift(
156            numpy.packbits(numpy.unpackbits(lsbs).reshape(h,w/4,4,2),3), 6)
157    # Pair the LSB bits group to 0th pixel instead of 3rd pixel
158    lsbs = lsbs.reshape(h,w/4,4)[:,:,::-1]
159    lsbs = lsbs.reshape(h,w)
160    # Fuse the MSBs and LSBs back together
161    img16 = numpy.bitwise_or(msbs, lsbs).reshape(h,w)
162    return img16
163
164
165def unpack_raw12_capture(cap, props):
166    """Unpack a raw-12 capture to a raw-16 capture.
167
168    Args:
169        cap: A raw-12 capture object.
170        props: Camera properties object.
171
172    Returns:
173        New capture object with raw-16 data.
174    """
175    # Data is packed as 4x10b pixels in 5 bytes, with the first 4 bytes holding
176    # the MSBs of the pixels, and the 5th byte holding 4x2b LSBs.
177    w,h = cap["width"], cap["height"]
178    if w % 2 != 0:
179        raise its.error.Error('Invalid raw-12 buffer width')
180    cap = copy.deepcopy(cap)
181    cap["data"] = unpack_raw12_image(cap["data"].reshape(h,w*3/2))
182    cap["format"] = "raw"
183    return cap
184
185
186def unpack_raw12_image(img):
187    """Unpack a raw-12 image to a raw-16 image.
188
189    Output image will have the 12 LSBs filled in each 16b word, and the 4 MSBs
190    will be set to zero.
191
192    Args:
193        img: A raw-12 image, as a uint8 numpy array.
194
195    Returns:
196        Image as a uint16 numpy array, with all row padding stripped.
197    """
198    if img.shape[1] % 3 != 0:
199        raise its.error.Error('Invalid raw-12 buffer width')
200    w = img.shape[1]*2/3
201    h = img.shape[0]
202    # Cut out the 2x8b MSBs and shift to bits [11:4] in 16b words.
203    msbs = numpy.delete(img, numpy.s_[2::3], 1)
204    msbs = msbs.astype(numpy.uint16)
205    msbs = numpy.left_shift(msbs, 4)
206    msbs = msbs.reshape(h,w)
207    # Cut out the 2x4b LSBs and put each in bits [3:0] of their own 8b words.
208    lsbs = img[::, 2::3].reshape(h,w/2)
209    lsbs = numpy.right_shift(
210            numpy.packbits(numpy.unpackbits(lsbs).reshape(h,w/2,2,4),3), 4)
211    # Pair the LSB bits group to pixel 0 instead of pixel 1
212    lsbs = lsbs.reshape(h,w/2,2)[:,:,::-1]
213    lsbs = lsbs.reshape(h,w)
214    # Fuse the MSBs and LSBs back together
215    img16 = numpy.bitwise_or(msbs, lsbs).reshape(h,w)
216    return img16
217
218
219def convert_capture_to_planes(cap, props=None):
220    """Convert a captured image object to separate image planes.
221
222    Decompose an image into multiple images, corresponding to different planes.
223
224    For YUV420 captures ("yuv"):
225        Returns Y,U,V planes, where the Y plane is full-res and the U,V planes
226        are each 1/2 x 1/2 of the full res.
227
228    For Bayer captures ("raw", "raw10", "raw12", or "rawStats"):
229        Returns planes in the order R,Gr,Gb,B, regardless of the Bayer pattern
230        layout. For full-res raw images ("raw", "raw10", "raw12"), each plane
231        is 1/2 x 1/2 of the full res. For "rawStats" images, the mean image
232        is returned.
233
234    For JPEG captures ("jpeg"):
235        Returns R,G,B full-res planes.
236
237    Args:
238        cap: A capture object as returned by its.device.do_capture.
239        props: (Optional) camera properties object (of static values);
240            required for processing raw images.
241
242    Returns:
243        A tuple of float numpy arrays (one per plane), consisting of pixel
244            values in the range [0.0, 1.0].
245    """
246    w = cap["width"]
247    h = cap["height"]
248    if cap["format"] == "raw10":
249        assert(props is not None)
250        cap = unpack_raw10_capture(cap, props)
251    if cap["format"] == "raw12":
252        assert(props is not None)
253        cap = unpack_raw12_capture(cap, props)
254    if cap["format"] == "yuv":
255        y = cap["data"][0:w*h]
256        u = cap["data"][w*h:w*h*5/4]
257        v = cap["data"][w*h*5/4:w*h*6/4]
258        return ((y.astype(numpy.float32) / 255.0).reshape(h, w, 1),
259                (u.astype(numpy.float32) / 255.0).reshape(h/2, w/2, 1),
260                (v.astype(numpy.float32) / 255.0).reshape(h/2, w/2, 1))
261    elif cap["format"] == "jpeg":
262        rgb = decompress_jpeg_to_rgb_image(cap["data"]).reshape(w*h*3)
263        return (rgb[::3].reshape(h,w,1),
264                rgb[1::3].reshape(h,w,1),
265                rgb[2::3].reshape(h,w,1))
266    elif cap["format"] == "raw":
267        assert(props is not None)
268        white_level = float(props['android.sensor.info.whiteLevel'])
269        img = numpy.ndarray(shape=(h*w,), dtype='<u2',
270                            buffer=cap["data"][0:w*h*2])
271        img = img.astype(numpy.float32).reshape(h,w) / white_level
272        # Crop the raw image to the active array region.
273        if props.has_key("android.sensor.info.preCorrectionActiveArraySize") \
274                and props["android.sensor.info.preCorrectionActiveArraySize"] is not None \
275                and props.has_key("android.sensor.info.pixelArraySize") \
276                and props["android.sensor.info.pixelArraySize"] is not None:
277            # Note that the Rect class is defined such that the left,top values
278            # are "inside" while the right,bottom values are "outside"; that is,
279            # it's inclusive of the top,left sides only. So, the width is
280            # computed as right-left, rather than right-left+1, etc.
281            wfull = props["android.sensor.info.pixelArraySize"]["width"]
282            hfull = props["android.sensor.info.pixelArraySize"]["height"]
283            xcrop = props["android.sensor.info.preCorrectionActiveArraySize"]["left"]
284            ycrop = props["android.sensor.info.preCorrectionActiveArraySize"]["top"]
285            wcrop = props["android.sensor.info.preCorrectionActiveArraySize"]["right"]-xcrop
286            hcrop = props["android.sensor.info.preCorrectionActiveArraySize"]["bottom"]-ycrop
287            assert(wfull >= wcrop >= 0)
288            assert(hfull >= hcrop >= 0)
289            assert(wfull - wcrop >= xcrop >= 0)
290            assert(hfull - hcrop >= ycrop >= 0)
291            if w == wfull and h == hfull:
292                # Crop needed; extract the center region.
293                img = img[ycrop:ycrop+hcrop,xcrop:xcrop+wcrop]
294                w = wcrop
295                h = hcrop
296            elif w == wcrop and h == hcrop:
297                # No crop needed; image is already cropped to the active array.
298                None
299            else:
300                raise its.error.Error('Invalid image size metadata')
301        # Separate the image planes.
302        imgs = [img[::2].reshape(w*h/2)[::2].reshape(h/2,w/2,1),
303                img[::2].reshape(w*h/2)[1::2].reshape(h/2,w/2,1),
304                img[1::2].reshape(w*h/2)[::2].reshape(h/2,w/2,1),
305                img[1::2].reshape(w*h/2)[1::2].reshape(h/2,w/2,1)]
306        idxs = get_canonical_cfa_order(props)
307        return [imgs[i] for i in idxs]
308    elif cap["format"] == "rawStats":
309        assert(props is not None)
310        white_level = float(props['android.sensor.info.whiteLevel'])
311        mean_image, var_image = its.image.unpack_rawstats_capture(cap)
312        idxs = get_canonical_cfa_order(props)
313        return [mean_image[:,:,i] / white_level for i in idxs]
314    else:
315        raise its.error.Error('Invalid format %s' % (cap["format"]))
316
317
318def get_canonical_cfa_order(props):
319    """Returns a mapping from the Bayer 2x2 top-left grid in the CFA to
320    the standard order R,Gr,Gb,B.
321
322    Args:
323        props: Camera properties object.
324
325    Returns:
326        List of 4 integers, corresponding to the positions in the 2x2 top-
327            left Bayer grid of R,Gr,Gb,B, where the 2x2 grid is labeled as
328            0,1,2,3 in row major order.
329    """
330    # Note that raw streams aren't croppable, so the cropRegion doesn't need
331    # to be considered when determining the top-left pixel color.
332    cfa_pat = props['android.sensor.info.colorFilterArrangement']
333    if cfa_pat == 0:
334        # RGGB
335        return [0,1,2,3]
336    elif cfa_pat == 1:
337        # GRBG
338        return [1,0,3,2]
339    elif cfa_pat == 2:
340        # GBRG
341        return [2,3,0,1]
342    elif cfa_pat == 3:
343        # BGGR
344        return [3,2,1,0]
345    else:
346        raise its.error.Error("Not supported")
347
348
349def get_gains_in_canonical_order(props, gains):
350    """Reorders the gains tuple to the canonical R,Gr,Gb,B order.
351
352    Args:
353        props: Camera properties object.
354        gains: List of 4 values, in R,G_even,G_odd,B order.
355
356    Returns:
357        List of gains values, in R,Gr,Gb,B order.
358    """
359    cfa_pat = props['android.sensor.info.colorFilterArrangement']
360    if cfa_pat in [0,1]:
361        # RGGB or GRBG, so G_even is Gr
362        return gains
363    elif cfa_pat in [2,3]:
364        # GBRG or BGGR, so G_even is Gb
365        return [gains[0], gains[2], gains[1], gains[3]]
366    else:
367        raise its.error.Error("Not supported")
368
369
370def convert_raw_to_rgb_image(r_plane, gr_plane, gb_plane, b_plane,
371                             props, cap_res):
372    """Convert a Bayer raw-16 image to an RGB image.
373
374    Includes some extremely rudimentary demosaicking and color processing
375    operations; the output of this function shouldn't be used for any image
376    quality analysis.
377
378    Args:
379        r_plane,gr_plane,gb_plane,b_plane: Numpy arrays for each color plane
380            in the Bayer image, with pixels in the [0.0, 1.0] range.
381        props: Camera properties object.
382        cap_res: Capture result (metadata) object.
383
384    Returns:
385        RGB float-3 image array, with pixel values in [0.0, 1.0]
386    """
387    # Values required for the RAW to RGB conversion.
388    assert(props is not None)
389    white_level = float(props['android.sensor.info.whiteLevel'])
390    black_levels = props['android.sensor.blackLevelPattern']
391    gains = cap_res['android.colorCorrection.gains']
392    ccm = cap_res['android.colorCorrection.transform']
393
394    # Reorder black levels and gains to R,Gr,Gb,B, to match the order
395    # of the planes.
396    black_levels = [get_black_level(i,props,cap_res) for i in range(4)]
397    gains = get_gains_in_canonical_order(props, gains)
398
399    # Convert CCM from rational to float, as numpy arrays.
400    ccm = numpy.array(its.objects.rational_to_float(ccm)).reshape(3,3)
401
402    # Need to scale the image back to the full [0,1] range after subtracting
403    # the black level from each pixel.
404    scale = white_level / (white_level - max(black_levels))
405
406    # Three-channel black levels, normalized to [0,1] by white_level.
407    black_levels = numpy.array([b/white_level for b in [
408            black_levels[i] for i in [0,1,3]]])
409
410    # Three-channel gains.
411    gains = numpy.array([gains[i] for i in [0,1,3]])
412
413    h,w = r_plane.shape[:2]
414    img = numpy.dstack([r_plane,(gr_plane+gb_plane)/2.0,b_plane])
415    img = (((img.reshape(h,w,3) - black_levels) * scale) * gains).clip(0.0,1.0)
416    img = numpy.dot(img.reshape(w*h,3), ccm.T).reshape(h,w,3).clip(0.0,1.0)
417    return img
418
419
420def get_black_level(chan, props, cap_res=None):
421    """Return the black level to use for a given capture.
422
423    Uses a dynamic value from the capture result if available, else falls back
424    to the static global value in the camera characteristics.
425
426    Args:
427        chan: The channel index, in canonical order (R, Gr, Gb, B).
428        props: The camera properties object.
429        cap_res: A capture result object.
430
431    Returns:
432        The black level value for the specified channel.
433    """
434    if (cap_res is not None and cap_res.has_key('android.sensor.dynamicBlackLevel') and
435            cap_res['android.sensor.dynamicBlackLevel'] is not None):
436        black_levels = cap_res['android.sensor.dynamicBlackLevel']
437    else:
438        black_levels = props['android.sensor.blackLevelPattern']
439    idxs = its.image.get_canonical_cfa_order(props)
440    ordered_black_levels = [black_levels[i] for i in idxs]
441    return ordered_black_levels[chan]
442
443
444def convert_yuv420_planar_to_rgb_image(y_plane, u_plane, v_plane,
445                                       w, h,
446                                       ccm_yuv_to_rgb=DEFAULT_YUV_TO_RGB_CCM,
447                                       yuv_off=DEFAULT_YUV_OFFSETS):
448    """Convert a YUV420 8-bit planar image to an RGB image.
449
450    Args:
451        y_plane: The packed 8-bit Y plane.
452        u_plane: The packed 8-bit U plane.
453        v_plane: The packed 8-bit V plane.
454        w: The width of the image.
455        h: The height of the image.
456        ccm_yuv_to_rgb: (Optional) the 3x3 CCM to convert from YUV to RGB.
457        yuv_off: (Optional) offsets to subtract from each of Y,U,V values.
458
459    Returns:
460        RGB float-3 image array, with pixel values in [0.0, 1.0].
461    """
462    y = numpy.subtract(y_plane, yuv_off[0])
463    u = numpy.subtract(u_plane, yuv_off[1]).view(numpy.int8)
464    v = numpy.subtract(v_plane, yuv_off[2]).view(numpy.int8)
465    u = u.reshape(h/2, w/2).repeat(2, axis=1).repeat(2, axis=0)
466    v = v.reshape(h/2, w/2).repeat(2, axis=1).repeat(2, axis=0)
467    yuv = numpy.dstack([y, u.reshape(w*h), v.reshape(w*h)])
468    flt = numpy.empty([h, w, 3], dtype=numpy.float32)
469    flt.reshape(w*h*3)[:] = yuv.reshape(h*w*3)[:]
470    flt = numpy.dot(flt.reshape(w*h,3), ccm_yuv_to_rgb.T).clip(0, 255)
471    rgb = numpy.empty([h, w, 3], dtype=numpy.uint8)
472    rgb.reshape(w*h*3)[:] = flt.reshape(w*h*3)[:]
473    return rgb.astype(numpy.float32) / 255.0
474
475def convert_y8_to_rgb_image(y_plane, w, h):
476    """Convert a Y 8-bit image to an RGB image.
477
478    Args:
479        y_plane: The packed 8-bit Y plane.
480        w: The width of the image.
481        h: The height of the image.
482
483    Returns:
484        RGB float-3 image array, with pixel values in [0.0, 1.0].
485    """
486    y3 = numpy.dstack([y_plane, y_plane, y_plane])
487    rgb = numpy.empty([h, w, 3], dtype=numpy.uint8)
488    rgb.reshape(w*h*3)[:] = y3.reshape(w*h*3)[:]
489    return rgb.astype(numpy.float32) / 255.0
490
491def load_rgb_image(fname):
492    """Load a standard image file (JPG, PNG, etc.).
493
494    Args:
495        fname: The path of the file to load.
496
497    Returns:
498        RGB float-3 image array, with pixel values in [0.0, 1.0].
499    """
500    img = Image.open(fname)
501    w = img.size[0]
502    h = img.size[1]
503    a = numpy.array(img)
504    if len(a.shape) == 3 and a.shape[2] == 3:
505        # RGB
506        return a.reshape(h,w,3) / 255.0
507    elif len(a.shape) == 2 or len(a.shape) == 3 and a.shape[2] == 1:
508        # Greyscale; convert to RGB
509        return a.reshape(h*w).repeat(3).reshape(h,w,3) / 255.0
510    else:
511        raise its.error.Error('Unsupported image type')
512
513
514def load_yuv420_to_rgb_image(yuv_fname,
515                             w, h,
516                             layout="planar",
517                             ccm_yuv_to_rgb=DEFAULT_YUV_TO_RGB_CCM,
518                             yuv_off=DEFAULT_YUV_OFFSETS):
519    """Load a YUV420 image file, and return as an RGB image.
520
521    Supported layouts include "planar" and "nv21". The "yuv" formatted captures
522    returned from the device via do_capture are in the "planar" layout; other
523    layouts may only be needed for loading files from other sources.
524
525    Args:
526        yuv_fname: The path of the YUV420 file.
527        w: The width of the image.
528        h: The height of the image.
529        layout: (Optional) the layout of the YUV data (as a string).
530        ccm_yuv_to_rgb: (Optional) the 3x3 CCM to convert from YUV to RGB.
531        yuv_off: (Optional) offsets to subtract from each of Y,U,V values.
532
533    Returns:
534        RGB float-3 image array, with pixel values in [0.0, 1.0].
535    """
536    with open(yuv_fname, "rb") as f:
537        if layout == "planar":
538            # Plane of Y, plane of V, plane of U.
539            y = numpy.fromfile(f, numpy.uint8, w*h, "")
540            v = numpy.fromfile(f, numpy.uint8, w*h/4, "")
541            u = numpy.fromfile(f, numpy.uint8, w*h/4, "")
542        elif layout == "nv21":
543            # Plane of Y, plane of interleaved VUVUVU...
544            y = numpy.fromfile(f, numpy.uint8, w*h, "")
545            vu = numpy.fromfile(f, numpy.uint8, w*h/2, "")
546            v = vu[0::2]
547            u = vu[1::2]
548        else:
549            raise its.error.Error('Unsupported image layout')
550        return convert_yuv420_planar_to_rgb_image(
551                y,u,v,w,h,ccm_yuv_to_rgb,yuv_off)
552
553
554def load_yuv420_planar_to_yuv_planes(yuv_fname, w, h):
555    """Load a YUV420 planar image file, and return Y, U, and V plane images.
556
557    Args:
558        yuv_fname: The path of the YUV420 file.
559        w: The width of the image.
560        h: The height of the image.
561
562    Returns:
563        Separate Y, U, and V images as float-1 Numpy arrays, pixels in [0,1].
564        Note that pixel (0,0,0) is not black, since U,V pixels are centered at
565        0.5, and also that the Y and U,V plane images returned are different
566        sizes (due to chroma subsampling in the YUV420 format).
567    """
568    with open(yuv_fname, "rb") as f:
569        y = numpy.fromfile(f, numpy.uint8, w*h, "")
570        v = numpy.fromfile(f, numpy.uint8, w*h/4, "")
571        u = numpy.fromfile(f, numpy.uint8, w*h/4, "")
572        return ((y.astype(numpy.float32) / 255.0).reshape(h, w, 1),
573                (u.astype(numpy.float32) / 255.0).reshape(h/2, w/2, 1),
574                (v.astype(numpy.float32) / 255.0).reshape(h/2, w/2, 1))
575
576
577def decompress_jpeg_to_rgb_image(jpeg_buffer):
578    """Decompress a JPEG-compressed image, returning as an RGB image.
579
580    Args:
581        jpeg_buffer: The JPEG stream.
582
583    Returns:
584        A numpy array for the RGB image, with pixels in [0,1].
585    """
586    img = Image.open(cStringIO.StringIO(jpeg_buffer))
587    w = img.size[0]
588    h = img.size[1]
589    return numpy.array(img).reshape(h,w,3) / 255.0
590
591
592def apply_lut_to_image(img, lut):
593    """Applies a LUT to every pixel in a float image array.
594
595    Internally converts to a 16b integer image, since the LUT can work with up
596    to 16b->16b mappings (i.e. values in the range [0,65535]). The lut can also
597    have fewer than 65536 entries, however it must be sized as a power of 2
598    (and for smaller luts, the scale must match the bitdepth).
599
600    For a 16b lut of 65536 entries, the operation performed is:
601
602        lut[r * 65535] / 65535 -> r'
603        lut[g * 65535] / 65535 -> g'
604        lut[b * 65535] / 65535 -> b'
605
606    For a 10b lut of 1024 entries, the operation becomes:
607
608        lut[r * 1023] / 1023 -> r'
609        lut[g * 1023] / 1023 -> g'
610        lut[b * 1023] / 1023 -> b'
611
612    Args:
613        img: Numpy float image array, with pixel values in [0,1].
614        lut: Numpy table encoding a LUT, mapping 16b integer values.
615
616    Returns:
617        Float image array after applying LUT to each pixel.
618    """
619    n = len(lut)
620    if n <= 0 or n > MAX_LUT_SIZE or (n & (n - 1)) != 0:
621        raise its.error.Error('Invalid arg LUT size: %d' % (n))
622    m = float(n-1)
623    return (lut[(img * m).astype(numpy.uint16)] / m).astype(numpy.float32)
624
625
626def apply_matrix_to_image(img, mat):
627    """Multiplies a 3x3 matrix with each float-3 image pixel.
628
629    Each pixel is considered a column vector, and is left-multiplied by
630    the given matrix:
631
632        [     ]   r    r'
633        [ mat ] * g -> g'
634        [     ]   b    b'
635
636    Args:
637        img: Numpy float image array, with pixel values in [0,1].
638        mat: Numpy 3x3 matrix.
639
640    Returns:
641        The numpy float-3 image array resulting from the matrix mult.
642    """
643    h = img.shape[0]
644    w = img.shape[1]
645    img2 = numpy.empty([h, w, 3], dtype=numpy.float32)
646    img2.reshape(w*h*3)[:] = (numpy.dot(img.reshape(h*w, 3), mat.T)
647                             ).reshape(w*h*3)[:]
648    return img2
649
650
651def get_image_patch(img, xnorm, ynorm, wnorm, hnorm):
652    """Get a patch (tile) of an image.
653
654    Args:
655        img: Numpy float image array, with pixel values in [0,1].
656        xnorm,ynorm,wnorm,hnorm: Normalized (in [0,1]) coords for the tile.
657
658    Returns:
659        Float image array of the patch.
660    """
661    hfull = img.shape[0]
662    wfull = img.shape[1]
663    xtile = int(math.ceil(xnorm * wfull))
664    ytile = int(math.ceil(ynorm * hfull))
665    wtile = int(math.floor(wnorm * wfull))
666    htile = int(math.floor(hnorm * hfull))
667    if len(img.shape)==2:
668        return img[ytile:ytile+htile,xtile:xtile+wtile].copy()
669    else:
670        return img[ytile:ytile+htile,xtile:xtile+wtile,:].copy()
671
672
673def compute_image_means(img):
674    """Calculate the mean of each color channel in the image.
675
676    Args:
677        img: Numpy float image array, with pixel values in [0,1].
678
679    Returns:
680        A list of mean values, one per color channel in the image.
681    """
682    means = []
683    chans = img.shape[2]
684    for i in xrange(chans):
685        means.append(numpy.mean(img[:,:,i], dtype=numpy.float64))
686    return means
687
688
689def compute_image_variances(img):
690    """Calculate the variance of each color channel in the image.
691
692    Args:
693        img: Numpy float image array, with pixel values in [0,1].
694
695    Returns:
696        A list of mean values, one per color channel in the image.
697    """
698    variances = []
699    chans = img.shape[2]
700    for i in xrange(chans):
701        variances.append(numpy.var(img[:,:,i], dtype=numpy.float64))
702    return variances
703
704
705def compute_image_snrs(img):
706    """Calculate the SNR (db) of each color channel in the image.
707
708    Args:
709        img: Numpy float image array, with pixel values in [0,1].
710
711    Returns:
712        A list of SNR value, one per color channel in the image.
713    """
714    means = compute_image_means(img)
715    variances = compute_image_variances(img)
716    std_devs = [math.sqrt(v) for v in variances]
717    snr = [20 * math.log10(m/s) for m,s in zip(means, std_devs)]
718    return snr
719
720
721def compute_image_max_gradients(img):
722    """Calculate the maximum gradient of each color channel in the image.
723
724    Args:
725        img: Numpy float image array, with pixel values in [0,1].
726
727    Returns:
728        A list of gradient max values, one per color channel in the image.
729    """
730    grads = []
731    chans = img.shape[2]
732    for i in xrange(chans):
733        grads.append(numpy.amax(numpy.gradient(img[:, :, i])))
734    return grads
735
736
737def write_image(img, fname, apply_gamma=False):
738    """Save a float-3 numpy array image to a file.
739
740    Supported formats: PNG, JPEG, and others; see PIL docs for more.
741
742    Image can be 3-channel, which is interpreted as RGB, or can be 1-channel,
743    which is greyscale.
744
745    Can optionally specify that the image should be gamma-encoded prior to
746    writing it out; this should be done if the image contains linear pixel
747    values, to make the image look "normal".
748
749    Args:
750        img: Numpy image array data.
751        fname: Path of file to save to; the extension specifies the format.
752        apply_gamma: (Optional) apply gamma to the image prior to writing it.
753    """
754    if apply_gamma:
755        img = apply_lut_to_image(img, DEFAULT_GAMMA_LUT)
756    (h, w, chans) = img.shape
757    if chans == 3:
758        Image.fromarray((img * 255.0).astype(numpy.uint8), "RGB").save(fname)
759    elif chans == 1:
760        img3 = (img * 255.0).astype(numpy.uint8).repeat(3).reshape(h,w,3)
761        Image.fromarray(img3, "RGB").save(fname)
762    else:
763        raise its.error.Error('Unsupported image type')
764
765
766def downscale_image(img, f):
767    """Shrink an image by a given integer factor.
768
769    This function computes output pixel values by averaging over rectangular
770    regions of the input image; it doesn't skip or sample pixels, and all input
771    image pixels are evenly weighted.
772
773    If the downscaling factor doesn't cleanly divide the width and/or height,
774    then the remaining pixels on the right or bottom edge are discarded prior
775    to the downscaling.
776
777    Args:
778        img: The input image as an ndarray.
779        f: The downscaling factor, which should be an integer.
780
781    Returns:
782        The new (downscaled) image, as an ndarray.
783    """
784    h,w,chans = img.shape
785    f = int(f)
786    assert(f >= 1)
787    h = (h/f)*f
788    w = (w/f)*f
789    img = img[0:h:,0:w:,::]
790    chs = []
791    for i in xrange(chans):
792        ch = img.reshape(h*w*chans)[i::chans].reshape(h,w)
793        ch = ch.reshape(h,w/f,f).mean(2).reshape(h,w/f)
794        ch = ch.T.reshape(w/f,h/f,f).mean(2).T.reshape(h/f,w/f)
795        chs.append(ch.reshape(h*w/(f*f)))
796    img = numpy.vstack(chs).T.reshape(h/f,w/f,chans)
797    return img
798
799
800def compute_image_sharpness(img):
801    """Calculate the sharpness of input image.
802
803    Args:
804        img: Numpy float RGB/luma image array, with pixel values in [0,1].
805
806    Returns:
807        A sharpness estimation value based on the average of gradient magnitude.
808        Larger value means the image is sharper.
809    """
810    chans = img.shape[2]
811    assert(chans == 1 or chans == 3)
812    if (chans == 1):
813        luma = img[:, :, 0]
814    elif (chans == 3):
815        luma = 0.299 * img[:,:,0] + 0.587 * img[:,:,1] + 0.114 * img[:,:,2]
816
817    [gy, gx] = numpy.gradient(luma)
818    return numpy.average(numpy.sqrt(gy*gy + gx*gx))
819
820
821def normalize_img(img):
822    """Normalize the image values to between 0 and 1.
823
824    Args:
825        img: 2-D numpy array of image values
826    Returns:
827        Normalized image
828    """
829    return (img - numpy.amin(img))/(numpy.amax(img) - numpy.amin(img))
830
831
832def chart_located_per_argv():
833    """Determine if chart already located outside of test.
834
835    If chart info provided, return location and size. If not, return None.
836
837    Args:
838        None
839    Returns:
840        chart_loc:  float converted xnorm,ynorm,wnorm,hnorm,scale from argv text.
841                    argv is of form 'chart_loc=0.45,0.45,0.1,0.1,1.0'
842    """
843    for s in sys.argv[1:]:
844        if s[:10] == "chart_loc=" and len(s) > 10:
845            chart_loc = s[10:].split(",")
846            return map(float, chart_loc)
847    return None, None, None, None, None
848
849
850def rotate_img_per_argv(img):
851    """Rotate an image 180 degrees if "rotate" is in argv
852
853    Args:
854        img: 2-D numpy array of image values
855    Returns:
856        Rotated image
857    """
858    img_out = img
859    if "rotate180" in sys.argv:
860        img_out = numpy.fliplr(numpy.flipud(img_out))
861    return img_out
862
863
864def stationary_lens_cap(cam, req, fmt):
865    """Take up to NUM_TRYS caps and save the 1st one with lens stationary.
866
867    Args:
868        cam:    open device session
869        req:    capture request
870        fmt:    format for capture
871
872    Returns:
873        capture
874    """
875    trys = 0
876    done = False
877    reqs = [req] * NUM_FRAMES
878    while not done:
879        print 'Waiting for lens to move to correct location...'
880        cap = cam.do_capture(reqs, fmt)
881        done = (cap[NUM_FRAMES-1]['metadata']['android.lens.state'] == 0)
882        print ' status: ', done
883        trys += 1
884        if trys == NUM_TRYS:
885            raise its.error.Error('Cannot settle lens after %d trys!' % trys)
886    return cap[NUM_FRAMES-1]
887
888
889class __UnitTest(unittest.TestCase):
890    """Run a suite of unit tests on this module.
891    """
892
893    # TODO: Add more unit tests.
894
895    def test_apply_matrix_to_image(self):
896        """Unit test for apply_matrix_to_image.
897
898        Test by using a canned set of values on a 1x1 pixel image.
899
900            [ 1 2 3 ]   [ 0.1 ]   [ 1.4 ]
901            [ 4 5 6 ] * [ 0.2 ] = [ 3.2 ]
902            [ 7 8 9 ]   [ 0.3 ]   [ 5.0 ]
903               mat         x         y
904        """
905        mat = numpy.array([[1,2,3], [4,5,6], [7,8,9]])
906        x = numpy.array([0.1,0.2,0.3]).reshape(1,1,3)
907        y = apply_matrix_to_image(x, mat).reshape(3).tolist()
908        y_ref = [1.4,3.2,5.0]
909        passed = all([math.fabs(y[i] - y_ref[i]) < 0.001 for i in xrange(3)])
910        self.assertTrue(passed)
911
912    def test_apply_lut_to_image(self):
913        """Unit test for apply_lut_to_image.
914
915        Test by using a canned set of values on a 1x1 pixel image. The LUT will
916        simply double the value of the index:
917
918            lut[x] = 2*x
919        """
920        lut = numpy.array([2*i for i in xrange(65536)])
921        x = numpy.array([0.1,0.2,0.3]).reshape(1,1,3)
922        y = apply_lut_to_image(x, lut).reshape(3).tolist()
923        y_ref = [0.2,0.4,0.6]
924        passed = all([math.fabs(y[i] - y_ref[i]) < 0.001 for i in xrange(3)])
925        self.assertTrue(passed)
926
927    def test_unpack_raw10_image(self):
928        """Unit test for unpack_raw10_image.
929
930        RAW10 bit packing format
931                bit 7   bit 6   bit 5   bit 4   bit 3   bit 2   bit 1   bit 0
932        Byte 0: P0[9]   P0[8]   P0[7]   P0[6]   P0[5]   P0[4]   P0[3]   P0[2]
933        Byte 1: P1[9]   P1[8]   P1[7]   P1[6]   P1[5]   P1[4]   P1[3]   P1[2]
934        Byte 2: P2[9]   P2[8]   P2[7]   P2[6]   P2[5]   P2[4]   P2[3]   P2[2]
935        Byte 3: P3[9]   P3[8]   P3[7]   P3[6]   P3[5]   P3[4]   P3[3]   P3[2]
936        Byte 4: P3[1]   P3[0]   P2[1]   P2[0]   P1[1]   P1[0]   P0[1]   P0[0]
937        """
938        # test by using a random 4x4 10-bit image
939        H = 4
940        W = 4
941        check_list = random.sample(range(0, 1024), H*W)
942        img_check = numpy.array(check_list).reshape(H, W)
943        # pack bits
944        for row_start in range(0, len(check_list), W):
945            msbs = []
946            lsbs = ""
947            for pixel in range(W):
948                val = numpy.binary_repr(check_list[row_start+pixel], 10)
949                msbs.append(int(val[:8], base=2))
950                lsbs = val[8:] + lsbs
951            packed = msbs
952            packed.append(int(lsbs, base=2))
953            chunk_raw10 = numpy.array(packed, dtype="uint8").reshape(1, 5)
954            if row_start == 0:
955                img_raw10 = chunk_raw10
956            else:
957                img_raw10 = numpy.vstack((img_raw10, chunk_raw10))
958        # unpack and check against original
959        self.assertTrue(numpy.array_equal(unpack_raw10_image(img_raw10),
960                                          img_check))
961
962if __name__ == "__main__":
963    unittest.main()
964