"""Train a local CPU/GPU portable point-wise semantic classifier from human labels.
|
|
The training input is an immutable annotation revision created in the workbench.
|
It deliberately does not read the rule-classification labels as training truth.
|
"""
|
|
from __future__ import annotations
|
|
import argparse
|
import hashlib
|
import json
|
import time
|
from datetime import UTC, datetime
|
from pathlib import Path
|
from typing import Any
|
|
import numpy as np
|
import open3d as o3d
|
import torch
|
from sklearn.metrics import classification_report, confusion_matrix
|
|
|
CLASS_SCHEMA = {
|
1: {"key": "other_unknown", "label": "Other / unknown", "color": [128, 128, 128]},
|
2: {"key": "ground", "label": "Ground", "color": [151, 111, 51]},
|
5: {"key": "vegetation", "label": "Vegetation", "color": [59, 163, 87]},
|
6: {"key": "building_structure", "label": "Building / structure", "color": [224, 115, 55]},
|
15: {"key": "pole_tower", "label": "Pole / tower", "color": [149, 89, 210]},
|
16: {"key": "power_line", "label": "Power line", "color": [231, 196, 61]},
|
}
|
MIN_POINTS_PER_CLASS = 500
|
|
|
class PointWiseNet(torch.nn.Module):
|
"""Small PointNet-style shared MLP; same state dict runs on CPU or CUDA."""
|
|
def __init__(self, class_count: int) -> None:
|
super().__init__()
|
self.layers = torch.nn.Sequential(
|
torch.nn.Linear(6, 96), torch.nn.ReLU(), torch.nn.BatchNorm1d(96),
|
torch.nn.Linear(96, 128), torch.nn.ReLU(), torch.nn.Dropout(0.15),
|
torch.nn.Linear(128, 96), torch.nn.ReLU(), torch.nn.Linear(96, class_count),
|
)
|
|
def forward(self, values: torch.Tensor) -> torch.Tensor:
|
return self.layers(values)
|
|
|
def sha256(path: Path) -> str:
|
digest = hashlib.sha256()
|
with path.open("rb") as stream:
|
for chunk in iter(lambda: stream.read(8 * 1024 * 1024), b""):
|
digest.update(chunk)
|
return digest.hexdigest()
|
|
|
def load_features(source: Path) -> tuple[np.ndarray, dict[str, Any]]:
|
cloud = o3d.io.read_point_cloud(str(source))
|
xyz = np.asarray(cloud.points, dtype=np.float32)
|
rgb = np.asarray(cloud.colors, dtype=np.float32)
|
if len(xyz) < 1 or xyz.shape != rgb.shape:
|
raise ValueError("Annotation source must be a coloured point-cloud PLY.")
|
center = xyz.mean(axis=0)
|
scale = float(np.maximum(np.abs(xyz - center).max(), 1e-6))
|
return np.column_stack(((xyz - center) / scale, np.clip(rgb, 0.0, 1.0))).astype(np.float32), {
|
"xyz_center": center.tolist(), "xyz_scale": scale,
|
}
|
|
|
def spatial_split(xyz: np.ndarray, indices: np.ndarray) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
|
"""Split complete XY blocks so adjacent tower/line points cannot leak."""
|
block_size = max(float(np.ptp(xyz[:, 0])) / 8.0, float(np.ptp(xyz[:, 1])) / 8.0, 1.0)
|
blocks = np.floor(xyz[indices, :2] / block_size).astype(np.int64)
|
hashes = (blocks[:, 0] * 73856093 + blocks[:, 1] * 19349663) % 10
|
train, validation, test = indices[hashes < 7], indices[(hashes >= 7) & (hashes < 9)], indices[hashes >= 9]
|
if not len(validation) or not len(test) or not len(train):
|
ordered = np.sort(indices)
|
train, validation, test = ordered[::3], ordered[1::3], ordered[2::3]
|
return train, validation, test
|
|
|
def metrics(y_true: np.ndarray, y_pred: np.ndarray, classes: list[int]) -> dict[str, Any]:
|
names = [CLASS_SCHEMA[code]["key"] for code in classes]
|
return {
|
"class_codes": classes,
|
"class_keys": names,
|
"confusion_matrix": confusion_matrix(y_true, y_pred, labels=classes).tolist(),
|
"report": classification_report(y_true, y_pred, labels=classes, target_names=names, output_dict=True, zero_division=0),
|
"sample_count": int(len(y_true)),
|
}
|
|
|
def main() -> int:
|
parser = argparse.ArgumentParser(description="Train semantic point-cloud model from human-confirmed labels.")
|
parser.add_argument("--annotation", type=Path, required=True)
|
parser.add_argument("--output", type=Path, required=True)
|
parser.add_argument("--device", choices={"auto", "cpu", "cuda"}, default="auto")
|
parser.add_argument("--epochs", type=int, default=40)
|
parser.add_argument("--batch-size", type=int, default=4096)
|
args = parser.parse_args()
|
if args.output.exists() and any(args.output.iterdir()):
|
raise SystemExit("Output directory must be new or empty.")
|
annotation = json.loads(args.annotation.read_text(encoding="utf-8"))
|
if annotation.get("schema_version") != 1:
|
raise SystemExit("Unsupported annotation schema.")
|
source = Path(str(annotation.get("source_path") or ""))
|
if not source.is_file() or sha256(source) != annotation.get("source_sha256"):
|
raise SystemExit("Annotation source is unavailable or its checksum changed.")
|
labels = annotation.get("labels")
|
if not isinstance(labels, list):
|
raise SystemExit("Annotation labels are missing.")
|
features, normalizer = load_features(source)
|
source_cloud = o3d.io.read_point_cloud(str(source))
|
xyz = np.asarray(source_cloud.points, dtype=np.float32)
|
assigned: dict[int, int] = {}
|
for item in labels:
|
if not isinstance(item, list) or len(item) != 2 or not all(isinstance(value, int) for value in item):
|
raise SystemExit("Annotation contains an invalid point label.")
|
index, code = item
|
if not 0 <= index < len(features) or code not in CLASS_SCHEMA:
|
raise SystemExit("Annotation contains an out-of-range point label.")
|
assigned[index] = code
|
indices = np.asarray(sorted(assigned), dtype=np.int64)
|
target = np.asarray([assigned[int(index)] for index in indices], dtype=np.int64)
|
class_codes = sorted(set(target.tolist()))
|
per_class = {code: int((target == code).sum()) for code in class_codes}
|
if len(class_codes) < 2 or any(count < MIN_POINTS_PER_CLASS for count in per_class.values()):
|
raise SystemExit(f"Need at least two classes and {MIN_POINTS_PER_CLASS} confirmed points per class; current={per_class}.")
|
train_idx, validation_idx, test_idx = spatial_split(xyz, indices)
|
if min(len(train_idx), len(validation_idx), len(test_idx)) < len(class_codes):
|
raise SystemExit("Labels do not cover enough spatial blocks for train/validation/test evaluation.")
|
code_to_class = {code: offset for offset, code in enumerate(class_codes)}
|
label_by_index = {int(index): code_to_class[int(code)] for index, code in zip(indices, target, strict=True)}
|
device_name = "cuda" if args.device == "cuda" or (args.device == "auto" and torch.cuda.is_available()) else "cpu"
|
if args.device == "cuda" and not torch.cuda.is_available():
|
raise SystemExit("CUDA was requested but is unavailable.")
|
device = torch.device(device_name)
|
torch.manual_seed(42)
|
model = PointWiseNet(len(class_codes)).to(device)
|
optimizer = torch.optim.AdamW(model.parameters(), lr=0.001, weight_decay=1e-4)
|
weight = torch.tensor([len(train_idx) / max(1, sum(label_by_index[int(i)] == group for i in train_idx)) for group in range(len(class_codes))], dtype=torch.float32, device=device)
|
criterion = torch.nn.CrossEntropyLoss(weight=weight)
|
train_labels = np.asarray([label_by_index[int(index)] for index in train_idx], dtype=np.int64)
|
started = time.perf_counter()
|
for _ in range(args.epochs):
|
order = np.random.default_rng(42).permutation(len(train_idx))
|
model.train()
|
for start in range(0, len(order), args.batch_size):
|
subset = train_idx[order[start:start + args.batch_size]]
|
x = torch.from_numpy(features[subset]).to(device)
|
y = torch.from_numpy(np.asarray([label_by_index[int(index)] for index in subset], dtype=np.int64)).to(device)
|
optimizer.zero_grad(set_to_none=True)
|
criterion(model(x), y).backward()
|
optimizer.step()
|
def predict(indices_to_predict: np.ndarray) -> np.ndarray:
|
model.eval(); parts: list[np.ndarray] = []
|
with torch.no_grad():
|
for start in range(0, len(indices_to_predict), args.batch_size):
|
logits = model(torch.from_numpy(features[indices_to_predict[start:start + args.batch_size]]).to(device))
|
parts.append(logits.argmax(dim=1).cpu().numpy())
|
return np.concatenate(parts)
|
validation_true = np.asarray([target[np.searchsorted(indices, index)] for index in validation_idx], dtype=np.int64)
|
test_true = np.asarray([target[np.searchsorted(indices, index)] for index in test_idx], dtype=np.int64)
|
validation_pred = np.asarray([class_codes[value] for value in predict(validation_idx)], dtype=np.int64)
|
test_pred = np.asarray([class_codes[value] for value in predict(test_idx)], dtype=np.int64)
|
all_pred = np.asarray([class_codes[value] for value in predict(np.arange(len(features), dtype=np.int64))], dtype=np.uint8)
|
args.output.mkdir(parents=True, exist_ok=True)
|
predicted = o3d.geometry.PointCloud()
|
predicted.points = o3d.utility.Vector3dVector(xyz)
|
predicted.colors = o3d.utility.Vector3dVector(np.asarray([CLASS_SCHEMA[int(code)]["color"] for code in all_pred], dtype=np.float64) / 255.0)
|
o3d.io.write_point_cloud(str(args.output / "predicted-semantic-preview.ply"), predicted, write_ascii=False)
|
payload = {
|
"capability": "05-3d-pointcloud", "classification": "B", "created_at": datetime.now(UTC).isoformat(),
|
"model": "PointWiseNet shared MLP (human-confirmed point labels)", "device": device_name,
|
"torch_version": torch.__version__, "annotation": str(args.annotation), "source": str(source),
|
"source_sha256": annotation["source_sha256"], "classes": {str(code): CLASS_SCHEMA[code] for code in class_codes},
|
"label_counts": {str(code): count for code, count in per_class.items()},
|
"split_counts": {"train": int(len(train_idx)), "validation": int(len(validation_idx)), "test": int(len(test_idx))},
|
"validation": metrics(validation_true, validation_pred, class_codes), "test": metrics(test_true, test_pred, class_codes),
|
"normalizer": normalizer, "epochs": args.epochs, "batch_size": args.batch_size,
|
"elapsed_seconds": round(time.perf_counter() - started, 3),
|
"limitations": ["Metrics cover only human-confirmed points in this annotation revision.", "Spatial blocks reduce leakage but one small source cannot establish field-wide generalization.", "Rule candidate colours were not used as labels or input features."],
|
}
|
torch.save({"state_dict": model.cpu().state_dict(), "class_codes": class_codes, "normalizer": normalizer, "schema_version": 1}, args.output / "model.pt")
|
(args.output / "metrics.json").write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8")
|
print(json.dumps(payload, ensure_ascii=False))
|
return 0
|
|
|
if __name__ == "__main__":
|
raise SystemExit(main())
|
