1# Copyright 2020 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"""Verifies android.jpeg.quality increases JPEG image quality."""
15
16
17import logging
18import math
19import os.path
20
21from matplotlib import pylab
22import matplotlib.pyplot
23from mobly import test_runner
24import numpy as np
25
26import its_base_test
27import camera_properties_utils
28import capture_request_utils
29import image_processing_utils
30import its_session_utils
31
32_JPEG_APPN_MARKERS = [[255, 224], [255, 225], [255, 226], [255, 227],
33                      [255, 228], [255, 229], [255, 230], [255, 231],
34                      [255, 232], [255, 235]]
35_JPEG_DHT_MARKER = [255, 196]  # JPEG Define Huffman Table
36_JPEG_DQT_MARKER = [255, 219]  # JPEG Define Quantization Table
37_JPEG_DQT_RTOL = 0.8  # -20% for each +20 in jpeg.quality (empirical number)
38_JPEG_EOI_MARKER = [255, 217]  # JPEG End of Image
39_JPEG_SOI_MARKER = [255, 216]  # JPEG Start of Image
40_JPEG_SOS_MARKER = [255, 218]  # JPEG Start of Scan
41_NAME = os.path.splitext(os.path.basename(__file__))[0]
42_QUALITIES = [25, 45, 65, 85]
43_SYMBOLS = ['o', 's', 'v', '^', '<', '>']
44
45
46def is_square(integer):
47  root = math.sqrt(integer)
48  return integer == int(root + 0.5)**2
49
50
51def strip_soi_marker(jpeg):
52  """Strip off start of image marker.
53
54  SOI is of form [xFF xD8] and JPEG needs to start with marker.
55
56  Args:
57   jpeg: 1-D numpy int [0:255] array; values from JPEG capture
58
59  Returns:
60    jpeg with SOI marker stripped off.
61  """
62
63  soi = jpeg[0:2]
64  if list(soi) != _JPEG_SOI_MARKER:
65    raise AssertionError('JPEG has no Start Of Image marker')
66  return jpeg[2:]
67
68
69def strip_appn_data(jpeg):
70  """Strip off application specific data at beginning of JPEG.
71
72  APPN markers are of form [xFF, xE*, size_msb, size_lsb] and should follow
73  SOI marker.
74
75  Args:
76   jpeg: 1-D numpy int [0:255] array; values from JPEG capture
77
78  Returns:
79    jpeg with APPN marker(s) and data stripped off.
80  """
81
82  length = 0
83  i = 0
84  # find APPN markers and strip off payloads at beginning of jpeg
85  while i < len(jpeg) - 1:
86    if [jpeg[i], jpeg[i + 1]] in _JPEG_APPN_MARKERS:
87      length = jpeg[i + 2] * 256 + jpeg[i + 3] + 2
88      logging.debug('stripped APPN length:%d', length)
89      jpeg = np.concatenate((jpeg[0:i], jpeg[length:]), axis=None)
90    elif ([jpeg[i], jpeg[i + 1]] == _JPEG_DQT_MARKER or
91          [jpeg[i], jpeg[i + 1]] == _JPEG_DHT_MARKER):
92      break
93    else:
94      i += 1
95
96  return jpeg
97
98
99def find_dqt_markers(marker, jpeg):
100  """Find location(s) of marker list in jpeg.
101
102  DQT marker is of form [xFF, xDB].
103
104  Args:
105    marker: list; marker values
106    jpeg: 1-D numpy int [0:255] array; JPEG capture w/ SOI & APPN stripped
107
108  Returns:
109    locs: list; marker locations in jpeg
110  """
111  locs = []
112  marker_len = len(marker)
113  for i in range(len(jpeg) - marker_len + 1):
114    if list(jpeg[i:i + marker_len]) == marker:
115      locs.append(i)
116  return locs
117
118
119def extract_dqts(jpeg, debug=False):
120  """Find and extract the DQT info in the JPEG.
121
122  SOI marker and APPN markers plus data are stripped off front of JPEG.
123  DQT marker is of form [xFF, xDB] followed by [size_msb, size_lsb].
124  Size includes the size values, but not the marker values.
125  Luma DQT is prefixed by 0, Chroma DQT by 1.
126  DQTs can have both luma & chroma or each individually.
127  There can be more than one DQT table for luma and chroma.
128
129  Args:
130   jpeg: 1-D numpy int [0:255] array; values from JPEG capture
131   debug: bool; command line flag to print debug data
132
133  Returns:
134    lumas,chromas: lists of numpy means of luma & chroma DQT matrices.
135    Higher values represent higher compression.
136  """
137
138  dqt_markers = find_dqt_markers(_JPEG_DQT_MARKER, jpeg)
139  logging.debug('DQT header loc(s):%s', dqt_markers)
140  lumas = []
141  chromas = []
142  for i, dqt in enumerate(dqt_markers):
143    if debug:
144      logging.debug('DQT %d start: %d, marker: %s, length: %s', i, dqt,
145                    jpeg[dqt:dqt + 2], jpeg[dqt + 2:dqt + 4])
146    dqt_size = jpeg[dqt + 2] * 256 + jpeg[dqt + 3] - 2  # strip off size marker
147    if dqt_size % 2 == 0:  # even payload means luma & chroma
148      logging.debug(' both luma & chroma DQT matrices in marker')
149      dqt_size = (dqt_size - 2) // 2  # subtact off luma/chroma markers
150      if not is_square(dqt_size):
151        raise AssertionError(f'DQT size: {dqt_size}')
152      luma_start = dqt + 5  # skip header, length, & matrix id
153      chroma_start = luma_start + dqt_size + 1  # skip lumen &  matrix_id
154      luma = np.array(jpeg[luma_start: luma_start + dqt_size])
155      chroma = np.array(jpeg[chroma_start: chroma_start + dqt_size])
156      lumas.append(np.mean(luma))
157      chromas.append(np.mean(chroma))
158      if debug:
159        h = int(math.sqrt(dqt_size))
160        logging.debug(' luma:%s', luma.reshape(h, h))
161        logging.debug(' chroma:%s', chroma.reshape(h, h))
162    else:  # odd payload means only 1 matrix
163      logging.debug(' single DQT matrix in marker')
164      dqt_size = dqt_size - 1  # subtract off luma/chroma marker
165      if not is_square(dqt_size):
166        raise AssertionError(f'DQT size: {dqt_size}')
167      start = dqt + 5
168      matrix = np.array(jpeg[start:start + dqt_size])
169      if jpeg[dqt + 4]:  # chroma == 1
170        chromas.append(np.mean(matrix))
171        if debug:
172          h = int(math.sqrt(dqt_size))
173          logging.debug(' chroma:%s', matrix.reshape(h, h))
174      else:  # luma == 0
175        lumas.append(np.mean(matrix))
176        if debug:
177          h = int(math.sqrt(dqt_size))
178          logging.debug(' luma:%s', matrix.reshape(h, h))
179
180  return lumas, chromas
181
182
183def plot_data(qualities, lumas, chromas, img_name):
184  """Create plot of data."""
185  logging.debug('qualities: %s', str(qualities))
186  logging.debug('luma DQT avgs: %s', str(lumas))
187  logging.debug('chroma DQT avgs: %s', str(chromas))
188  pylab.title(_NAME)
189  for i in range(lumas.shape[1]):
190    pylab.plot(
191        qualities, lumas[:, i], '-g' + _SYMBOLS[i], label='luma_dqt' + str(i))
192    pylab.plot(
193        qualities,
194        chromas[:, i],
195        '-r' + _SYMBOLS[i],
196        label='chroma_dqt' + str(i))
197  pylab.xlim([0, 100])
198  pylab.ylim([0, None])
199  pylab.xlabel('jpeg.quality')
200  pylab.ylabel('DQT luma/chroma matrix averages')
201  pylab.legend(loc='upper right', numpoints=1, fancybox=True)
202  matplotlib.pyplot.savefig(f'{img_name}_plot.png')
203
204
205class JpegQualityTest(its_base_test.ItsBaseTest):
206  """Test the camera JPEG compression quality.
207
208  Step JPEG qualities through android.jpeg.quality. Ensure quanitization
209  matrix decreases with quality increase. Matrix should decrease as the
210  matrix represents the division factor. Higher numbers --> fewer quantization
211  levels.
212  """
213
214  def test_jpeg_quality(self):
215    logging.debug('Starting %s', _NAME)
216    # init variables
217    lumas = []
218    chromas = []
219
220    with its_session_utils.ItsSession(
221        device_id=self.dut.serial,
222        camera_id=self.camera_id,
223        hidden_physical_id=self.hidden_physical_id) as cam:
224
225      props = cam.get_camera_properties()
226      props = cam.override_with_hidden_physical_camera_props(props)
227      debug = self.debug_mode
228
229      # Load chart for scene
230      its_session_utils.load_scene(
231          cam, props, self.scene, self.tablet,
232          its_session_utils.CHART_DISTANCE_NO_SCALING)
233
234      # Check skip conditions
235      camera_properties_utils.skip_unless(
236          camera_properties_utils.jpeg_quality(props))
237      cam.do_3a()
238
239      # do captures over jpeg quality range
240      req = capture_request_utils.auto_capture_request()
241      for q in _QUALITIES:
242        logging.debug('jpeg.quality: %.d', q)
243        req['android.jpeg.quality'] = q
244        cap = cam.do_capture(req, cam.CAP_JPEG)
245        jpeg = cap['data']
246
247        # strip off start of image
248        jpeg = strip_soi_marker(jpeg)
249
250        # strip off application specific data
251        jpeg = strip_appn_data(jpeg)
252        logging.debug('remaining JPEG header:%s', jpeg[0:4])
253
254        # find and extract DQTs
255        lumas_i, chromas_i = extract_dqts(jpeg, debug)
256        lumas.append(lumas_i)
257        chromas.append(chromas_i)
258
259        # save JPEG image
260        img = image_processing_utils.convert_capture_to_rgb_image(
261            cap, props=props)
262        img_name = os.path.join(self.log_path, _NAME)
263        image_processing_utils.write_image(img, f'{img_name}_{q}.jpg')
264
265    # turn lumas/chromas into np array to ease multi-dimensional plots/asserts
266    lumas = np.array(lumas)
267    chromas = np.array(chromas)
268
269    # create plot of luma & chroma averages vs quality
270    plot_data(_QUALITIES, lumas, chromas, img_name)
271
272    # assert decreasing luma/chroma with improved jpeg quality
273    for i in range(lumas.shape[1]):
274      l = lumas[:, i]
275      c = chromas[:, i]
276      if not all(y < x * _JPEG_DQT_RTOL for x, y in zip(l, l[1:])):
277        raise AssertionError(f'luma DQT avgs: {l}, RTOL: {_JPEG_DQT_RTOL}')
278
279      if not all(y < x * _JPEG_DQT_RTOL for x, y in zip(c, c[1:])):
280        raise AssertionError(f'chroma DQT avgs: {c}, RTOL: {_JPEG_DQT_RTOL}')
281
282if __name__ == '__main__':
283  test_runner.main()
284