"""Apply a human-labelled point-cloud semantic model to a new RGB point cloud.
|
|
This script deliberately accepts only explicit model/input/output paths. The
|
workbench server supplies those paths from fixed directories; it never forwards
|
browser paths or commands to this process.
|
"""
|
|
from __future__ import annotations
|
|
import argparse
|
import csv
|
import hashlib
|
import json
|
import sys
|
import time
|
from datetime import UTC, datetime
|
from pathlib import Path
|
from typing import Any
|
|
import laspy
|
import numpy as np
|
import open3d as o3d
|
import torch
|
|
from train_pointcloud_semantic_model import CLASS_SCHEMA, PointWiseNet
|
|
|
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 source_has_rgb(path: Path) -> bool:
|
suffix = path.suffix.lower()
|
if suffix in {".las", ".laz"}:
|
header = laspy.read(path).header
|
return all(name in header.point_format.dimension_names for name in ("red", "green", "blue"))
|
if suffix == ".ply":
|
with path.open("rb") as stream:
|
header = stream.read(64 * 1024).decode("latin-1", errors="ignore").lower()
|
return "end_header" in header and all(f" {name}" in header for name in ("red", "green", "blue"))
|
if suffix == ".pcd":
|
with path.open("rb") as stream:
|
header = stream.read(64 * 1024).decode("latin-1", errors="ignore").lower()
|
fields = next((line for line in header.splitlines() if line.startswith("fields ")), "")
|
return " rgb" in f" {fields}" or " rgba" in f" {fields}"
|
return False
|
|
|
def load_cloud(path: Path) -> tuple[np.ndarray, np.ndarray, Any | None]:
|
if path.suffix.lower() in {".las", ".laz"}:
|
las = laspy.read(path)
|
dimensions = set(las.point_format.dimension_names)
|
if not {"red", "green", "blue"}.issubset(dimensions):
|
raise ValueError("LAS/LAZ input has no RGB dimensions; this model requires observed RGB.")
|
xyz = np.column_stack((las.x, las.y, las.z)).astype(np.float32)
|
raw_rgb = np.column_stack((las.red, las.green, las.blue)).astype(np.float32)
|
if len(xyz) < 1 or xyz.shape != raw_rgb.shape:
|
raise ValueError(
|
"LAS/LAZ 中没有可读取的点。请选择完整的点云导出文件，不要上传空文件或未完成下载的分块。"
|
)
|
maximum = float(raw_rgb.max())
|
if maximum <= 0:
|
raise ValueError("LAS/LAZ input RGB values are all zero; observed RGB is required.")
|
return xyz, np.clip(raw_rgb / max(maximum, 1.0), 0.0, 1.0), las
|
if not source_has_rgb(path):
|
raise ValueError("Input has no readable RGB fields; XYZ-only point clouds cannot use this trained model.")
|
cloud = o3d.io.read_point_cloud(str(path))
|
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("Input point cloud has no readable XYZ/RGB points.")
|
return xyz, np.clip(rgb, 0.0, 1.0), None
|
|
|
def load_model(path: Path, device: torch.device) -> tuple[PointWiseNet, list[int], dict[int, dict[str, Any]]]:
|
payload = torch.load(path, map_location=device, weights_only=False)
|
if not isinstance(payload, dict) or payload.get("schema_version") not in {1, 2}:
|
raise ValueError("Unsupported model schema.")
|
codes = payload.get("class_codes")
|
if not isinstance(codes, list) or len(codes) < 2 or any(not isinstance(code, int) or not 1 <= code <= 255 for code in codes):
|
raise ValueError("Model class schema is invalid.")
|
saved_classes = payload.get("classes")
|
classes: dict[int, dict[str, Any]] = {}
|
for code in codes:
|
raw = saved_classes.get(str(code)) if isinstance(saved_classes, dict) else CLASS_SCHEMA.get(code)
|
if not isinstance(raw, dict):
|
raise ValueError("Model is missing class metadata for a predicted class.")
|
key, label, color = raw.get("key"), raw.get("label"), raw.get("color")
|
if not isinstance(key, str) or not isinstance(label, str) or not isinstance(color, list) or len(color) != 3 or not all(isinstance(item, int) and 0 <= item <= 255 for item in color):
|
raise ValueError("Model class metadata is invalid.")
|
classes[code] = {"code": code, "key": key, "label": label, "color": color}
|
model = PointWiseNet(len(codes)).to(device)
|
try:
|
model.load_state_dict(payload["state_dict"], strict=True)
|
except (KeyError, RuntimeError) as exc:
|
raise ValueError("Model weights do not match the supported PointWiseNet architecture.") from exc
|
model.eval()
|
return model, codes, classes
|
|
|
def main() -> int:
|
parser = argparse.ArgumentParser(description="Apply a local point-cloud semantic model to a new RGB point cloud.")
|
parser.add_argument("--model", type=Path, required=True)
|
parser.add_argument("--input", 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("--batch-size", type=int, default=4096)
|
parser.add_argument("--annotation-source-id", type=str, default="")
|
parser.add_argument("--candidate-confidence", type=float, default=0.95)
|
args = parser.parse_args()
|
if args.device == "cuda" and not torch.cuda.is_available():
|
raise SystemExit("CUDA was requested but is unavailable.")
|
if args.batch_size < 1 or args.batch_size > 262_144:
|
raise SystemExit("Batch size must be between 1 and 262144.")
|
if not 0.5 <= args.candidate_confidence < 1.0:
|
raise SystemExit("candidate-confidence must be between 0.5 and 1.0.")
|
if not args.model.is_file() or not args.input.is_file():
|
raise SystemExit("Model or input point cloud is unavailable.")
|
if args.output.exists() and any(args.output.iterdir()):
|
raise SystemExit("Output directory must be new or empty.")
|
|
started = time.perf_counter()
|
device_name = "cuda" if args.device == "cuda" or (args.device == "auto" and torch.cuda.is_available()) else "cpu"
|
device = torch.device(device_name)
|
model, class_codes, class_schema = load_model(args.model, device)
|
xyz, rgb, source_las = load_cloud(args.input)
|
center = xyz.mean(axis=0)
|
scale = float(max(np.abs(xyz - center).max(), 1e-6))
|
features = np.column_stack(((xyz - center) / scale, rgb)).astype(np.float32)
|
prediction_pieces: list[np.ndarray] = []
|
confidence_pieces: list[np.ndarray] = []
|
with torch.no_grad():
|
for start in range(0, len(features), args.batch_size):
|
values = torch.from_numpy(features[start:start + args.batch_size]).to(device)
|
probabilities = torch.softmax(model(values), dim=1)
|
confidence, predicted = probabilities.max(dim=1)
|
prediction_pieces.append(predicted.cpu().numpy())
|
confidence_pieces.append(confidence.cpu().numpy())
|
predictions = np.asarray([class_codes[index] for index in np.concatenate(prediction_pieces)], dtype=np.uint8)
|
confidence = np.concatenate(confidence_pieces).astype(np.float32, copy=False)
|
counts = {int(code): int((predictions == code).sum()) for code in class_codes}
|
|
args.output.mkdir(parents=True, exist_ok=True)
|
preview_selection = np.arange(len(predictions), dtype=np.int64)
|
preview = o3d.geometry.PointCloud()
|
preview.points = o3d.utility.Vector3dVector(xyz[preview_selection])
|
preview.colors = o3d.utility.Vector3dVector(np.asarray([class_schema[int(code)]["color"] for code in predictions[preview_selection]], dtype=np.float64) / 255.0)
|
preview_path = args.output / "predicted-semantic-preview.ply"
|
if not o3d.io.write_point_cloud(str(preview_path), preview, write_ascii=False):
|
raise RuntimeError("Could not write predicted PLY preview.")
|
|
classified_las = args.output / "predicted-semantic-classified.las"
|
if source_las is not None:
|
source_las.classification = predictions
|
source_las.write(classified_las)
|
else:
|
header = laspy.LasHeader(point_format=3, version="1.2")
|
header.offsets = xyz.min(axis=0)
|
header.scales = np.array([0.001, 0.001, 0.001])
|
generated = laspy.LasData(header)
|
generated.x, generated.y, generated.z = xyz[:, 0], xyz[:, 1], xyz[:, 2]
|
generated.red = np.rint(rgb[:, 0] * 65535).astype(np.uint16)
|
generated.green = np.rint(rgb[:, 1] * 65535).astype(np.uint16)
|
generated.blue = np.rint(rgb[:, 2] * 65535).astype(np.uint16)
|
generated.classification = predictions
|
generated.write(classified_las)
|
|
csv_path = args.output / "class-counts.csv"
|
with csv_path.open("w", newline="", encoding="utf-8") as stream:
|
writer = csv.DictWriter(stream, fieldnames=["class_code", "class_key", "class_label", "point_count"])
|
writer.writeheader()
|
for code in class_codes:
|
writer.writerow({"class_code": code, "class_key": class_schema[code]["key"], "class_label": class_schema[code]["label"], "point_count": counts[code]})
|
summary = {"class_codes": class_codes, "class_counts": {str(code): counts[code] for code in class_codes}, "input_points": int(len(xyz)), "preview_points": int(len(preview_selection)), "preview_sampling": "complete prediction point set without display sampling", "input_has_rgb": True}
|
candidate_path: Path | None = None
|
if args.annotation_source_id:
|
selected = np.flatnonzero(confidence >= args.candidate_confidence)
|
candidate_counts = {
|
int(code): int(np.count_nonzero(predictions[selected] == code))
|
for code in class_codes
|
}
|
candidate_path = args.output / "automatic-annotation-candidates.json"
|
candidate_path.write_text(json.dumps({
|
"schema_version": 1,
|
"source_id": args.annotation_source_id,
|
"source_sha256": sha256(args.input),
|
"model_sha256": sha256(args.model),
|
"candidate_confidence": args.candidate_confidence,
|
"input_points": int(len(xyz)),
|
"candidate_count": int(len(selected)),
|
"candidate_class_counts": {str(code): candidate_counts[code] for code in class_codes},
|
"labels": [[int(index), int(predictions[index]), round(float(confidence[index]), 6)] for index in selected],
|
}, ensure_ascii=False), encoding="utf-8")
|
summary["automatic_annotation"] = {
|
"candidate_confidence": args.candidate_confidence,
|
"candidate_count": int(len(selected)),
|
"candidate_class_counts": {str(code): candidate_counts[code] for code in class_codes},
|
"candidate_file": candidate_path.name,
|
}
|
summary_path = args.output / "prediction-summary.json"
|
summary_path.write_text(json.dumps(summary, ensure_ascii=False, indent=2), encoding="utf-8")
|
artifacts = {"preview": preview_path.name, "classified_las": classified_las.name, "class_counts": csv_path.name, "summary": summary_path.name}
|
if candidate_path:
|
artifacts["automatic_candidates"] = candidate_path.name
|
metadata = {"capability": "05-3d-pointcloud", "classification": "B", "created_at": datetime.now(UTC).isoformat(), "model": {"path": str(args.model), "sha256": sha256(args.model), "architecture": "PointWiseNet shared MLP"}, "input": {"path": str(args.input), "sha256": sha256(args.input), "bytes": args.input.stat().st_size, "points": int(len(xyz)), "has_rgb": True}, "classes": {str(code): class_schema[code] for code in class_codes}, "prediction": summary, "processing": {"requested_device": args.device, "device": device_name, "batch_size": args.batch_size, "normalization": {"method": "source-local per input", "xyz_center": center.tolist(), "xyz_scale": scale}}, "versions": {"python": sys.version.split()[0], "torch": torch.__version__, "open3d": o3d.__version__, "laspy": laspy.__version__}, "artifacts": artifacts, "elapsed_seconds": round(time.perf_counter() - started, 3), "limitations": ["Predictions are model candidates, not asset inventory or inspection conclusions.", "Automatic candidates retain only points at or above the recorded confidence threshold; user confirmation is required before merging them into a training revision.", "This model requires observed RGB; it cannot infer labels for XYZ-only point clouds.", "Model metrics apply only to the labelled source spatial blocks. The current pole/tower class has high false-positive risk and requires review.", "New inputs are normalized with their own XYZ centre and scale to match the training feature definition; this preserves their coordinates but does not prove cross-site generalization."]}
|
(args.output / "run_metadata.json").write_text(json.dumps(metadata, ensure_ascii=False, indent=2), encoding="utf-8")
|
print(json.dumps(metadata, ensure_ascii=False))
|
return 0
|
|
|
if __name__ == "__main__":
|
raise SystemExit(main())
|
