| from __future__ import annotations |
|
|
| import io |
| import json |
| import pickle |
| import sys |
| import tarfile |
| from pathlib import Path |
|
|
| import numpy as np |
| from PIL import Image |
|
|
|
|
| REPO_ROOT = Path(__file__).resolve().parents[1] |
| STAGES = REPO_ROOT / "scripts" / "stages" |
| if str(STAGES) not in sys.path: |
| sys.path.insert(0, str(STAGES)) |
|
|
| from evaluate_mlperf_image_quality import ( |
| canonical_digest, |
| load_checkpoint, |
| load_cifar10_samples, |
| load_vww_samples, |
| ) |
|
|
|
|
| def add_bytes(archive: tarfile.TarFile, name: str, payload: bytes) -> None: |
| member = tarfile.TarInfo(name) |
| member.size = len(payload) |
| archive.addfile(member, io.BytesIO(payload)) |
|
|
|
|
| def test_cifar_loader_applies_pinned_indices_and_planar_to_nhwc(tmp_path: Path) -> None: |
| data = np.zeros((10000, 3072), dtype=np.uint8) |
| labels = [0] * 10000 |
| labels[17] = 3 |
| labels[991] = 8 |
| data[17, :1024] = 11 |
| data[17, 1024:2048] = 22 |
| data[17, 2048:] = 33 |
| data[991, :1024] = 44 |
| archive_path = tmp_path / "cifar.tar.gz" |
| with tarfile.open(archive_path, "w:gz") as archive: |
| add_bytes( |
| archive, |
| "cifar-10-batches-py/test_batch", |
| pickle.dumps({b"data": data, b"labels": labels}), |
| ) |
| indices_path = tmp_path / "indices.npy" |
| np.save(indices_path, np.asarray([17, 991], dtype=np.int64), allow_pickle=False) |
| labels_path = tmp_path / "labels.csv" |
| labels_path.write_text("first.bin,10,3\nsecond.bin,10,8\n") |
|
|
| samples = load_cifar10_samples(archive_path, labels_path, indices_path) |
|
|
| assert [sample.label for sample in samples] == [3, 8] |
| assert [sample.sample_id for sample in samples] == ["cifar10_test_00017", "cifar10_test_00991"] |
| assert samples[0].image.shape == (32, 32, 3) |
| assert samples[0].image[0, 0].tolist() == [11, 22, 33] |
| assert samples[1].image[0, 0].tolist() == [44, 0, 0] |
|
|
|
|
| def test_vww_loader_matches_coco_ids_and_decodes_rgb(tmp_path: Path) -> None: |
| archive_path = tmp_path / "vww.tar.gz" |
| expected = { |
| "000000000123": np.full((96, 96, 3), [255, 0, 0], dtype=np.uint8), |
| "000000999999": np.full((96, 96, 3), [0, 255, 0], dtype=np.uint8), |
| } |
| with tarfile.open(archive_path, "w:gz") as archive: |
| for image_id, array in expected.items(): |
| payload = io.BytesIO() |
| Image.fromarray(array).save(payload, format="PNG") |
| add_bytes(archive, f"vw_coco2014_96/class/COCO_val2014_{image_id}.png", payload.getvalue()) |
| labels_path = tmp_path / "labels.csv" |
| labels_path.write_text("000000000123.bin,2,1\n000000999999.bin,2,0\n") |
|
|
| samples = load_vww_samples(archive_path, labels_path) |
|
|
| assert [sample.sample_id for sample in samples] == ["coco_000000000123", "coco_000000999999"] |
| assert [sample.label for sample in samples] == [1, 0] |
| assert np.array_equal(samples[0].image, expected["000000000123"]) |
| assert np.array_equal(samples[1].image, expected["000000999999"]) |
|
|
|
|
| def test_checkpoint_requires_same_canonical_config(tmp_path: Path) -> None: |
| digest = canonical_digest({"model_id": "VC02", "threshold": 0.85}) |
| checkpoint = tmp_path / "checkpoint.jsonl" |
| checkpoint.write_text( |
| json.dumps({"config_sha256": digest, "sample_id": "one", "label": 1}) + "\n" |
| ) |
| assert load_checkpoint(checkpoint, digest)["one"]["label"] == 1 |
|
|
| try: |
| load_checkpoint(checkpoint, "0" * 64) |
| except ValueError as error: |
| assert "config mismatch" in str(error) |
| else: |
| raise AssertionError("checkpoint from a different config was accepted") |
|
|
