zjs81/meshcore-open
1
0
Fork 0
mirror of https://github.com/zjs81/meshcore-open.git synced 2026-08-11 18:26:27 +10:00
Code Issues Packages Projects Releases 2 Wiki Activity
Files
c78c77309e366dea05f528ba6cd7b6e40df5bdd6
meshcore-open/test/services/image_codec_entropy_test.dart
T
zjs81 cf6c93c50a rebuild translations
2026-08-11 00:15:45 -07:00

1026 lines
29 KiB
Dart
Raw Blame History

import 'dart:io';
import 'dart:math' as math;
import 'dart:typed_data';
import 'package:flutter_test/flutter_test.dart';
import 'package:meshcore_open/services/entropy_tables.dart';
import 'package:meshcore_open/services/image_codec_backend.dart'
show AeicRansCoders;
import 'package:meshcore_open/services/image_codec_entropy.dart';
/// Golden vectors for the entropy layer — the arithmetic between the ONNX
/// tensors and the rANS coder.
///
/// GENERATED BY: `aic/exp/export_entropy_layer_golden.py`, which runs the real
/// `torch` ops from `aic/aeic/src/codec/codec_practical.py`
/// (`get_mask_four_parts`, `sequeeze`, `torch.round`, `my_build_indexes`) on
/// inputs built from a closed-form integer recipe. Every constant in that
/// recipe is a power-of-two fraction or a plain float64 division, so Dart and
/// numpy reproduce the float32 inputs bit-for-bit and the comparison is real
/// rather than a re-implementation checking itself.
///
/// Regenerate with:
/// cd /Users/Zach/Documents/mycode/aic
/// AEIC_DEVICE=cpu .venv/bin/python exp/export_entropy_layer_golden.py
///
/// WHY THIS MATTERS: none of these failures are loud. A wrong mask permutation,
/// a `sequeeze` that folds the wrong channels, or a rounding tie resolved away
/// from zero instead of to even does not throw — it desynchronises rANS and
/// produces a sharp, plausible, wrong image.
class _Golden {
// --- small case: C = 8, H = 4, W = 4, squeezed length 32 ---
/// `mask_i` flattened over `[1, 8, 4, 4]`, '1' where live. Straight from
/// `get_mask_four_parts(1, 8, 4, 4)`.
static const List<String> masks = <String>[
'10100000101000001010000010100000010100000101000001010000010100000000101000001010000010100000101000000101000001010000010100000101',
'00000101000001010000010100000101000010100000101000001010000010100101000001010000010100000101000010100000101000001010000010100000',
'00001010000010100000101000001010000001010000010100000101000001011010000010100000101000001010000001010000010100000101000001010000',
'01010000010100000101000001010000101000001010000010100000101000000000010100000101000001010000010100001010000010100000101000001010',
];
static const List<List<int>> symbols = <List<int>>[
<int>[
-1,
-1,
-2,
4,
0,
0,
6,
-3,
2,
1,
1,
-1,
-5,
1,
0,
0,
-1,
-1,
-2,
4,
0,
0,
5,
-3,
2,
1,
1,
-1,
-5,
1,
0,
0,
],
<int>[
-1,
-1,
4,
4,
0,
0,
0,
5,
2,
1,
-1,
-1,
3,
-5,
1,
0,
-1,
-2,
4,
4,
0,
0,
0,
5,
1,
1,
-1,
-1,
3,
-5,
1,
0,
],
<int>[
-1,
-2,
-2,
4,
1,
0,
0,
5,
2,
0,
-1,
-1,
3,
-5,
1,
0,
-1,
-2,
-2,
4,
0,
0,
0,
5,
2,
-1,
-1,
-1,
3,
-5,
1,
0,
],
<int>[
-1,
-1,
-2,
-2,
0,
0,
5,
-3,
2,
1,
-1,
-1,
-5,
-5,
0,
0,
-1,
-1,
-2,
-2,
0,
0,
5,
-3,
2,
1,
-1,
-1,
-5,
1,
0,
0,
],
];
static const List<List<int>> indexes = <List<int>>[
<int>[
-1,
6,
-1,
8,
12,
15,
13,
16,
0,
9,
3,
10,
13,
16,
14,
17,
0,
9,
3,
10,
13,
16,
14,
17,
5,
11,
7,
12,
15,
17,
15,
18,
],
<int>[
14,
11,
15,
12,
7,
0,
8,
2,
15,
13,
16,
14,
9,
4,
11,
6,
15,
13,
16,
14,
10,
4,
11,
6,
16,
14,
17,
15,
11,
7,
12,
9,
],
<int>[
11,
14,
12,
15,
-1,
8,
0,
9,
12,
16,
13,
16,
2,
10,
5,
11,
13,
16,
13,
16,
3,
10,
5,
11,
14,
17,
15,
17,
7,
12,
8,
13,
],
<int>[
5,
-1,
7,
0,
15,
12,
15,
13,
8,
1,
10,
4,
16,
14,
16,
14,
8,
2,
10,
4,
16,
14,
16,
14,
11,
6,
12,
8,
17,
15,
17,
16,
],
];
// --- full case: C = 256, H = W = 16 (the shipping geometry) ---
//
// 16,384 symbols per stage is too much to embed, so the golden is a
// fingerprint: total, extremes, and the first and last twelve values. Any
// permutation error moves the head or the tail; any arithmetic error moves
// the sum.
static const List<int> fullSymbolSum = <int>[-17, 40, -10, 29];
static const List<int> fullIndexSum = <int>[412240, 412285, 412231, 412219];
static const List<List<int>> fullSymbolHead = <List<int>>[
<int>[-1, -4, -2, 1, 4, -1, -5, -2, 0, 3, -1, -5],
<int>[-3, 0, 3, 5, 0, -4, -1, 2, 4, 1, -4, -2],
<int>[0, 3, -1, 2, -3, 0, 2, -1, 1, -4, -1, 2],
<int>[-3, -1, 2, -2, 1, 3, -2, 1, 4, 0, 3, 5],
];
static const List<List<int>> fullSymbolTail = <List<int>>[
<int>[-4, -2, 1, 4, -2, 1, -3, 0, 3, -2, -6, -3],
<int>[-2, 0, -3, -1, 2, 5, 0, -4, -1, 1, 4, 6],
<int>[-5, -2, 0, 3, -1, -6, -3, 0, 2, 5, 1, -4],
<int>[-1, 1, 4, 0, -4, -2, 1, 3, 6, 2, -2, 0],
];
static const List<List<int>> fullIndexHead = <List<int>>[
<int>[-1, -1, -1, 1, 0, 4, 4, 6, 6, 8, 8, 9],
<int>[0, 0, 4, 3, 6, 6, 8, 7, 9, 9, 10, 10],
<int>[0, 2, 2, 5, 4, 7, 7, 8, 8, 10, 10, 11],
<int>[-1, -1, 0, 0, 3, 2, 5, 5, 7, 7, 9, 8],
];
static const List<List<int>> fullIndexTail = <List<int>>[
<int>[16, 17, 17, 18, 17, 18, 18, 18, 18, 19, 19, 19],
<int>[16, 16, 17, 16, 17, 17, 18, 18, 18, 18, 19, 18],
<int>[16, 16, 16, 17, 17, 17, 17, 18, 18, 18, 18, 19],
<int>[17, 17, 17, 17, 18, 18, 18, 18, 19, 19, 19, 19],
];
/// `my_build_indexes` on a spread of scales, isolated from everything else.
static const List<double> probeScales = <double>[
0.0,
9.999999974752427e-07,
0.05000000074505806,
0.07989999651908875,
0.07999999821186066,
0.10999999940395355,
0.11000010371208191,
0.5,
1.0,
2.0,
7.5,
63.900001525878906,
255.89999389648438,
256.0,
10000.0,
];
static const List<int> probeIndexes = <int>[
-1,
-1,
-1,
-1,
0,
0,
0,
12,
17,
23,
34,
51,
62,
63,
63,
];
}
/// The same closed-form recipe the generator uses, in Dart.
///
/// `y`, `means_supp` and `scales_supp` for a `[1, C, H, W]` tensor. Kept
/// byte-identical to `recipe()` in `export_entropy_layer_golden.py`.
({Float32List y, Float32List means, Float32List scales}) _recipe(
int channels,
int height,
int width,
) {
final n = channels * height * width;
final y = Float32List(n);
final means = Float32List(n);
final scales = Float32List(n);
var i = 0;
for (var c = 0; c < channels; c++) {
for (var h = 0; h < height; h++) {
for (var w = 0; w < width; w++, i++) {
y[i] = ((c * 3 + h * 17 + w * 11) % 61 - 30) * 0.125;
means[i] = ((c * 7 + h * 13 + w * 29) % 97 - 48) * 0.0625;
scales[i] = ((c * 11 + h * 5 + w * 3) % 700) * 0.01 + 0.001;
}
}
}
return (y: y, means: means, scales: scales);
}
/// One stage of `compress()`: mask, quantize, fold, index.
({Int16List symbols, Int16List indexes}) _runStage(
AeicMaskSet masks,
Float32List y,
Float32List meansSupp,
Float32List scalesSupp,
int stage,
) {
final means = masks.applyMask(meansSupp, stage);
final scales = masks.applyMask(scalesSupp, stage);
final maskedY = masks.applyMask(y, stage);
final yq = Float32List(y.length);
for (var i = 0; i < y.length; i++) {
yq[i] = roundHalfToEven(f32(maskedY[i] - means[i]));
}
return (
symbols: aeicToSymbols(masks.squeeze(yq)),
indexes: aeicBuildIndexes(masks.squeeze(scales)),
);
}
void main() {
group('AeicEntropyGeometry', () {
test('512x512 ft32 matches the shapes the bitstream format assumes', () {
final g = AeicEntropyGeometry.forResolution(512);
expect(g.yShape, <int>[1, 256, 16, 16]);
expect(g.zShape, <int>[1, 128, 4, 4]);
expect(g.squeezedChannels, 64);
// From aic/results/rans_port_spec.md §1.
expect(g.zElements, 2048);
expect(g.symbolsPerStage, 16384);
expect(g.totalEntries, 67584);
});
test('z is ceil(y/4), not floor', () {
// 320/32 = 10 -> z must be 3, matching compress()'s reflect padding.
final g = AeicEntropyGeometry.forResolution(320);
expect(g.yHeight, 10);
expect(g.zHeight, 3);
});
test('rejects a resolution g_a cannot downsample by 32', () {
expect(() => AeicEntropyGeometry.forResolution(500), throwsArgumentError);
expect(() => AeicEntropyGeometry.forResolution(0), throwsArgumentError);
});
});
group('AeicMaskSet', () {
test('the four masks are get_mask_four_parts, element for element', () {
final masks = AeicMaskSet(
AeicEntropyGeometry.forResolution(128, yChannels: 8),
);
// 128/32 = 4, so this is exactly the C=8 H=W=4 case in the golden.
for (var stage = 0; stage < 4; stage++) {
final tensor = masks.maskTensor(stage);
final actual = tensor.map((v) => v == 1.0 ? '1' : '0').join();
expect(actual, _Golden.masks[stage], reason: 'mask_$stage');
}
});
test('every position is claimed by exactly one stage, per channel', () {
final geometry = AeicEntropyGeometry.forResolution(512);
final masks = AeicMaskSet(geometry);
final counts = Uint8List(geometry.yElements);
for (var stage = 0; stage < 4; stage++) {
final tensor = masks.maskTensor(stage);
for (var i = 0; i < tensor.length; i++) {
counts[i] += tensor[i].toInt();
}
}
expect(counts.every((c) => c == 1), isTrue);
});
test('each mask carries exactly a quarter of the tensor', () {
final geometry = AeicEntropyGeometry.forResolution(512);
final masks = AeicMaskSet(geometry);
for (var stage = 0; stage < 4; stage++) {
final live = masks.maskTensor(stage).where((v) => v == 1.0).length;
expect(live, geometry.symbolsPerStage);
}
});
test('squeeze then unsqueeze is the identity on a masked tensor', () {
final geometry = AeicEntropyGeometry.forResolution(512);
final masks = AeicMaskSet(geometry);
final recipe = _recipe(256, 16, 16);
for (var stage = 0; stage < 4; stage++) {
final masked = masks.applyMask(recipe.y, stage);
final restored = masks.unsqueeze(masks.squeeze(masked), stage);
expect(restored, masked, reason: 'stage $stage');
}
});
test('mergeContext replaces exactly the stage mask', () {
final geometry = AeicEntropyGeometry.forResolution(512);
final masks = AeicMaskSet(geometry);
final base = _recipe(256, 16, 16).y;
final stageLatent = masks.applyMask(_recipe(256, 16, 16).means, 2);
final merged = masks.mergeContext(base, stageLatent, 2);
final mask = masks.maskTensor(2);
for (var i = 0; i < merged.length; i++) {
expect(merged[i], mask[i] == 1.0 ? stageLatent[i] : base[i]);
}
});
});
group('roundHalfToEven', () {
test('ties go to even, unlike Dart round()', () {
expect(roundHalfToEven(0.5), 0.0);
expect(roundHalfToEven(1.5), 2.0);
expect(roundHalfToEven(2.5), 2.0);
expect(roundHalfToEven(-0.5), 0.0);
expect(roundHalfToEven(-1.5), -2.0);
expect(roundHalfToEven(-2.5), -2.0);
// Dart disagrees on every one of those ties.
expect((-0.5).roundToDouble(), -1.0);
});
test('non-ties are ordinary rounding', () {
expect(roundHalfToEven(0.49), 0.0);
expect(roundHalfToEven(0.51), 1.0);
expect(roundHalfToEven(-1.51), -2.0);
expect(roundHalfToEven(-1.49), -1.0);
expect(roundHalfToEven(7.0), 7.0);
});
});
group('aeicBuildIndexes', () {
test('matches my_build_indexes on the probe scales', () {
final scales = Float32List.fromList(_Golden.probeScales);
expect(aeicBuildIndexes(scales), _Golden.probeIndexes);
});
test('the 0.08 threshold is strict, and clamps do not leak', () {
// 0.08 itself is NOT skipped; it clamps to row 0 because ln(0.08) is
// below ln(0.11).
expect(aeicBuildIndexes(Float32List.fromList(<double>[0.08])).first, 0);
expect(
aeicBuildIndexes(Float32List.fromList(<double>[0.0799])).first,
-1,
);
expect(
aeicBuildIndexes(Float32List.fromList(<double>[1e9])).first,
kAeicScalesLevels - 1,
);
});
});
group('aeicZIndexes', () {
test('is the channel arange broadcast over H*W', () {
final geometry = AeicEntropyGeometry.forResolution(512);
final indexes = aeicZIndexes(geometry);
expect(indexes.length, 2048);
expect(indexes.first, 0);
expect(indexes[15], 0);
expect(indexes[16], 1);
expect(indexes.last, 127);
// Verified against results/golden/vectors/kodim01.gv, which stores the
// exact int16 array the C++ coder was given.
for (var i = 0; i < indexes.length; i++) {
expect(indexes[i], i ~/ 16);
}
});
});
group('aeicToSymbols', () {
test('rejects a value the int16 wire format cannot carry', () {
expect(
() => aeicToSymbols(Float32List.fromList(<double>[40000.0])),
throwsStateError,
);
expect(
() => aeicToSymbols(Float32List.fromList(<double>[1.5])),
throwsStateError,
);
expect(aeicToSymbols(Float32List.fromList(<double>[-3.0])), <int>[-3]);
});
});
group('aeicRgbToChw', () {
test('reproduces ToTensor + Normalize([0.5], [0.5])', () {
final chw = aeicRgbToChw(Uint8List.fromList(<int>[0, 128, 255]), 1);
expect(chw.length, 3);
expect(chw[0], -1.0);
expect(chw[1], closeTo(0.00392, 1e-4));
expect(chw[2], 1.0);
});
test('is channel-planar, not interleaved', () {
final rgb = Uint8List(4 * 3);
for (var i = 0; i < 4; i++) {
rgb[i * 3] = 255; // R
rgb[i * 3 + 1] = 0; // G
rgb[i * 3 + 2] = 0; // B
}
final chw = aeicRgbToChw(rgb, 2);
expect(chw.sublist(0, 4), <double>[1.0, 1.0, 1.0, 1.0]);
expect(chw.sublist(4, 8), <double>[-1.0, -1.0, -1.0, -1.0]);
});
test('rejects a byte count that is not the stated square', () {
expect(() => aeicRgbToChw(Uint8List(11), 2), throwsArgumentError);
});
});
group('four-stage symbol packing (golden)', () {
test('C=8 H=4 W=4: symbols and indexes match torch exactly', () {
final geometry = AeicEntropyGeometry.forResolution(128, yChannels: 8);
final masks = AeicMaskSet(geometry);
final recipe = _recipe(8, 4, 4);
for (var stage = 0; stage < 4; stage++) {
final out = _runStage(
masks,
recipe.y,
recipe.means,
recipe.scales,
stage,
);
expect(out.symbols, _Golden.symbols[stage], reason: 'symbols $stage');
expect(out.indexes, _Golden.indexes[stage], reason: 'indexes $stage');
}
});
test('C=256 H=W=16 (shipping geometry): fingerprint matches torch', () {
final geometry = AeicEntropyGeometry.forResolution(512);
final masks = AeicMaskSet(geometry);
final recipe = _recipe(256, 16, 16);
for (var stage = 0; stage < 4; stage++) {
final out = _runStage(
masks,
recipe.y,
recipe.means,
recipe.scales,
stage,
);
expect(out.symbols.length, 16384);
expect(
out.symbols.fold<int>(0, (a, b) => a + b),
_Golden.fullSymbolSum[stage],
reason: 'symbol sum $stage',
);
expect(out.symbols.sublist(0, 12), _Golden.fullSymbolHead[stage]);
expect(out.symbols.sublist(16372), _Golden.fullSymbolTail[stage]);
expect(
out.indexes.fold<int>(0, (a, b) => a + b),
_Golden.fullIndexSum[stage],
reason: 'index sum $stage',
);
expect(out.indexes.sublist(0, 12), _Golden.fullIndexHead[stage]);
expect(out.indexes.sublist(16372), _Golden.fullIndexTail[stage]);
}
});
});
group('AeicEntropyCodec', () {
test(
'encode pushes z, y0, y1, y2, y3 in that order and no other',
() async {
final geometry = AeicEntropyGeometry.forResolution(512);
final network = _FakeNetwork(geometry);
final coders = _RecordingCoders();
final codec = AeicEntropyCodec(
geometry: geometry,
network: network,
coders: coders,
);
final progress = <double>[];
final stream = await codec.encode(
Uint8List(512 * 512 * 3),
onProgress: progress.add,
);
expect(coders.encoder.groups, <int>[
kAeicZCdfGroup,
kAeicYCdfGroup,
kAeicYCdfGroup,
kAeicYCdfGroup,
kAeicYCdfGroup,
]);
expect(coders.encoder.lengths, <int>[2048, 16384, 16384, 16384, 16384]);
expect(stream, isNotEmpty);
expect(progress.last, 1.0);
// The image really was normalized before the graph saw it: 0 -> -1.
expect(network.lastInput!.first, -1.0);
},
);
test('encode refuses a graph that returns the wrong z size', () async {
final geometry = AeicEntropyGeometry.forResolution(512);
final codec = AeicEntropyCodec(
geometry: geometry,
network: _FakeNetwork(geometry, zElements: 7),
coders: _RecordingCoders(),
);
await expectLater(
codec.encode(Uint8List(512 * 512 * 3)),
throwsStateError,
);
});
test('encode honours shouldCancel between stages', () async {
final geometry = AeicEntropyGeometry.forResolution(512);
final codec = AeicEntropyCodec(
geometry: geometry,
network: _FakeNetwork(geometry),
coders: _RecordingCoders(),
);
await expectLater(
codec.encode(Uint8List(512 * 512 * 3), shouldCancel: () => true),
throwsA(isA<AeicEntropyCancelled>()),
);
});
test('decode reports a send-side-only graph instead of guessing', () async {
final geometry = AeicEntropyGeometry.forResolution(512);
final codec = AeicEntropyCodec(
geometry: geometry,
network: _FakeNetwork(geometry, decodeSide: false),
coders: _RecordingCoders(),
);
await expectLater(
codec.decodeToLatent(Uint8List(16)),
throwsA(isA<AeicEntropyUnavailable>()),
);
});
test('decode walks the stages in order and rebuilds y_hat', () async {
final geometry = AeicEntropyGeometry.forResolution(512);
final network = _FakeNetwork(geometry);
final coders = _RecordingCoders();
final codec = AeicEntropyCodec(
geometry: geometry,
network: network,
coders: coders,
);
final yHat = await codec.decodeToLatent(Uint8List(16));
expect(network.stageCalls, <int>[0, 1, 2, 3]);
expect(coders.decoder.groups, <int>[
kAeicZCdfGroup,
kAeicYCdfGroup,
kAeicYCdfGroup,
kAeicYCdfGroup,
kAeicYCdfGroup,
]);
expect(yHat.length, geometry.yElements);
// The fake decoder returns symbol 1 everywhere and the fake network
// returns means 0, so every position of y_hat must be exactly 1 — which
// only holds if the four masks tile the tensor and each stage's squeezed
// symbols were unsqueezed back into the right channel group.
expect(yHat.every((v) => v == 1.0), isTrue);
});
});
// The end-to-end proof: the entropy layer plus the pure-Dart range coder
// reproduce, byte for byte, the bitstreams the C++ coder produced for real
// images — and decode them back to the same symbols.
group('bitstream round trip against the C++ golden vectors', () {
final Directory goldenDir = _resolveGoldenDir();
final EntropyTables tables = EntropyTables.parse(
File('${goldenDir.path}/aeic_cdf_ft32.bin').readAsBytesSync(),
);
final geometry = AeicEntropyGeometry.forResolution(512);
final masks = AeicMaskSet(geometry);
// kodim01 is the plain case; kodim23 and image2 both carry index -1
// (scales below 0.08), which the coder must skip on encode and read back as
// a literal 0.
for (final name in <String>['kodim01', 'kodim23', 'image2']) {
test('$name: symbols -> bitstream -> symbols', () async {
final vector = _readGoldenVector(
File('${goldenDir.path}/vectors/$name.gv').readAsBytesSync(),
);
final expectedStream = File(
'${goldenDir.path}/vectors/$name.bin',
).readAsBytesSync();
// Rebuild the pre-fold tensors the graph would have produced. `squeeze`
// is a bijection on masked tensors, so unsqueezing the golden arrays
// recovers a legitimate input and the loop's own fold has to invert it.
final yQ = <Float32List>[];
final scales = <Float32List>[];
for (var stage = 0; stage < 4; stage++) {
final symbols = vector['y_q$stage']!;
final indexes = vector['y_indexes$stage']!;
final squeezedY = Float32List(symbols.length);
final squeezedS = Float32List(indexes.length);
for (var i = 0; i < symbols.length; i++) {
squeezedY[i] = symbols[i].toDouble();
squeezedS[i] = _scaleForIndex(indexes[i]);
}
// Self-check: the synthesized scales must land back on the exact
// golden indexes, or this test is measuring the wrong thing.
expect(aeicBuildIndexes(squeezedS), indexes, reason: 'stage $stage');
yQ.add(masks.unsqueeze(squeezedY, stage));
scales.add(masks.unsqueeze(squeezedS, stage));
}
final zQ = Float32List(vector['z_q']!.length);
for (var i = 0; i < zQ.length; i++) {
zQ[i] = vector['z_q']![i].toDouble();
}
final network = _ReplayNetwork(
geometry: geometry,
tensors: AeicEncodeSideTensors(zQ: zQ, yQ: yQ, scales: scales),
);
final codec = AeicEntropyCodec(
geometry: geometry,
network: network,
coders: AeicRansCoders(tables),
);
// ENCODE: byte-for-byte against the C++ coder's output.
final stream = await codec.encode(Uint8List(512 * 512 * 3));
expect(stream, expectedStream, reason: 'bitstream for $name');
// DECODE: the same bytes back to the same symbols. means are zero, so
// y_hat is exactly the four stages' symbols tiled back into place —
// which only holds if every mask, fold and unfold agrees with encode.
final yHat = await codec.decodeToLatent(stream);
var expected = Float32List(geometry.yElements);
for (var stage = 0; stage < 4; stage++) {
expected = masks.mergeContext(expected, yQ[stage], stage);
}
expect(yHat, expected, reason: 'y_hat for $name');
expect(network.stageCalls, <int>[0, 1, 2, 3]);
});
}
});
}
Directory _resolveGoldenDir() {
for (final candidate in <String>[
'test/services/golden',
'../test/services/golden',
'golden',
]) {
final dir = Directory(candidate);
if (dir.existsSync()) {
return dir;
}
}
throw StateError(
'golden vectors not found; expected test/services/golden relative to the '
'package root',
);
}
/// Parses the `.gv` container documented in `aic/results/rans_port_spec.md` §10.
Map<String, Int16List> _readGoldenVector(Uint8List bytes) {
final data = ByteData.sublistView(bytes);
const magic = <int>[0x41, 0x45, 0x49, 0x43, 0x47, 0x56, 0x00, 0x01];
for (var i = 0; i < magic.length; i++) {
if (bytes[i] != magic[i]) {
throw StateError('not a .gv container');
}
}
final count = data.getUint32(12, Endian.little);
final names = <String>[];
final counts = <int>[];
final dtypes = <int>[];
var off = 16;
for (var i = 0; i < count; i++) {
final raw = bytes.sublist(off, off + 16);
final end = raw.indexOf(0);
names.add(String.fromCharCodes(raw.sublist(0, end < 0 ? 16 : end)));
dtypes.add(data.getUint32(off + 16, Endian.little));
counts.add(data.getUint32(off + 20, Endian.little));
off += 24;
}
final out = <String, Int16List>{};
for (var i = 0; i < count; i++) {
if (dtypes[i] != 0) {
throw StateError('${names[i]} is not int16');
}
final values = Int16List(counts[i]);
for (var j = 0; j < counts[i]; j++) {
values[j] = data.getInt16(off + j * 2, Endian.little);
}
off += counts[i] * 2;
out[names[i]] = values;
}
return out;
}
/// A scale that `my_build_indexes` maps back to exactly [index].
///
/// Row centres, so float32 rounding cannot push one over a boundary; index -1
/// means "skipped", which any scale below 0.08 produces.
double _scaleForIndex(int index) {
if (index < 0) {
return 0.0;
}
return f32(math.exp(kAeicLogScaleMin + (index + 0.5) * kAeicLogScaleStep));
}
/// Replays fixed tensors as if they came from the graph, for both directions.
class _ReplayNetwork implements AeicEntropyNetwork {
final AeicEntropyGeometry geometry;
final AeicEncodeSideTensors tensors;
final List<int> stageCalls = <int>[];
_ReplayNetwork({required this.geometry, required this.tensors});
@override
bool get supportsDecodeSide => true;
@override
Future<AeicEncodeSideTensors> runEncodeSide(Float32List imageChw) async =>
tensors;
@override
Future<Float32List> runHyperSynthesis(Float32List zQ) async {
// The real h_s consumes z_hat; here the only thing under test is that the
// decoded z symbols reach it. Assert that and hand back a zero context.
expect(zQ.length, geometry.zElements);
for (var i = 0; i < zQ.length; i++) {
expect(zQ[i], tensors.zQ[i], reason: 'z symbol $i');
}
return Float32List(geometry.yElements);
}
@override
Future<AeicStageParams> runStage(int stage, Float32List base) async {
stageCalls.add(stage);
return AeicStageParams(
meansSupp: Float32List(geometry.yElements),
scalesSupp: tensors.scales[stage],
);
}
}
/// A stand-in for the ONNX graph: shapes and call order, no arithmetic.
class _FakeNetwork implements AeicEntropyNetwork {
final AeicEntropyGeometry geometry;
final int? zElements;
final bool decodeSide;
final List<int> stageCalls = <int>[];
Float32List? lastInput;
_FakeNetwork(this.geometry, {this.zElements, this.decodeSide = true});
@override
bool get supportsDecodeSide => decodeSide;
@override
Future<AeicEncodeSideTensors> runEncodeSide(Float32List imageChw) async {
lastInput = imageChw;
return AeicEncodeSideTensors(
zQ: Float32List(zElements ?? geometry.zElements),
yQ: <Float32List>[
for (var i = 0; i < 4; i++) Float32List(geometry.yElements),
],
scales: <Float32List>[
for (var i = 0; i < 4; i++)
Float32List(geometry.yElements)
..fillRange(0, geometry.yElements, 1.0),
],
);
}
@override
Future<Float32List> runHyperSynthesis(Float32List zQ) async =>
Float32List(geometry.yElements);
@override
Future<AeicStageParams> runStage(int stage, Float32List base) async {
stageCalls.add(stage);
return AeicStageParams(
meansSupp: Float32List(geometry.yElements),
scalesSupp: Float32List(geometry.yElements)
..fillRange(0, geometry.yElements, 1.0),
);
}
}
class _RecordingEncoder implements AeicRansEncoder {
final List<int> groups = <int>[];
final List<int> lengths = <int>[];
@override
void pushSymbols(Int16List symbols, Int16List indexes, int cdfGroup) {
expect(symbols.length, indexes.length);
groups.add(cdfGroup);
lengths.add(symbols.length);
}
@override
Uint8List finish() => Uint8List.fromList(<int>[0x11, 4, 0]);
}
class _RecordingDecoder implements AeicRansDecoder {
final List<int> groups = <int>[];
@override
Int16List decodeStream(Int16List indexes, int cdfGroup) {
groups.add(cdfGroup);
return Int16List(indexes.length)..fillRange(0, indexes.length, 1);
}
}
class _RecordingCoders implements AeicRansCoderFactory {
final _RecordingEncoder encoder = _RecordingEncoder();
final _RecordingDecoder decoder = _RecordingDecoder();
@override
AeicRansEncoder createEncoder() => encoder;
@override
AeicRansDecoder createDecoder(Uint8List bitstream) => decoder;
}
Reference in New Issue View Git Blame Copy Permalink