File size: 2,071 Bytes
953508f ae754fb 953508f ae754fb 953508f ae754fb 953508f ae754fb 953508f |
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 |
from pathlib import Path
import json
import numpy as np
from PIL import Image
from wireseghr.data.ttpla_to_masks import convert_ttpla_jsons_to_masks
def _read_dims(json_path: Path):
with open(json_path, "r") as f:
data = json.load(f)
return (
int(data["imageHeight"]),
int(data["imageWidth"]),
Path(data["imagePath"]).stem,
)
def test_convert_single_json_cable_only(tmp_path: Path):
# Use the provided example JSON at repo root
src_json = Path("/workspace/wire-seg-hr-impl/tests/1_00186_from_ttpla_dataset.json")
assert src_json.exists()
H, W, stem = _read_dims(src_json)
out_dir = tmp_path / "masks"
written = convert_ttpla_jsons_to_masks(src_json, out_dir, label="cable")
assert len(written) == 1
out_path = written[0]
# Converter writes numeric-only stems
expected_stem = "".join([c for c in stem if c.isdigit()])
assert out_path.name == f"{expected_stem}.png"
assert out_path.exists()
mask = np.array(Image.open(out_path).convert("L"))
assert mask is not None
assert mask.shape == (H, W)
assert mask.dtype == np.uint8
# Binary with values in {0,255}
uniq = np.unique(mask)
assert all(int(v) in (0, 255) for v in uniq)
assert (mask > 0).any(), "Expected some positive pixels for cable"
def test_convert_different_labels(tmp_path: Path):
src_json = Path("/workspace/wire-seg-hr-impl/tests/1_00186_from_ttpla_dataset.json")
assert src_json.exists()
out_dir_cable = tmp_path / "masks_cable"
out_dir_tower = tmp_path / "masks_tower"
written_cable = convert_ttpla_jsons_to_masks(src_json, out_dir_cable, label="cable")
written_tower = convert_ttpla_jsons_to_masks(
src_json, out_dir_tower, label="tower_wooden"
)
mc = np.array(Image.open(written_cable[0]).convert("L"))
mt = np.array(Image.open(written_tower[0]).convert("L"))
# Both masks should have some positives and should not be identical
assert (mc > 0).any()
assert (mt > 0).any()
assert not np.array_equal(mc, mt)
|