Files
esp-idf/examples/peripherals/jpeg/jpeg_decode/pytest_jpeg_decode.py

240 lines
8.9 KiB
Python

# SPDX-FileCopyrightText: 2026 Espressif Systems (Shanghai) CO LTD
# SPDX-License-Identifier: CC0-1.0
import base64
import hashlib
import logging
import re
from dataclasses import dataclass
from pathlib import Path
import pytest
from pytest_embedded import Dut
from pytest_embedded_idf.utils import idf_parametrize
from pytest_embedded_idf.utils import soc_filtered_targets
DECODE_OUTPUT_NAME = 'jpeg_decode_result.ppm'
GOLDEN_OUTPUT_NAME = 'golden_output.ppm'
GOLDEN_OUTPUT_PATH = Path(__file__).with_name(GOLDEN_OUTPUT_NAME)
EXPECTED_PIXEL_FORMAT = 'RGB888'
EXPECTED_ENCODING = 'base64'
RGB888_BYTES_PER_PIXEL = 3
PPM_MAGIC = b'P6'
PPM_MAX_VALUE = b'255'
DECODE_INFO_PATTERN = (
r'JPEG_DECODE_INFO width=(?P<width>\d+) height=(?P<height>\d+) '
r'padded_width=(?P<padded_width>\d+) padded_height=(?P<padded_height>\d+) '
r'format=(?P<format>\w+) encoding=(?P<encoding>\w+) size=(?P<size>\d+)'
)
DECODE_INFO_RE = re.compile(DECODE_INFO_PATTERN)
DECODE_CHUNK_PATTERN = r'JPEG_DECODE_BASE64 (?P<payload>[A-Za-z0-9+/=]+)'
DECODE_CHUNK_RE = re.compile(DECODE_CHUNK_PATTERN)
PPM_HEADER_RE = re.compile(rb'^P6\s+(?P<width>\d+)\s+(?P<height>\d+)\s+(?P<max_value>\d+)\s')
@dataclass(frozen=True, slots=True)
class DecodeMetadata:
width: int
height: int
padded_width: int
padded_height: int
pixel_format: str
encoding: str
size: int
def __post_init__(self) -> None:
if self.width <= 0 or self.height <= 0:
raise ValueError(f'Invalid dimensions: {self.width}x{self.height}')
if self.padded_width < self.width or self.padded_height < self.height:
raise ValueError(
f'Padded size ({self.padded_width}x{self.padded_height}) '
f'smaller than visible size ({self.width}x{self.height})'
)
if self.pixel_format != EXPECTED_PIXEL_FORMAT:
raise ValueError(f'Unsupported pixel format: {self.pixel_format}')
if self.encoding != EXPECTED_ENCODING:
raise ValueError(f'Unsupported encoding: {self.encoding}')
@property
def padded_image_size(self) -> int:
return self.padded_width * self.padded_height * RGB888_BYTES_PER_PIXEL
@dataclass(frozen=True, slots=True)
class RgbImage:
width: int
height: int
pixels_rgb888: bytes
def __post_init__(self) -> None:
expected_size = self.width * self.height * RGB888_BYTES_PER_PIXEL
if len(self.pixels_rgb888) != expected_size:
raise ValueError(f'Expected {expected_size} RGB bytes, got {len(self.pixels_rgb888)}')
def parse_decode_metadata(meta_line: str) -> DecodeMetadata:
match = DECODE_INFO_RE.fullmatch(meta_line)
if not match:
raise ValueError(f'Invalid decode metadata line: {meta_line}')
return DecodeMetadata(
width=int(match.group('width')),
height=int(match.group('height')),
padded_width=int(match.group('padded_width')),
padded_height=int(match.group('padded_height')),
pixel_format=match.group('format'),
encoding=match.group('encoding'),
size=int(match.group('size')),
)
def collect_base64_payload(dut: Dut) -> list[str]:
payload_lines: list[str] = []
while True:
# The example prints the decoded frame as multiple short UART lines
# instead of one giant base64 blob, so collect and join them here.
match = dut.expect(rf'(?P<line>JPEG_DECODE_BASE64_END|{DECODE_CHUNK_PATTERN}\r?\n)', timeout=60)
line = match.group('line').decode('utf-8').strip()
if line == 'JPEG_DECODE_BASE64_END':
return payload_lines
chunk_match = DECODE_CHUNK_RE.fullmatch(line)
assert chunk_match is not None
payload_lines.append(chunk_match.group('payload'))
def _crop_visible_bgr888(raw_bytes: bytes, metadata: DecodeMetadata) -> bytes:
# The hardware can write into a padded decode buffer whose width/height are
# rounded up to JPEG block boundaries. Pytest only wants the visible image,
# so keep the useful bytes from each row and discard the padded tail rows.
if len(raw_bytes) != metadata.padded_image_size:
raise ValueError(f'Expected {metadata.padded_image_size} padded BGR bytes, got {len(raw_bytes)}')
visible_row_size = metadata.width * RGB888_BYTES_PER_PIXEL
padded_row_size = metadata.padded_width * RGB888_BYTES_PER_PIXEL
return b''.join(
raw_bytes[offset : offset + visible_row_size]
for offset in range(0, padded_row_size * metadata.height, padded_row_size)
)
def _bgr888_to_rgb888(raw_bytes: bytes) -> bytes:
# The decoder's RGB888 mode uses BGR24 byte layout by default. Swap the
# first and third byte in each pixel so the PPM artifact becomes standard
# RGB order that common desktop image tools expect.
rgb_bytes = bytearray(raw_bytes)
rgb_bytes[0::3], rgb_bytes[2::3] = raw_bytes[2::3], raw_bytes[0::3]
return bytes(rgb_bytes)
def decode_base64_image(metadata: DecodeMetadata, payload_lines: list[str]) -> RgbImage:
# The DUT sends the raw decode buffer as base64 over UART because the test
# environment only observes text logs. Rebuild bytes on the host, crop away
# decoder padding, then normalize the pixel order for image comparison.
raw_bytes = base64.b64decode(''.join(payload_lines), validate=True)
if len(raw_bytes) != metadata.size:
raise ValueError(f'Expected {metadata.size} decoded bytes, got {len(raw_bytes)}')
visible_bgr888 = _crop_visible_bgr888(raw_bytes, metadata)
return RgbImage(
width=metadata.width,
height=metadata.height,
pixels_rgb888=_bgr888_to_rgb888(visible_bgr888),
)
def _encode_ppm(image: RgbImage) -> bytes:
header = b'%s\n%d %d\n%s\n' % (PPM_MAGIC, image.width, image.height, PPM_MAX_VALUE)
return header + image.pixels_rgb888
def _load_ppm(path: Path) -> RgbImage:
ppm_bytes = path.read_bytes()
header_match = PPM_HEADER_RE.match(ppm_bytes)
if not header_match:
raise ValueError('Invalid PPM header')
width = int(header_match.group('width'))
height = int(header_match.group('height'))
max_value = header_match.group('max_value')
if width <= 0 or height <= 0:
raise ValueError('Unsupported PPM dimensions')
if max_value != PPM_MAX_VALUE:
raise ValueError(f'Unsupported PPM max value: {max_value.decode("ascii", errors="replace")}')
pixel_data = ppm_bytes[header_match.end() :]
return RgbImage(width=width, height=height, pixels_rgb888=pixel_data)
def save_ppm_artifact(image: RgbImage, output_path: Path) -> None:
output_path.parent.mkdir(parents=True, exist_ok=True)
try:
output_path.write_bytes(_encode_ppm(image))
except OSError:
logging.exception('Failed to save JPEG decode artifact to %s', output_path)
return
logging.info('Saved JPEG decode artifact to %s', output_path)
def image_digest(image: RgbImage) -> str:
digest = hashlib.sha256()
digest.update(image.width.to_bytes(4, 'big'))
digest.update(image.height.to_bytes(4, 'big'))
digest.update(image.pixels_rgb888)
return digest.hexdigest()
def assert_image_matches_golden(result_image: RgbImage, golden_path: Path) -> None:
assert golden_path.is_file(), f'Golden PPM not found: {golden_path}'
golden_image = _load_ppm(golden_path)
assert image_digest(result_image) == image_digest(golden_image), (
f'Generated image does not match golden file: {golden_path.name}'
)
def run_jpeg_decode_example(dut: Dut) -> None:
dut.expect_exact('Loading embedded JPEG from flash...')
dut.expect(r'Embedded JPEG size: \d+ bytes')
dut.expect(r'JPEG header parsed: width=\d+ height=\d+')
dut.expect_exact('Decoding JPEG -> RGB888...')
dut.expect(r'Decoded RGB888 size: \d+ bytes')
metadata_line = dut.expect(DECODE_INFO_PATTERN).group(0).decode('utf-8')
metadata = parse_decode_metadata(metadata_line)
dut.expect_exact('JPEG_DECODE_BASE64_BEGIN')
# Collect the machine-readable payload before the example prints its final
# completion line so we keep the UART parsing strictly in output order.
payload_lines = collect_base64_payload(dut)
dut.expect_exact('JPEG decode demo done.')
result_image = decode_base64_image(metadata, payload_lines)
output_path = Path(dut.logdir) / DECODE_OUTPUT_NAME
save_ppm_artifact(result_image, output_path)
assert_image_matches_golden(result_image, GOLDEN_OUTPUT_PATH)
@pytest.mark.generic
@idf_parametrize('target', soc_filtered_targets('SOC_JPEG_DECODE_SUPPORTED == 1'), indirect=['target'])
def test_jpeg_decode_example(dut: Dut) -> None:
run_jpeg_decode_example(dut)
@pytest.mark.flash_encryption
@pytest.mark.parametrize(
'config',
[
'flash_enc',
],
indirect=True,
)
@idf_parametrize(
'target',
soc_filtered_targets('SOC_JPEG_DECODE_SUPPORTED == 1 and SOC_FLASH_ENC_SUPPORTED == 1'),
indirect=['target'],
)
def test_jpeg_decode_example_with_flash_encryption(dut: Dut) -> None:
run_jpeg_decode_example(dut)