Home | History | Annotate | Download | only in its
      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 
     15 import matplotlib
     16 matplotlib.use('Agg')
     17 
     18 import its.error
     19 import pylab
     20 import sys
     21 import Image
     22 import numpy
     23 import math
     24 import unittest
     25 import cStringIO
     26 import scipy.stats
     27 import copy
     28 
     29 DEFAULT_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 
     34 DEFAULT_YUV_OFFSETS = numpy.array([0, 128, 128])
     35 
     36 DEFAULT_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 
     40 DEFAULT_INVGAMMA_LUT = numpy.array(
     41         [math.floor(65535 * math.pow(i/65535.0, 2.2) + 0.5)
     42          for i in xrange(65536)])
     43 
     44 MAX_LUT_SIZE = 65536
     45 
     46 def 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"] == "yuv":
     68         y = cap["data"][0:w*h]
     69         u = cap["data"][w*h:w*h*5/4]
     70         v = cap["data"][w*h*5/4:w*h*6/4]
     71         return convert_yuv420_to_rgb_image(y, u, v, w, h)
     72     elif cap["format"] == "jpeg":
     73         return decompress_jpeg_to_rgb_image(cap["data"])
     74     elif cap["format"] == "raw":
     75         assert(props is not None)
     76         r,gr,gb,b = convert_capture_to_planes(cap, props)
     77         return convert_raw_to_rgb_image(r,gr,gb,b, props, cap["metadata"])
     78     else:
     79         raise its.error.Error('Invalid format %s' % (cap["format"]))
     80 
     81 def unpack_raw10_capture(cap, props):
     82     """Unpack a raw-10 capture to a raw-16 capture.
     83 
     84     Args:
     85         cap: A raw-10 capture object.
     86         props: Camera propertis object.
     87 
     88     Returns:
     89         New capture object with raw-16 data.
     90     """
     91     # Data is packed as 4x10b pixels in 5 bytes, with the first 4 bytes holding
     92     # the MSPs of the pixels, and the 5th byte holding 4x2b LSBs.
     93     w,h = cap["width"], cap["height"]
     94     if w % 4 != 0:
     95         raise its.error.Error('Invalid raw-10 buffer width')
     96     cap = copy.deepcopy(cap)
     97     cap["data"] = unpack_raw10_image(cap["data"].reshape(h,w*5/4))
     98     cap["format"] = "raw"
     99     return cap
    100 
    101 def unpack_raw10_image(img):
    102     """Unpack a raw-10 image to a raw-16 image.
    103 
    104     Output image will have the 10 LSBs filled in each 16b word, and the 6 MSBs
    105     will be set to zero.
    106 
    107     Args:
    108         img: A raw-10 image, as a uint8 numpy array.
    109 
    110     Returns:
    111         Image as a uint16 numpy array, with all row padding stripped.
    112     """
    113     if img.shape[1] % 5 != 0:
    114         raise its.error.Error('Invalid raw-10 buffer width')
    115     w = img.shape[1]*4/5
    116     h = img.shape[0]
    117     # Cut out the 4x8b MSBs and shift to bits [10:2] in 16b words.
    118     msbs = numpy.delete(img, numpy.s_[4::5], 1)
    119     msbs = msbs.astype(numpy.uint16)
    120     msbs = numpy.left_shift(msbs, 2)
    121     msbs = msbs.reshape(h,w)
    122     # Cut out the 4x2b LSBs and put each in bits [2:0] of their own 8b words.
    123     lsbs = img[::, 4::5].reshape(h,w/4)
    124     lsbs = numpy.right_shift(
    125             numpy.packbits(numpy.unpackbits(lsbs).reshape(h,w/4,4,2),3), 6)
    126     lsbs = lsbs.reshape(h,w)
    127     # Fuse the MSBs and LSBs back together
    128     img16 = numpy.bitwise_or(msbs, lsbs).reshape(h,w)
    129     return img16
    130 
    131 def convert_capture_to_planes(cap, props=None):
    132     """Convert a captured image object to separate image planes.
    133 
    134     Decompose an image into multiple images, corresponding to different planes.
    135 
    136     For YUV420 captures ("yuv"):
    137         Returns Y,U,V planes, where the Y plane is full-res and the U,V planes
    138         are each 1/2 x 1/2 of the full res.
    139 
    140     For Bayer captures ("raw" or "raw10"):
    141         Returns planes in the order R,Gr,Gb,B, regardless of the Bayer pattern
    142         layout. Each plane is 1/2 x 1/2 of the full res.
    143 
    144     For JPEG captures ("jpeg"):
    145         Returns R,G,B full-res planes.
    146 
    147     Args:
    148         cap: A capture object as returned by its.device.do_capture.
    149         props: (Optional) camera properties object (of static values);
    150             required for processing raw images.
    151 
    152     Returns:
    153         A tuple of float numpy arrays (one per plane), consisting of pixel
    154             values in the range [0.0, 1.0].
    155     """
    156     w = cap["width"]
    157     h = cap["height"]
    158     if cap["format"] == "raw10":
    159         assert(props is not None)
    160         cap = unpack_raw10_capture(cap, props)
    161     if cap["format"] == "yuv":
    162         y = cap["data"][0:w*h]
    163         u = cap["data"][w*h:w*h*5/4]
    164         v = cap["data"][w*h*5/4:w*h*6/4]
    165         return ((y.astype(numpy.float32) / 255.0).reshape(h, w, 1),
    166                 (u.astype(numpy.float32) / 255.0).reshape(h/2, w/2, 1),
    167                 (v.astype(numpy.float32) / 255.0).reshape(h/2, w/2, 1))
    168     elif cap["format"] == "jpeg":
    169         rgb = decompress_jpeg_to_rgb_image(cap["data"]).reshape(w*h*3)
    170         return (rgb[::3].reshape(h,w,1),
    171                 rgb[1::3].reshape(h,w,1),
    172                 rgb[2::3].reshape(h,w,1))
    173     elif cap["format"] == "raw":
    174         assert(props is not None)
    175         white_level = float(props['android.sensor.info.whiteLevel'])
    176         img = numpy.ndarray(shape=(h*w,), dtype='<u2',
    177                             buffer=cap["data"][0:w*h*2])
    178         img = img.astype(numpy.float32).reshape(h,w) / white_level
    179         imgs = [img[::2].reshape(w*h/2)[::2].reshape(h/2,w/2,1),
    180                 img[::2].reshape(w*h/2)[1::2].reshape(h/2,w/2,1),
    181                 img[1::2].reshape(w*h/2)[::2].reshape(h/2,w/2,1),
    182                 img[1::2].reshape(w*h/2)[1::2].reshape(h/2,w/2,1)]
    183         idxs = get_canonical_cfa_order(props)
    184         return [imgs[i] for i in idxs]
    185     else:
    186         raise its.error.Error('Invalid format %s' % (cap["format"]))
    187 
    188 def get_canonical_cfa_order(props):
    189     """Returns a mapping from the Bayer 2x2 top-left grid in the CFA to
    190     the standard order R,Gr,Gb,B.
    191 
    192     Args:
    193         props: Camera properties object.
    194 
    195     Returns:
    196         List of 4 integers, corresponding to the positions in the 2x2 top-
    197             left Bayer grid of R,Gr,Gb,B, where the 2x2 grid is labeled as
    198             0,1,2,3 in row major order.
    199     """
    200     # Note that raw streams aren't croppable, so the cropRegion doesn't need
    201     # to be considered when determining the top-left pixel color.
    202     cfa_pat = props['android.sensor.info.colorFilterArrangement']
    203     if cfa_pat == 0:
    204         # RGGB
    205         return [0,1,2,3]
    206     elif cfa_pat == 1:
    207         # GRBG
    208         return [1,0,3,2]
    209     elif cfa_pat == 2:
    210         # GBRG
    211         return [2,3,0,1]
    212     elif cfa_pat == 3:
    213         # BGGR
    214         return [3,2,1,0]
    215     else:
    216         raise its.error.Error("Not supported")
    217 
    218 def get_gains_in_canonical_order(props, gains):
    219     """Reorders the gains tuple to the canonical R,Gr,Gb,B order.
    220 
    221     Args:
    222         props: Camera properties object.
    223         gains: List of 4 values, in R,G_even,G_odd,B order.
    224 
    225     Returns:
    226         List of gains values, in R,Gr,Gb,B order.
    227     """
    228     cfa_pat = props['android.sensor.info.colorFilterArrangement']
    229     if cfa_pat in [0,1]:
    230         # RGGB or GRBG, so G_even is Gr
    231         return gains
    232     elif cfa_pat in [2,3]:
    233         # GBRG or BGGR, so G_even is Gb
    234         return [gains[0], gains[2], gains[1], gains[3]]
    235     else:
    236         raise its.error.Error("Not supported")
    237 
    238 def convert_raw_to_rgb_image(r_plane, gr_plane, gb_plane, b_plane,
    239                              props, cap_res):
    240     """Convert a Bayer raw-16 image to an RGB image.
    241 
    242     Includes some extremely rudimentary demosaicking and color processing
    243     operations; the output of this function shouldn't be used for any image
    244     quality analysis.
    245 
    246     Args:
    247         r_plane,gr_plane,gb_plane,b_plane: Numpy arrays for each color plane
    248             in the Bayer image, with pixels in the [0.0, 1.0] range.
    249         props: Camera properties object.
    250         cap_res: Capture result (metadata) object.
    251 
    252     Returns:
    253         RGB float-3 image array, with pixel values in [0.0, 1.0]
    254     """
    255     # Values required for the RAW to RGB conversion.
    256     assert(props is not None)
    257     white_level = float(props['android.sensor.info.whiteLevel'])
    258     black_levels = props['android.sensor.blackLevelPattern']
    259     gains = cap_res['android.colorCorrection.gains']
    260     ccm = cap_res['android.colorCorrection.transform']
    261 
    262     # Reorder black levels and gains to R,Gr,Gb,B, to match the order
    263     # of the planes.
    264     idxs = get_canonical_cfa_order(props)
    265     black_levels = [black_levels[i] for i in idxs]
    266     gains = get_gains_in_canonical_order(props, gains)
    267 
    268     # Convert CCM from rational to float, as numpy arrays.
    269     ccm = numpy.array(its.objects.rational_to_float(ccm)).reshape(3,3)
    270 
    271     # Need to scale the image back to the full [0,1] range after subtracting
    272     # the black level from each pixel.
    273     scale = white_level / (white_level - max(black_levels))
    274 
    275     # Three-channel black levels, normalized to [0,1] by white_level.
    276     black_levels = numpy.array([b/white_level for b in [
    277             black_levels[i] for i in [0,1,3]]])
    278 
    279     # Three-channel gains.
    280     gains = numpy.array([gains[i] for i in [0,1,3]])
    281 
    282     h,w = r_plane.shape[:2]
    283     img = numpy.dstack([r_plane,(gr_plane+gb_plane)/2.0,b_plane])
    284     img = (((img.reshape(h,w,3) - black_levels) * scale) * gains).clip(0.0,1.0)
    285     img = numpy.dot(img.reshape(w*h,3), ccm.T).reshape(h,w,3).clip(0.0,1.0)
    286     return img
    287 
    288 def convert_yuv420_to_rgb_image(y_plane, u_plane, v_plane,
    289                                 w, h,
    290                                 ccm_yuv_to_rgb=DEFAULT_YUV_TO_RGB_CCM,
    291                                 yuv_off=DEFAULT_YUV_OFFSETS):
    292     """Convert a YUV420 8-bit planar image to an RGB image.
    293 
    294     Args:
    295         y_plane: The packed 8-bit Y plane.
    296         u_plane: The packed 8-bit U plane.
    297         v_plane: The packed 8-bit V plane.
    298         w: The width of the image.
    299         h: The height of the image.
    300         ccm_yuv_to_rgb: (Optional) the 3x3 CCM to convert from YUV to RGB.
    301         yuv_off: (Optional) offsets to subtract from each of Y,U,V values.
    302 
    303     Returns:
    304         RGB float-3 image array, with pixel values in [0.0, 1.0].
    305     """
    306     y = numpy.subtract(y_plane, yuv_off[0])
    307     u = numpy.subtract(u_plane, yuv_off[1]).view(numpy.int8)
    308     v = numpy.subtract(v_plane, yuv_off[2]).view(numpy.int8)
    309     u = u.reshape(h/2, w/2).repeat(2, axis=1).repeat(2, axis=0)
    310     v = v.reshape(h/2, w/2).repeat(2, axis=1).repeat(2, axis=0)
    311     yuv = numpy.dstack([y, u.reshape(w*h), v.reshape(w*h)])
    312     flt = numpy.empty([h, w, 3], dtype=numpy.float32)
    313     flt.reshape(w*h*3)[:] = yuv.reshape(h*w*3)[:]
    314     flt = numpy.dot(flt.reshape(w*h,3), ccm_yuv_to_rgb.T).clip(0, 255)
    315     rgb = numpy.empty([h, w, 3], dtype=numpy.uint8)
    316     rgb.reshape(w*h*3)[:] = flt.reshape(w*h*3)[:]
    317     return rgb.astype(numpy.float32) / 255.0
    318 
    319 def load_yuv420_to_rgb_image(yuv_fname,
    320                              w, h,
    321                              ccm_yuv_to_rgb=DEFAULT_YUV_TO_RGB_CCM,
    322                              yuv_off=DEFAULT_YUV_OFFSETS):
    323     """Load a YUV420 image file, and return as an RGB image.
    324 
    325     Args:
    326         yuv_fname: The path of the YUV420 file.
    327         w: The width of the image.
    328         h: The height of the image.
    329         ccm_yuv_to_rgb: (Optional) the 3x3 CCM to convert from YUV to RGB.
    330         yuv_off: (Optional) offsets to subtract from each of Y,U,V values.
    331 
    332     Returns:
    333         RGB float-3 image array, with pixel values in [0.0, 1.0].
    334     """
    335     with open(yuv_fname, "rb") as f:
    336         y = numpy.fromfile(f, numpy.uint8, w*h, "")
    337         v = numpy.fromfile(f, numpy.uint8, w*h/4, "")
    338         u = numpy.fromfile(f, numpy.uint8, w*h/4, "")
    339         return convert_yuv420_to_rgb_image(y,u,v,w,h,ccm_yuv_to_rgb,yuv_off)
    340 
    341 def load_yuv420_to_yuv_planes(yuv_fname, w, h):
    342     """Load a YUV420 image file, and return separate Y, U, and V plane images.
    343 
    344     Args:
    345         yuv_fname: The path of the YUV420 file.
    346         w: The width of the image.
    347         h: The height of the image.
    348 
    349     Returns:
    350         Separate Y, U, and V images as float-1 Numpy arrays, pixels in [0,1].
    351         Note that pixel (0,0,0) is not black, since U,V pixels are centered at
    352         0.5, and also that the Y and U,V plane images returned are different
    353         sizes (due to chroma subsampling in the YUV420 format).
    354     """
    355     with open(yuv_fname, "rb") as f:
    356         y = numpy.fromfile(f, numpy.uint8, w*h, "")
    357         v = numpy.fromfile(f, numpy.uint8, w*h/4, "")
    358         u = numpy.fromfile(f, numpy.uint8, w*h/4, "")
    359         return ((y.astype(numpy.float32) / 255.0).reshape(h, w, 1),
    360                 (u.astype(numpy.float32) / 255.0).reshape(h/2, w/2, 1),
    361                 (v.astype(numpy.float32) / 255.0).reshape(h/2, w/2, 1))
    362 
    363 def decompress_jpeg_to_rgb_image(jpeg_buffer):
    364     """Decompress a JPEG-compressed image, returning as an RGB image.
    365 
    366     Args:
    367         jpeg_buffer: The JPEG stream.
    368 
    369     Returns:
    370         A numpy array for the RGB image, with pixels in [0,1].
    371     """
    372     img = Image.open(cStringIO.StringIO(jpeg_buffer))
    373     w = img.size[0]
    374     h = img.size[1]
    375     return numpy.array(img).reshape(h,w,3) / 255.0
    376 
    377 def apply_lut_to_image(img, lut):
    378     """Applies a LUT to every pixel in a float image array.
    379 
    380     Internally converts to a 16b integer image, since the LUT can work with up
    381     to 16b->16b mappings (i.e. values in the range [0,65535]). The lut can also
    382     have fewer than 65536 entries, however it must be sized as a power of 2
    383     (and for smaller luts, the scale must match the bitdepth).
    384 
    385     For a 16b lut of 65536 entries, the operation performed is:
    386 
    387         lut[r * 65535] / 65535 -> r'
    388         lut[g * 65535] / 65535 -> g'
    389         lut[b * 65535] / 65535 -> b'
    390 
    391     For a 10b lut of 1024 entries, the operation becomes:
    392 
    393         lut[r * 1023] / 1023 -> r'
    394         lut[g * 1023] / 1023 -> g'
    395         lut[b * 1023] / 1023 -> b'
    396 
    397     Args:
    398         img: Numpy float image array, with pixel values in [0,1].
    399         lut: Numpy table encoding a LUT, mapping 16b integer values.
    400 
    401     Returns:
    402         Float image array after applying LUT to each pixel.
    403     """
    404     n = len(lut)
    405     if n <= 0 or n > MAX_LUT_SIZE or (n & (n - 1)) != 0:
    406         raise its.error.Error('Invalid arg LUT size: %d' % (n))
    407     m = float(n-1)
    408     return (lut[(img * m).astype(numpy.uint16)] / m).astype(numpy.float32)
    409 
    410 def apply_matrix_to_image(img, mat):
    411     """Multiplies a 3x3 matrix with each float-3 image pixel.
    412 
    413     Each pixel is considered a column vector, and is left-multiplied by
    414     the given matrix:
    415 
    416         [     ]   r    r'
    417         [ mat ] * g -> g'
    418         [     ]   b    b'
    419 
    420     Args:
    421         img: Numpy float image array, with pixel values in [0,1].
    422         mat: Numpy 3x3 matrix.
    423 
    424     Returns:
    425         The numpy float-3 image array resulting from the matrix mult.
    426     """
    427     h = img.shape[0]
    428     w = img.shape[1]
    429     img2 = numpy.empty([h, w, 3], dtype=numpy.float32)
    430     img2.reshape(w*h*3)[:] = (numpy.dot(img.reshape(h*w, 3), mat.T)
    431                              ).reshape(w*h*3)[:]
    432     return img2
    433 
    434 def get_image_patch(img, xnorm, ynorm, wnorm, hnorm):
    435     """Get a patch (tile) of an image.
    436 
    437     Args:
    438         img: Numpy float image array, with pixel values in [0,1].
    439         xnorm,ynorm,wnorm,hnorm: Normalized (in [0,1]) coords for the tile.
    440 
    441     Returns:
    442         Float image array of the patch.
    443     """
    444     hfull = img.shape[0]
    445     wfull = img.shape[1]
    446     xtile = math.ceil(xnorm * wfull)
    447     ytile = math.ceil(ynorm * hfull)
    448     wtile = math.floor(wnorm * wfull)
    449     htile = math.floor(hnorm * hfull)
    450     return img[ytile:ytile+htile,xtile:xtile+wtile,:].copy()
    451 
    452 def compute_image_means(img):
    453     """Calculate the mean of each color channel in the image.
    454 
    455     Args:
    456         img: Numpy float image array, with pixel values in [0,1].
    457 
    458     Returns:
    459         A list of mean values, one per color channel in the image.
    460     """
    461     means = []
    462     chans = img.shape[2]
    463     for i in xrange(chans):
    464         means.append(numpy.mean(img[:,:,i], dtype=numpy.float64))
    465     return means
    466 
    467 def compute_image_variances(img):
    468     """Calculate the variance of each color channel in the image.
    469 
    470     Args:
    471         img: Numpy float image array, with pixel values in [0,1].
    472 
    473     Returns:
    474         A list of mean values, one per color channel in the image.
    475     """
    476     variances = []
    477     chans = img.shape[2]
    478     for i in xrange(chans):
    479         variances.append(numpy.var(img[:,:,i], dtype=numpy.float64))
    480     return variances
    481 
    482 def write_image(img, fname, apply_gamma=False):
    483     """Save a float-3 numpy array image to a file.
    484 
    485     Supported formats: PNG, JPEG, and others; see PIL docs for more.
    486 
    487     Image can be 3-channel, which is interpreted as RGB, or can be 1-channel,
    488     which is greyscale.
    489 
    490     Can optionally specify that the image should be gamma-encoded prior to
    491     writing it out; this should be done if the image contains linear pixel
    492     values, to make the image look "normal".
    493 
    494     Args:
    495         img: Numpy image array data.
    496         fname: Path of file to save to; the extension specifies the format.
    497         apply_gamma: (Optional) apply gamma to the image prior to writing it.
    498     """
    499     if apply_gamma:
    500         img = apply_lut_to_image(img, DEFAULT_GAMMA_LUT)
    501     (h, w, chans) = img.shape
    502     if chans == 3:
    503         Image.fromarray((img * 255.0).astype(numpy.uint8), "RGB").save(fname)
    504     elif chans == 1:
    505         img3 = (img * 255.0).astype(numpy.uint8).repeat(3).reshape(h,w,3)
    506         Image.fromarray(img3, "RGB").save(fname)
    507     else:
    508         raise its.error.Error('Unsupported image type')
    509 
    510 def downscale_image(img, f):
    511     """Shrink an image by a given integer factor.
    512 
    513     This function computes output pixel values by averaging over rectangular
    514     regions of the input image; it doesn't skip or sample pixels, and all input
    515     image pixels are evenly weighted.
    516 
    517     If the downscaling factor doesn't cleanly divide the width and/or height,
    518     then the remaining pixels on the right or bottom edge are discarded prior
    519     to the downscaling.
    520 
    521     Args:
    522         img: The input image as an ndarray.
    523         f: The downscaling factor, which should be an integer.
    524 
    525     Returns:
    526         The new (downscaled) image, as an ndarray.
    527     """
    528     h,w,chans = img.shape
    529     f = int(f)
    530     assert(f >= 1)
    531     h = (h/f)*f
    532     w = (w/f)*f
    533     img = img[0:h:,0:w:,::]
    534     chs = []
    535     for i in xrange(chans):
    536         ch = img.reshape(h*w*chans)[i::chans].reshape(h,w)
    537         ch = ch.reshape(h,w/f,f).mean(2).reshape(h,w/f)
    538         ch = ch.T.reshape(w/f,h/f,f).mean(2).T.reshape(h/f,w/f)
    539         chs.append(ch.reshape(h*w/(f*f)))
    540     img = numpy.vstack(chs).T.reshape(h/f,w/f,chans)
    541     return img
    542 
    543 def __measure_color_checker_patch(img, xc,yc, patch_size):
    544     r = patch_size/2
    545     tile = img[yc-r:yc+r+1:, xc-r:xc+r+1:, ::]
    546     means = tile.mean(1).mean(0)
    547     return means
    548 
    549 def get_color_checker_chart_patches(img, debug_fname_prefix=None):
    550     """Return the center coords of each patch in a color checker chart.
    551 
    552     Assumptions:
    553     * Chart is vertical or horizontal w.r.t. camera, but not diagonal.
    554     * Chart is (roughly) planar-parallel to the camera.
    555     * Chart is centered in frame (roughly).
    556     * Around/behind chart is white/grey background.
    557     * The only black pixels in the image are from the chart.
    558     * Chart is 100% visible and contained within image.
    559     * No other objects within image.
    560     * Image is well-exposed.
    561     * Standard color checker chart with standard-sized black borders.
    562 
    563     The values returned are in the coordinate system of the chart; that is,
    564     the "origin" patch is the brown patch that is in the chart's top-left
    565     corner when it is in the normal upright/horizontal orientation. (The chart
    566     may be any of the four main orientations in the image.)
    567 
    568     The chart is 6x4 patches in the normal upright orientation. The return
    569     values of this function are the center coordinate of the top-left patch,
    570     and the displacement vectors to the next patches to the right and below
    571     the top-left patch. From these pieces of data, the center coordinates of
    572     any of the patches can be computed.
    573 
    574     Args:
    575         img: Input image, as a numpy array with pixels in [0,1].
    576         debug_fname_prefix: If not None, the (string) name of a file prefix to
    577             use to save a number of debug images for visulaizing the output of
    578             this function; can be used to see if the patches are being found
    579             successfully.
    580 
    581     Returns:
    582         6x4 list of lists of integer (x,y) coords of the center of each patch,
    583         ordered in the "chart order" (6x4 row major).
    584     """
    585 
    586     # Shrink the original image.
    587     DOWNSCALE_FACTOR = 4
    588     img_small = downscale_image(img, DOWNSCALE_FACTOR)
    589 
    590     # Make a threshold image, which is 1.0 where the image is black,
    591     # and 0.0 elsewhere.
    592     BLACK_PIXEL_THRESH = 0.2
    593     mask_img = scipy.stats.threshold(
    594                 img_small.max(2), BLACK_PIXEL_THRESH, 1.1, 0.0)
    595     mask_img = 1.0 - scipy.stats.threshold(mask_img, -0.1, 0.1, 1.0)
    596 
    597     if debug_fname_prefix is not None:
    598         h,w = mask_img.shape
    599         write_image(img, debug_fname_prefix+"_0.jpg")
    600         write_image(mask_img.repeat(3).reshape(h,w,3),
    601                 debug_fname_prefix+"_1.jpg")
    602 
    603     # Mask image flattened to a single row or column (by averaging).
    604     # Also apply a threshold to these arrays.
    605     FLAT_PIXEL_THRESH = 0.05
    606     flat_row = mask_img.mean(0)
    607     flat_col = mask_img.mean(1)
    608     flat_row = [0 if v < FLAT_PIXEL_THRESH else 1 for v in flat_row]
    609     flat_col = [0 if v < FLAT_PIXEL_THRESH else 1 for v in flat_col]
    610 
    611     # Start and end of the non-zero region of the flattened row/column.
    612     flat_row_nonzero = [i for i in range(len(flat_row)) if flat_row[i]>0]
    613     flat_col_nonzero = [i for i in range(len(flat_col)) if flat_col[i]>0]
    614     flat_row_min, flat_row_max = min(flat_row_nonzero), max(flat_row_nonzero)
    615     flat_col_min, flat_col_max = min(flat_col_nonzero), max(flat_col_nonzero)
    616 
    617     # Orientation of chart, and number of grid cells horz. and vertically.
    618     orient = "h" if flat_row_max-flat_row_min>flat_col_max-flat_col_min else "v"
    619     xgrids = 6 if orient=="h" else 4
    620     ygrids = 6 if orient=="v" else 4
    621 
    622     # Get better bounds on the patches region, lopping off some of the excess
    623     # black border.
    624     HRZ_BORDER_PAD_FRAC = 0.0138
    625     VERT_BORDER_PAD_FRAC = 0.0395
    626     xpad = HRZ_BORDER_PAD_FRAC if orient=="h" else VERT_BORDER_PAD_FRAC
    627     ypad = HRZ_BORDER_PAD_FRAC if orient=="v" else VERT_BORDER_PAD_FRAC
    628     xchart = flat_row_min + (flat_row_max - flat_row_min) * xpad
    629     ychart = flat_col_min + (flat_col_max - flat_col_min) * ypad
    630     wchart = (flat_row_max - flat_row_min) * (1 - 2*xpad)
    631     hchart = (flat_col_max - flat_col_min) * (1 - 2*ypad)
    632 
    633     # Get the colors of the 4 corner patches, in clockwise order, by measuring
    634     # the average value of a small patch at each of the 4 patch centers.
    635     colors = []
    636     centers = []
    637     for (x,y) in [(0,0), (xgrids-1,0), (xgrids-1,ygrids-1), (0,ygrids-1)]:
    638         xc = xchart + (x + 0.5)*wchart/xgrids
    639         yc = ychart + (y + 0.5)*hchart/ygrids
    640         xc = int(xc * DOWNSCALE_FACTOR + 0.5)
    641         yc = int(yc * DOWNSCALE_FACTOR + 0.5)
    642         centers.append((xc,yc))
    643         chan_means = __measure_color_checker_patch(img, xc,yc, 32)
    644         colors.append(sum(chan_means) / len(chan_means))
    645 
    646     # The brightest corner is the white patch, the darkest is the black patch.
    647     # The black patch should be counter-clockwise from the white patch.
    648     white_patch_index = None
    649     for i in range(4):
    650         if colors[i] == max(colors) and \
    651                 colors[(i-1+4)%4] == min(colors):
    652             white_patch_index = i%4
    653     assert(white_patch_index is not None)
    654 
    655     # Return the coords of the origin (top-left when the chart is in the normal
    656     # upright orientation) patch's center, and the vector displacement to the
    657     # center of the second patch on the first row of the chart (when in the
    658     # normal upright orienation).
    659     origin_index = (white_patch_index+1)%4
    660     prev_index = (origin_index-1+4)%4
    661     next_index = (origin_index+1)%4
    662     origin_center = centers[origin_index]
    663     prev_center = centers[prev_index]
    664     next_center = centers[next_index]
    665     vec_across = tuple([(next_center[i]-origin_center[i])/5.0 for i in [0,1]])
    666     vec_down = tuple([(prev_center[i]-origin_center[i])/3.0 for i in [0,1]])
    667 
    668     # Compute the center of each patch.
    669     patches = [[],[],[],[]]
    670     for yi in range(4):
    671         for xi in range(6):
    672             x0,y0 = origin_center
    673             dxh,dyh = vec_across
    674             dxv,dyv = vec_down
    675             xc = int(x0 + dxh*xi + dxv*yi)
    676             yc = int(y0 + dyh*xi + dyv*yi)
    677             patches[yi].append((xc,yc))
    678 
    679     # Sanity check: test that the R,G,B,black,white patches are correct.
    680     patch_info = [(2,2,[0]), # Red
    681                   (2,1,[1]), # Green
    682                   (2,0,[2]), # Blue
    683                   (3,0,[0,1,2]), # White
    684                   (3,5,[])] # Black
    685     for i in range(len(patch_info)):
    686         yi,xi,high_chans = patch_info[i]
    687         low_chans = [i for i in [0,1,2] if i not in high_chans]
    688         xc,yc = patches[yi][xi]
    689         means = __measure_color_checker_patch(img, xc,yc, 64)
    690         if (min([means[i] for i in high_chans]+[1]) < \
    691                 max([means[i] for i in low_chans]+[0])):
    692             print "Color patch sanity check failed: patch", i
    693             # If the debug info is requested, then don't assert that the patches
    694             # are matched, to allow the caller to see the output.
    695             if debug_fname_prefix is None:
    696                 assert(0)
    697 
    698     if debug_fname_prefix is not None:
    699         for (xc,yc) in sum(patches,[]):
    700             img[yc,xc] = 1.0
    701         write_image(img, debug_fname_prefix+"_2.jpg")
    702 
    703     return patches
    704 
    705 class __UnitTest(unittest.TestCase):
    706     """Run a suite of unit tests on this module.
    707     """
    708 
    709     # TODO: Add more unit tests.
    710 
    711     def test_apply_matrix_to_image(self):
    712         """Unit test for apply_matrix_to_image.
    713 
    714         Test by using a canned set of values on a 1x1 pixel image.
    715 
    716             [ 1 2 3 ]   [ 0.1 ]   [ 1.4 ]
    717             [ 4 5 6 ] * [ 0.2 ] = [ 3.2 ]
    718             [ 7 8 9 ]   [ 0.3 ]   [ 5.0 ]
    719                mat         x         y
    720         """
    721         mat = numpy.array([[1,2,3],[4,5,6],[7,8,9]])
    722         x = numpy.array([0.1,0.2,0.3]).reshape(1,1,3)
    723         y = apply_matrix_to_image(x, mat).reshape(3).tolist()
    724         y_ref = [1.4,3.2,5.0]
    725         passed = all([math.fabs(y[i] - y_ref[i]) < 0.001 for i in xrange(3)])
    726         self.assertTrue(passed)
    727 
    728     def test_apply_lut_to_image(self):
    729         """ Unit test for apply_lut_to_image.
    730 
    731         Test by using a canned set of values on a 1x1 pixel image. The LUT will
    732         simply double the value of the index:
    733 
    734             lut[x] = 2*x
    735         """
    736         lut = numpy.array([2*i for i in xrange(65536)])
    737         x = numpy.array([0.1,0.2,0.3]).reshape(1,1,3)
    738         y = apply_lut_to_image(x, lut).reshape(3).tolist()
    739         y_ref = [0.2,0.4,0.6]
    740         passed = all([math.fabs(y[i] - y_ref[i]) < 0.001 for i in xrange(3)])
    741         self.assertTrue(passed)
    742 
    743 if __name__ == '__main__':
    744     unittest.main()
    745 
    746