feat(avatar): add SVG, PNG and HDR avatar exports

tests/test_hdr.py

@@ -0,0 +1,93 @@
+"""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()