• 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"] == "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
81def 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
101def 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
131def 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
188def 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
218def 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
238def 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
288def 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
319def 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
341def 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
363def 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
377def 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
410def 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
434def 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
452def 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
467def 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
482def 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
510def 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
543def __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
549def 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
705class __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
743if __name__ == '__main__':
744    unittest.main()
745
746