from __future__ import annotations
|
|
import hashlib
|
import json
|
import sys
|
import tempfile
|
import unittest
|
from pathlib import Path
|
|
ROOT = Path(__file__).resolve().parents[3]
|
sys.path.insert(0, str(ROOT / "capabilities" / "05-3d-pointcloud"))
|
|
from train_pointcloud_semantic_model import load_training_dataset # noqa: E402
|
|
|
def checksum(path: Path) -> str:
|
return hashlib.sha256(path.read_bytes()).hexdigest()
|
|
|
def write_rgb_ply(path: Path, offset: float) -> None:
|
points = []
|
for index in range(1_000):
|
x, y = float(index % 100) + offset, float(index // 100)
|
red, green, blue = (210, 40, 40) if index % 2 else (40, 180, 60)
|
points.append(f"{x} {y} 0 {red} {green} {blue}")
|
path.write_text(
|
"ply\nformat ascii 1.0\nelement vertex 1000\nproperty float x\nproperty float y\nproperty float z\nproperty uchar red\nproperty uchar green\nproperty uchar blue\nend_header\n"
|
+ "\n".join(points)
|
+ "\n",
|
encoding="ascii",
|
)
|
|
|
class ConcentratedTrainingDatasetTests(unittest.TestCase):
|
def test_combines_two_sources_and_keeps_source_scoped_splits(self) -> None:
|
with tempfile.TemporaryDirectory() as temporary:
|
directory = Path(temporary)
|
revisions = []
|
for source_index in range(2):
|
source = directory / f"source-{source_index}.ply"
|
write_rgb_ply(source, source_index * 1_000.0)
|
revision = directory / f"annotation-{source_index}.json"
|
revision.write_text(json.dumps({
|
"schema_version": 1,
|
"source_id": f"source-{source_index}",
|
"source_path": str(source),
|
"source_sha256": checksum(source),
|
"labels": [[index, 2 if index % 2 else 5] for index in range(1_000)],
|
"class_schema": {
|
"2": {"code": 2, "key": "ground", "label": "Ground", "color": [151, 111, 51]},
|
"5": {"code": 5, "key": "vegetation", "label": "Vegetation", "color": [59, 163, 87]},
|
},
|
}), encoding="utf-8")
|
revisions.append(revision)
|
|
features, xyz, labels, schema, splits, sources, class_codes = load_training_dataset(revisions)
|
|
self.assertEqual(features.shape, (2_000, 6))
|
self.assertEqual(xyz.shape, (2_000, 3))
|
self.assertEqual(len(labels), 2_000)
|
self.assertEqual(class_codes.tolist(), [2, 5])
|
self.assertEqual({value["key"] for value in schema.values()}, {"ground", "vegetation"})
|
self.assertEqual(len(sources), 2)
|
self.assertTrue(all(source["split_counts"]["train"] > 0 for source in sources))
|
for split in splits.values():
|
self.assertEqual({labels[int(index)] for index in split}, {2, 5})
|
|
|
if __name__ == "__main__":
|
unittest.main()
|
