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 27 DEFAULT_YUV_TO_RGB_CCM = numpy.matrix([ 28 [1.000, 0.000, 1.402], 29 [1.000, -0.344, -0.714], 30 [1.000, 1.772, 0.000]]) 31 32 DEFAULT_YUV_OFFSETS = numpy.array([0, 128, 128]) 33 34 DEFAULT_GAMMA_LUT = numpy.array( 35 [math.floor(65535 * math.pow(i/65535.0, 1/2.2) + 0.5) for i in xrange(65536)]) 36 37 DEFAULT_INVGAMMA_LUT = numpy.array( 38 [math.floor(65535 * math.pow(i/65535.0, 2.2) + 0.5) for i in xrange(65536)]) 39 40 MAX_LUT_SIZE = 65536 41 42 def convert_capture_to_rgb_image(cap, 43 ccm_yuv_to_rgb=DEFAULT_YUV_TO_RGB_CCM, 44 yuv_off=DEFAULT_YUV_OFFSETS): 45 """Convert a captured image object to a RGB image. 46 47 Args: 48 cap: A capture object as returned by its.device.do_capture. 49 ccm_yuv_to_rgb: (Optional) the 3x3 CCM to convert from YUV to RGB. 50 yuv_off: (Optional) offsets to subtract from each of Y,U,V values. 51 52 Returns: 53 RGB float-3 image array, with pixel values in [0.0, 1.0]. 54 """ 55 w = cap["width"] 56 h = cap["height"] 57 if cap["format"] == "yuv": 58 y = cap["data"][0:w*h] 59 u = cap["data"][w*h:w*h*5/4] 60 v = cap["data"][w*h*5/4:w*h*6/4] 61 return convert_yuv420_to_rgb_image(y, u, v, w, h) 62 elif cap["format"] == "jpeg": 63 # TODO: Convert JPEG to RGB. 64 raise its.error.Error('Invalid format %s' % (cap["format"])) 65 else: 66 raise its.error.Error('Invalid format %s' % (cap["format"])) 67 68 def convert_capture_to_yuv_planes(cap): 69 """Convert a captured image object to separate Y,U,V image planes. 70 71 The only input format that is supported is planar YUV420, and the planes 72 that are returned are such that the U,V planes are 1/2 x 1/2 of the Y 73 plane size. 74 75 Args: 76 cap: A capture object as returned by its.device.do_capture. 77 78 Returns: 79 Three float arrays, for the Y,U,V planes, with pixel values in [0,1]. 80 """ 81 w = cap["width"] 82 h = cap["height"] 83 if cap["format"] == "yuv": 84 y = cap["data"][0:w*h] 85 u = cap["data"][w*h:w*h*5/4] 86 v = cap["data"][w*h*5/4:w*h*6/4] 87 return ((y.astype(numpy.float32) / 255.0).reshape(h, w, 1), 88 (u.astype(numpy.float32) / 255.0).reshape(h/2, w/2, 1), 89 (v.astype(numpy.float32) / 255.0).reshape(h/2, w/2, 1)) 90 else: 91 raise its.error.Error('Invalid format %s' % (cap["format"])) 92 93 def convert_yuv420_to_rgb_image(y_plane, u_plane, v_plane, 94 w, h, 95 ccm_yuv_to_rgb=DEFAULT_YUV_TO_RGB_CCM, 96 yuv_off=DEFAULT_YUV_OFFSETS): 97 """Convert a YUV420 8-bit planar image to an RGB image. 98 99 Args: 100 y_plane: The packed 8-bit Y plane. 101 u_plane: The packed 8-bit U plane. 102 v_plane: The packed 8-bit V plane. 103 w: The width of the image. 104 h: The height of the image. 105 ccm_yuv_to_rgb: (Optional) the 3x3 CCM to convert from YUV to RGB. 106 yuv_off: (Optional) offsets to subtract from each of Y,U,V values. 107 108 Returns: 109 RGB float-3 image array, with pixel values in [0.0, 1.0]. 110 """ 111 y = numpy.subtract(y_plane, yuv_off[0]) 112 u = numpy.subtract(u_plane, yuv_off[1]).view(numpy.int8) 113 v = numpy.subtract(v_plane, yuv_off[2]).view(numpy.int8) 114 u = u.reshape(h/2, w/2).repeat(2, axis=1).repeat(2, axis=0) 115 v = v.reshape(h/2, w/2).repeat(2, axis=1).repeat(2, axis=0) 116 yuv = numpy.dstack([y, u.reshape(w*h), v.reshape(w*h)]) 117 flt = numpy.empty([h, w, 3], dtype=numpy.float32) 118 flt.reshape(w*h*3)[:] = yuv.reshape(h*w*3)[:] 119 flt = numpy.dot(flt.reshape(w*h,3), ccm_yuv_to_rgb.T).clip(0, 255) 120 rgb = numpy.empty([h, w, 3], dtype=numpy.uint8) 121 rgb.reshape(w*h*3)[:] = flt.reshape(w*h*3)[:] 122 return rgb.astype(numpy.float32) / 255.0 123 124 def load_yuv420_to_rgb_image(yuv_fname, 125 w, h, 126 ccm_yuv_to_rgb=DEFAULT_YUV_TO_RGB_CCM, 127 yuv_off=DEFAULT_YUV_OFFSETS): 128 """Load a YUV420 image file, and return as an RGB image. 129 130 Args: 131 yuv_fname: The path of the YUV420 file. 132 w: The width of the image. 133 h: The height of the image. 134 ccm_yuv_to_rgb: (Optional) the 3x3 CCM to convert from YUV to RGB. 135 yuv_off: (Optional) offsets to subtract from each of Y,U,V values. 136 137 Returns: 138 RGB float-3 image array, with pixel values in [0.0, 1.0]. 139 """ 140 with open(yuv_fname, "rb") as f: 141 y = numpy.fromfile(f, numpy.uint8, w*h, "") 142 v = numpy.fromfile(f, numpy.uint8, w*h/4, "") 143 u = numpy.fromfile(f, numpy.uint8, w*h/4, "") 144 return convert_yuv420_to_rgb_image(y,u,v,w,h,ccm_yuv_to_rgb,yuv_off) 145 146 def load_yuv420_to_yuv_planes(yuv_fname, w, h): 147 """Load a YUV420 image file, and return separate Y, U, and V plane images. 148 149 Args: 150 yuv_fname: The path of the YUV420 file. 151 w: The width of the image. 152 h: The height of the image. 153 154 Returns: 155 Separate Y, U, and V images as float-1 Numpy arrays, pixels in [0,1]. 156 Note that pixel (0,0,0) is not black, since U,V pixels are centered at 157 0.5, and also that the Y and U,V plane images returned are different 158 sizes (due to chroma subsampling in the YUV420 format). 159 """ 160 with open(yuv_fname, "rb") as f: 161 y = numpy.fromfile(f, numpy.uint8, w*h, "") 162 v = numpy.fromfile(f, numpy.uint8, w*h/4, "") 163 u = numpy.fromfile(f, numpy.uint8, w*h/4, "") 164 return ((y.astype(numpy.float32) / 255.0).reshape(h, w, 1), 165 (u.astype(numpy.float32) / 255.0).reshape(h/2, w/2, 1), 166 (v.astype(numpy.float32) / 255.0).reshape(h/2, w/2, 1)) 167 168 def decompress_jpeg_to_rgb_image(jpeg_buffer): 169 """Decompress a JPEG-compressed image, returning as an RGB image. 170 171 Args: 172 jpeg_buffer: The JPEG stream. 173 174 Returns: 175 A numpy array for the RGB image, with pixels in [0,1]. 176 """ 177 img = Image.open(cStringIO.StringIO(jpeg_buffer)) 178 w = img.size[0] 179 h = img.size[1] 180 return numpy.array(img).reshape(h,w,3) / 255.0 181 182 def apply_lut_to_image(img, lut): 183 """Applies a LUT to every pixel in a float image array. 184 185 Internally converts to a 16b integer image, since the LUT can work with up 186 to 16b->16b mappings (i.e. values in the range [0,65535]). The lut can also 187 have fewer than 65536 entries, however it must be sized as a power of 2 188 (and for smaller luts, the scale must match the bitdepth). 189 190 For a 16b lut of 65536 entries, the operation performed is: 191 192 lut[r * 65535] / 65535 -> r' 193 lut[g * 65535] / 65535 -> g' 194 lut[b * 65535] / 65535 -> b' 195 196 For a 10b lut of 1024 entries, the operation becomes: 197 198 lut[r * 1023] / 1023 -> r' 199 lut[g * 1023] / 1023 -> g' 200 lut[b * 1023] / 1023 -> b' 201 202 Args: 203 img: Numpy float image array, with pixel values in [0,1]. 204 lut: Numpy table encoding a LUT, mapping 16b integer values. 205 206 Returns: 207 Float image array after applying LUT to each pixel. 208 """ 209 n = len(lut) 210 if n <= 0 or n > MAX_LUT_SIZE or (n & (n - 1)) != 0: 211 raise its.error.Error('Invalid arg LUT size: %d' % (n)) 212 m = float(n-1) 213 return (lut[(img * m).astype(numpy.uint16)] / m).astype(numpy.float32) 214 215 def apply_matrix_to_image(img, mat): 216 """Multiplies a 3x3 matrix with each float-3 image pixel. 217 218 Each pixel is considered a column vector, and is left-multiplied by 219 the given matrix: 220 221 [ ] r r' 222 [ mat ] * g -> g' 223 [ ] b b' 224 225 Args: 226 img: Numpy float image array, with pixel values in [0,1]. 227 mat: Numpy 3x3 matrix. 228 229 Returns: 230 The numpy float-3 image array resulting from the matrix mult. 231 """ 232 h = img.shape[0] 233 w = img.shape[1] 234 img2 = numpy.empty([h, w, 3], dtype=numpy.float32) 235 img2.reshape(w*h*3)[:] = (numpy.dot(img.reshape(h*w, 3), mat.T) 236 ).reshape(w*h*3)[:] 237 return img2 238 239 def get_image_patch(img, xnorm, ynorm, wnorm, hnorm): 240 """Get a patch (tile) of an image. 241 242 Args: 243 img: Numpy float image array, with pixel values in [0,1]. 244 xnorm,ynorm,wnorm,hnorm: Normalized (in [0,1]) coords for the tile. 245 246 Returns: 247 Float image array of the patch. 248 """ 249 hfull = img.shape[0] 250 wfull = img.shape[1] 251 xtile = math.ceil(xnorm * wfull) 252 ytile = math.ceil(ynorm * hfull) 253 wtile = math.floor(wnorm * wfull) 254 htile = math.floor(hnorm * hfull) 255 return img[ytile:ytile+htile,xtile:xtile+wtile,:].copy() 256 257 def compute_image_means(img): 258 """Calculate the mean of each color channel in the image. 259 260 Args: 261 img: Numpy float image array, with pixel values in [0,1]. 262 263 Returns: 264 A list of mean values, one per color channel in the image. 265 """ 266 means = [] 267 chans = img.shape[2] 268 for i in xrange(chans): 269 means.append(numpy.mean(img[:,:,i], dtype=numpy.float64)) 270 return means 271 272 def compute_image_variances(img): 273 """Calculate the variance of each color channel in the image. 274 275 Args: 276 img: Numpy float image array, with pixel values in [0,1]. 277 278 Returns: 279 A list of mean values, one per color channel in the image. 280 """ 281 variances = [] 282 chans = img.shape[2] 283 for i in xrange(chans): 284 variances.append(numpy.var(img[:,:,i], dtype=numpy.float64)) 285 return variances 286 287 def write_image(img, fname, apply_gamma=False): 288 """Save a float-3 numpy array image to a file. 289 290 Supported formats: PNG, JPEG, and others; see PIL docs for more. 291 292 Image can be 3-channel, which is interpreted as RGB, or can be 1-channel, 293 which is greyscale. 294 295 Can optionally specify that the image should be gamma-encoded prior to 296 writing it out; this should be done if the image contains linear pixel 297 values, to make the image look "normal". 298 299 Args: 300 img: Numpy image array data. 301 fname: Path of file to save to; the extension specifies the format. 302 apply_gamma: (Optional) apply gamma to the image prior to writing it. 303 """ 304 if apply_gamma: 305 img = apply_lut_to_image(img, DEFAULT_GAMMA_LUT) 306 (h, w, chans) = img.shape 307 if chans == 3: 308 Image.fromarray((img * 255.0).astype(numpy.uint8), "RGB").save(fname) 309 elif chans == 1: 310 img3 = (img * 255.0).astype(numpy.uint8).repeat(3).reshape(h,w,3) 311 Image.fromarray(img3, "RGB").save(fname) 312 else: 313 raise its.error.Error('Unsupported image type') 314 315 class __UnitTest(unittest.TestCase): 316 """Run a suite of unit tests on this module. 317 """ 318 319 # TODO: Add more unit tests. 320 321 def test_apply_matrix_to_image(self): 322 """Unit test for apply_matrix_to_image. 323 324 Test by using a canned set of values on a 1x1 pixel image. 325 326 [ 1 2 3 ] [ 0.1 ] [ 1.4 ] 327 [ 4 5 6 ] * [ 0.2 ] = [ 3.2 ] 328 [ 7 8 9 ] [ 0.3 ] [ 5.0 ] 329 mat x y 330 """ 331 mat = numpy.array([[1,2,3],[4,5,6],[7,8,9]]) 332 x = numpy.array([0.1,0.2,0.3]).reshape(1,1,3) 333 y = apply_matrix_to_image(x, mat).reshape(3).tolist() 334 y_ref = [1.4,3.2,5.0] 335 passed = all([math.fabs(y[i] - y_ref[i]) < 0.001 for i in xrange(3)]) 336 self.assertTrue(passed) 337 338 def test_apply_lut_to_image(self): 339 """ Unit test for apply_lut_to_image. 340 341 Test by using a canned set of values on a 1x1 pixel image. The LUT will 342 simply double the value of the index: 343 344 lut[x] = 2*x 345 """ 346 lut = numpy.array([2*i for i in xrange(65536)]) 347 x = numpy.array([0.1,0.2,0.3]).reshape(1,1,3) 348 y = apply_lut_to_image(x, lut).reshape(3).tolist() 349 y_ref = [0.2,0.4,0.6] 350 passed = all([math.fabs(y[i] - y_ref[i]) < 0.001 for i in xrange(3)]) 351 self.assertTrue(passed) 352 353 if __name__ == '__main__': 354 unittest.main() 355 356
