| | |
| | | from train_pointcloud_semantic_model import CLASS_SCHEMA, PointWiseNet |
| | | |
| | | |
| | | PREVIEW_POINT_LIMIT = 400_000 |
| | | |
| | | |
| | | def sha256(path: Path) -> str: |
| | | digest = hashlib.sha256() |
| | | with path.open("rb") as stream: |
| | |
| | | return xyz, np.clip(rgb, 0.0, 1.0), None |
| | | |
| | | |
| | | def preview_indices(predictions: np.ndarray, limit: int) -> np.ndarray: |
| | | """Deterministically retain every small predicted class before filling the budget.""" |
| | | count = len(predictions) |
| | | if count <= limit: |
| | | return np.arange(count, dtype=np.int64) |
| | | rng = np.random.default_rng(42) |
| | | selected: list[np.ndarray] = [] |
| | | remaining = limit |
| | | groups = sorted((np.flatnonzero(predictions == code) for code in np.unique(predictions)), key=len) |
| | | for group in groups: |
| | | keep = min(len(group), max(1, min(50_000, remaining // max(1, len(groups) - len(selected))))) |
| | | selected.append(group if keep == len(group) else np.sort(rng.choice(group, size=keep, replace=False))) |
| | | remaining -= keep |
| | | chosen = np.concatenate(selected) |
| | | if len(chosen) < limit: |
| | | mask = np.ones(count, dtype=bool) |
| | | mask[chosen] = False |
| | | fill = rng.choice(np.flatnonzero(mask), size=limit - len(chosen), replace=False) |
| | | chosen = np.concatenate((chosen, fill)) |
| | | if len(chosen) > limit: |
| | | chosen = rng.choice(chosen, size=limit, replace=False) |
| | | return np.sort(chosen.astype(np.int64)) |
| | | |
| | | |
| | | def load_model(path: Path, device: torch.device) -> tuple[PointWiseNet, list[int]]: |
| | | 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") != 1: |
| | | 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 code not in CLASS_SCHEMA for code in 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 |
| | | return model, codes, classes |
| | | |
| | | |
| | | def main() -> int: |
| | |
| | | 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()): |
| | |
| | | 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 = load_model(args.model, device) |
| | | 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) |
| | | pieces: list[np.ndarray] = [] |
| | | 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) |
| | | pieces.append(model(values).argmax(dim=1).cpu().numpy()) |
| | | predictions = np.asarray([class_codes[index] for index in np.concatenate(pieces)], dtype=np.uint8) |
| | | 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 = preview_indices(predictions, PREVIEW_POINT_LIMIT) |
| | | 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.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.") |
| | |
| | | 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": "deterministic class-aware cap; smaller predicted classes retained before the remaining budget is sampled", "input_has_rgb": True} |
| | | 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") |
| | | 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": {"preview": preview_path.name, "classified_las": classified_las.name, "class_counts": csv_path.name, "summary": summary_path.name}, "elapsed_seconds": round(time.perf_counter() - started, 3), "limitations": ["Predictions are model candidates, not asset inventory or inspection conclusions.", "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."]} |
| | | 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 |