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