geointel/backend/tests/test_grayscale_yolo_dataset.py

from __future__ import annotations

import json
import subprocess
import sys
from hashlib import sha256
from pathlib import Path

from PIL import Image


SCRIPT = Path(__file__).parents[2] / "scripts" / "build_grayscale_yolo_dataset.py"


def test_grayscale_builder_preserves_labels_and_split(tmp_path: Path) -> None:
    source = tmp_path / "source"
    entries = []
    for split, sample_slug, colour in (("train", "fixture-train", (255, 0, 0)), ("val", "fixture-val", (0, 255, 0))):
        image = source / "images" / split / f"{sample_slug}.png"
        label = source / "labels" / split / f"{sample_slug}.txt"
        image.parent.mkdir(parents=True, exist_ok=True)
        label.parent.mkdir(parents=True, exist_ok=True)
        Image.new("RGB", (8, 8), colour).save(image)
        label.write_text("0 0.5 0.5 0.5 0.5\n", encoding="utf-8")
        entries.append((split, sample_slug, image, label))
    manifest = source / "operator_samples_manifest.json"
    policy = "geointel-training-source-eligibility/v1"
    manifest.write_text(
        json.dumps(
            {
                "immutable": True,
                "training_eligibility": {"policy_version": policy, "status": "eligible", "fixture_mode": True},
                "samples": [
                    {
                        "sample_slug": sample_slug,
                        "split": split,
                        "raster_dataset_id": f"raster:{sample_slug}",
                        "reference_dataset_id": f"reference:{sample_slug}",
                        "training_eligibility": {
                            "policy_version": policy,
                            "eligible": True,
                            "fixture_mode": True,
                            "raster": {"eligible": True, "reasons": [], "evidence": {"dataset_id": f"raster:{sample_slug}", "checksum_sha256": "a" * 64, "source_registry_id": "fixture-raster", "source_snapshot_id": "fixture-raster-snapshot"}},
                            "reference": {"eligible": True, "reasons": [], "evidence": {"dataset_id": f"reference:{sample_slug}", "checksum_sha256": "b" * 64, "source_registry_id": "fixture-reference", "source_snapshot_id": "fixture-reference-snapshot"}},
                        },
                    }
                    for split, sample_slug, _image, _label in entries
                ],
            }
        ),
        encoding="utf-8",
    )
    (source / "corpus-freeze.json").write_text(
        json.dumps(
            {
                "schema_version": 2,
                "immutable": True,
                "fixture_mode": True,
                "training_eligibility_policy": policy,
                "manifest_sha256": sha256(manifest.read_bytes()).hexdigest(),
            }
        ),
        encoding="utf-8",
    )
    dataset_yaml = source / "dataset.yaml"
    dataset_yaml.write_text(
        f"path: {source}\ntrain: images/train\nval: images/val\nnames:\n  0: building\n",
        encoding="utf-8",
    )
    release_script = SCRIPT.parent / "training_release_manifest.py"
    subprocess.run(
        [
            sys.executable,
            str(release_script),
            "create",
            "--train-yaml",
            str(dataset_yaml),
            "--corpus-manifest",
            str(manifest),
            "--fixture-mode",
        ],
        check=True,
    )
    release_path = dataset_yaml.with_name(dataset_yaml.name + ".geointel-training-release.json")
    asset_path = dataset_yaml.with_name(dataset_yaml.name + ".geointel-training-assets.json")
    assets = json.loads(asset_path.read_text(encoding="utf-8"))
    summary = source / "yolo_tile_dataset_summary.json"
    summary.write_text(
        json.dumps(
            {
                "dataset_yaml": str(dataset_yaml.resolve()),
                "training_release_manifest": str(release_path.resolve()),
                "training_release_manifest_sha256": sha256(release_path.read_bytes()).hexdigest(),
                "training_asset_manifest": str(asset_path.resolve()),
                "source_manifest_sha256": sha256(manifest.read_bytes()).hexdigest(),
                "tiles": [
                    {"split": entry["split"], "image_path": entry["image_path"], "label_path": entry["label_path"]}
                    for entry in assets["entries"]
                ],
            }
        ),
        encoding="utf-8",
    )
    output = tmp_path / "gray"
    subprocess.run(
        [
            sys.executable,
            str(SCRIPT),
            "--summary",
            str(summary),
            "--train-yaml",
            str(dataset_yaml),
            "--corpus-manifest",
            str(manifest),
            "--fixture-mode",
            "--output-dir",
            str(output),
        ],
        check=True,
    )
    converted = Image.open(output / "images" / "train" / "fixture-train.png")
    r, g, b = converted.getpixel((0, 0))
    assert r == g == b
    assert (output / "labels" / "train" / "fixture-train.txt").read_text() == "0 0.5 0.5 0.5 0.5\n"
    evidence = json.loads((output / "grayscale-dataset-evidence.json").read_text())
    assert evidence["converted_tile_count"] == 2
    assert evidence["training_eligible"] is False