94 lines
4.4 KiB
Python
94 lines
4.4 KiB
Python
"""Check physical luminance, PQ encoding, and the actual PNG byte stream."""
|
|
|
|
import struct
|
|
import tempfile
|
|
from pathlib import Path
|
|
import unittest
|
|
import xml.etree.ElementTree as ET
|
|
import zlib
|
|
|
|
import numpy as np
|
|
|
|
from avatar_color import BT2020_Y, SRGB_Y, SDR_PEAK_Y, TO_2020, equal_luminance, gradient_level, linear_srgb, sdr_color
|
|
from avatar_hdr import export_hdr, pq_decode, pq_encode, render_linear_hdr
|
|
from cmatrix_avatar import BLUE, GREEN, YELLOW, CMatrixAvatar
|
|
|
|
|
|
def read_png(path):
|
|
data = Path(path).read_bytes()
|
|
assert data[:8] == b'\x89PNG\r\n\x1a\n'
|
|
position, chunks = 8, {}
|
|
while position < len(data):
|
|
size = struct.unpack_from('>I', data, position)[0]
|
|
kind = data[position + 4:position + 8]
|
|
payload = data[position + 8:position + 8 + size]
|
|
crc = struct.unpack_from('>I', data, position + 8 + size)[0]
|
|
assert crc == zlib.crc32(kind + payload) & 0xffffffff
|
|
chunks[kind] = chunks.get(kind, b'') + payload
|
|
position += 12 + size
|
|
return chunks
|
|
|
|
|
|
class HDRTests(unittest.TestCase):
|
|
def test_pq_reference_values_and_round_trip(self):
|
|
self.assertAlmostEqual(float(pq_encode(100)), .50807842, places=7)
|
|
self.assertAlmostEqual(float(pq_encode(1000)), .75182710, places=7)
|
|
self.assertAlmostEqual(float(pq_encode(10000)), 1, places=7)
|
|
nits = np.array([.02, 2, 100, 400, 1000, 4000])
|
|
np.testing.assert_allclose(pq_decode(pq_encode(nits)), nits, rtol=1e-10)
|
|
|
|
def test_equal_peaks_and_stronger_gradient(self):
|
|
self.assertEqual(sdr_color(BLUE), BLUE)
|
|
for base in (BLUE, GREEN, YELLOW):
|
|
rgb = equal_luminance(base, 400, 10000, hdr=True)
|
|
self.assertAlmostEqual(sum(a * b for a, b in zip(rgb, BT2020_Y)), 400)
|
|
reference = np.asarray(TO_2020) @ np.asarray(linear_srgb(base))
|
|
np.testing.assert_allclose(np.asarray(rgb) / sum(rgb), reference / reference.sum())
|
|
self.assertLessEqual(max(equal_luminance(base, 2000, 10000, hdr=True)), 10000)
|
|
sdr_y = sum(a * b for a, b in zip(linear_srgb(sdr_color(base)), SRGB_Y))
|
|
self.assertAlmostEqual(sdr_y, SDR_PEAK_Y, delta=.004)
|
|
levels = [gradient_level(i / 31) for i in range(32)]
|
|
self.assertEqual(len(set(levels)), 32)
|
|
self.assertEqual(levels, sorted(levels))
|
|
self.assertEqual(levels[0], .005)
|
|
self.assertEqual(levels[-1], 1)
|
|
avatar = CMatrixAvatar(42)
|
|
yellow = [c.level for row in avatar.cells for c in row if c.color == YELLOW]
|
|
self.assertGreater(len(set(yellow)), 5)
|
|
|
|
def test_hdr_file_encoding_and_decoded_luminance(self):
|
|
svg = '<svg xmlns="http://www.w3.org/2000/svg" width="128" height="16" viewBox="0 0 128 16">'
|
|
for index, (base, level) in enumerate(((GREEN, 1), (BLUE, 1), (YELLOW, 1), (GREEN, .005))):
|
|
x = index * 32
|
|
svg += f'<path data-base="{base}" data-level="{level}" d="M{x} 0h32v16h-32z"/>'
|
|
svg += '</svg>'
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
path = Path(directory) / 'hdr.png'
|
|
export_hdr(svg, path, 400)
|
|
chunks = read_png(path)
|
|
self.assertEqual(struct.unpack('>IIBBBBB', chunks[b'IHDR']), (128, 16, 16, 2, 0, 0, 0))
|
|
self.assertEqual(chunks[b'cICP'], bytes((9, 16, 0, 1)))
|
|
self.assertNotIn(b'sRGB', chunks)
|
|
self.assertNotIn(b'gAMA', chunks)
|
|
raw = np.frombuffer(zlib.decompress(chunks[b'IDAT']), dtype=np.uint8).reshape(16, 128 * 6 + 1)
|
|
self.assertTrue((raw[:, 0] == 0).all())
|
|
signal = np.frombuffer(raw[:, 1:].copy().tobytes(), dtype='>u2').reshape(16, 128, 3)
|
|
self.assertTrue((signal % 257 != 0).any())
|
|
decoded = pq_decode(signal.astype(float) / 65535)
|
|
y = decoded @ np.asarray(BT2020_Y)
|
|
for column in (16, 48, 80):
|
|
self.assertAlmostEqual(y[8, column], 400, delta=.1)
|
|
self.assertAlmostEqual(y[8, 112], 2, delta=.01)
|
|
max_cll, max_fall = struct.unpack('>II', chunks[b'cLLI'])
|
|
self.assertAlmostEqual(max_cll / 10000, y.max(), delta=.0001)
|
|
self.assertAlmostEqual(max_fall / 10000, y.mean(), delta=.0001)
|
|
|
|
def test_invalid_hdr_peak_is_rejected(self):
|
|
for value in (0, 3000, float('nan'), float('inf')):
|
|
with self.assertRaises(ValueError):
|
|
render_linear_hdr('<svg/>', value)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|