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