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