File size: 3,559 Bytes
ed3aeeb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 | 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 ( # noqa: E402
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")
|