| | |
|---|
| | | } |
|---|
| | | |
|---|
| | | |
|---|
| | | def load_training_dataset(annotation_paths: list[Path]) -> tuple[np.ndarray, np.ndarray, dict[int, int], dict[int, dict[str, Any]], dict[str, np.ndarray], list[dict[str, Any]], np.ndarray]: |
|---|
| | | """Load immutable RGB/XYZ revisions and split each source independently. |
|---|
| | | |
|---|
| | | Point indices and local coordinates are valid only within their source. A |
|---|
| | | source-scoped XY split prevents neighbouring points from one source leaking |
|---|
| | | into another split while still allowing several scenes to train one model. |
|---|
| | | """ |
|---|
| | | if not annotation_paths: |
|---|
| | | raise SystemExit("Provide at least one annotation revision.") |
|---|
| | | source_ids: set[str] = set() |
|---|
| | | key_to_code: dict[str, int] = {} |
|---|
| | | class_schema: dict[int, dict[str, Any]] = {} |
|---|
| | | used_codes: set[int] = set() |
|---|
| | | features_parts: list[np.ndarray] = [] |
|---|
| | | xyz_parts: list[np.ndarray] = [] |
|---|
| | | label_by_index: dict[int, int] = {} |
|---|
| | | split_parts: dict[str, list[np.ndarray]] = {"train": [], "validation": [], "test": []} |
|---|
| | | sources: list[dict[str, Any]] = [] |
|---|
| | | offset = 0 |
|---|
| | | |
|---|
| | | for annotation_path in annotation_paths: |
|---|
| | | annotation = json.loads(annotation_path.read_text(encoding="utf-8")) |
|---|
| | | if annotation.get("schema_version") != 1: |
|---|
| | | raise SystemExit(f"Unsupported annotation schema: {annotation_path}") |
|---|
| | | source_id = annotation.get("source_id") |
|---|
| | | if not isinstance(source_id, str) or not source_id: |
|---|
| | | raise SystemExit(f"Annotation source id is missing: {annotation_path}") |
|---|
| | | if source_id in source_ids: |
|---|
| | | raise SystemExit("Choose at most one saved annotation revision for each source.") |
|---|
| | | source_ids.add(source_id) |
|---|
| | | source = Path(str(annotation.get("source_path") or "")) |
|---|
| | | if not source.is_file() or sha256(source) != annotation.get("source_sha256"): |
|---|
| | | raise SystemExit(f"Annotation source is unavailable or its checksum changed: {annotation_path}") |
|---|
| | | labels = annotation.get("labels") |
|---|
| | | if not isinstance(labels, list): |
|---|
| | | raise SystemExit(f"Annotation labels are missing: {annotation_path}") |
|---|
| | | source_schema = annotation_schema(annotation) |
|---|
| | | 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] = {} |
|---|
| | | code_mapping: 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(f"Annotation contains an invalid point label: {annotation_path}") |
|---|
| | | index, source_code = item |
|---|
| | | if not 0 <= index < len(features) or source_code not in source_schema: |
|---|
| | | raise SystemExit(f"Annotation contains an out-of-range point label: {annotation_path}") |
|---|
| | | definition = source_schema[source_code] |
|---|
| | | key = definition["key"] |
|---|
| | | canonical_code = key_to_code.get(key) |
|---|
| | | if canonical_code is None: |
|---|
| | | canonical_code = source_code if source_code not in used_codes else next((code for code in range(1, 256) if code not in used_codes), None) |
|---|
| | | if canonical_code is None: |
|---|
| | | raise SystemExit("No unused LAS-compatible class code is available for the training dataset.") |
|---|
| | | key_to_code[key] = canonical_code |
|---|
| | | class_schema[canonical_code] = {"code": canonical_code, "key": key, "label": definition["label"], "color": definition["color"]} |
|---|
| | | used_codes.add(canonical_code) |
|---|
| | | code_mapping[source_code] = canonical_code |
|---|
| | | assigned[index] = canonical_code |
|---|
| | | indices = np.asarray(sorted(assigned), dtype=np.int64) |
|---|
| | | if not len(indices): |
|---|
| | | raise SystemExit(f"Annotation contains no confirmed point labels: {annotation_path}") |
|---|
| | | source_target = np.asarray([assigned[int(index)] for index in indices], dtype=np.int64) |
|---|
| | | source_splits = dict(zip(("train", "validation", "test"), spatial_split(xyz, indices), strict=True)) |
|---|
| | | for split_name, split_indices in source_splits.items(): |
|---|
| | | split_parts[split_name].append(offset + split_indices) |
|---|
| | | label_by_index.update({offset + int(index): int(code) for index, code in zip(indices, source_target, strict=True)}) |
|---|
| | | features_parts.append(features) |
|---|
| | | xyz_parts.append(xyz) |
|---|
| | | counts = {str(code): int((source_target == code).sum()) for code in sorted(set(source_target.tolist()))} |
|---|
| | | sources.append({ |
|---|
| | | "annotation": str(annotation_path), "annotation_sha256": sha256(annotation_path), |
|---|
| | | "source_id": source_id, "source": str(source), "source_sha256": annotation["source_sha256"], |
|---|
| | | "point_count": int(len(features)), "label_count": int(len(indices)), "label_counts": counts, |
|---|
| | | "class_code_mapping": {str(code): mapped for code, mapped in sorted(code_mapping.items())}, |
|---|
| | | "normalizer": normalizer, |
|---|
| | | "split_counts": {name: int(len(value)) for name, value in source_splits.items()}, |
|---|
| | | }) |
|---|
| | | offset += len(features) |
|---|
| | | |
|---|
| | | features = np.vstack(features_parts).astype(np.float32, copy=False) |
|---|
| | | xyz = np.vstack(xyz_parts).astype(np.float32, copy=False) |
|---|
| | | splits = {name: np.concatenate(parts).astype(np.int64, copy=False) for name, parts in split_parts.items()} |
|---|
| | | class_codes = sorted(set(label_by_index.values())) |
|---|
| | | if len(class_codes) < 2: |
|---|
| | | raise SystemExit("Need at least two confirmed classes across the selected annotation revisions.") |
|---|
| | | counts = {code: sum(value == code for value in label_by_index.values()) for code in class_codes} |
|---|
| | | if any(count < MIN_POINTS_PER_CLASS for count in counts.values()): |
|---|
| | | raise SystemExit(f"Need at least {MIN_POINTS_PER_CLASS} confirmed points per class across the selected revisions; current={counts}.") |
|---|
| | | for name, split in splits.items(): |
|---|
| | | split_classes = {label_by_index[int(index)] for index in split} |
|---|
| | | if not len(split) or split_classes != set(class_codes): |
|---|
| | | raise SystemExit(f"Every confirmed class must cover the {name} split across the selected sources. Add labels in more spatial areas or sources.") |
|---|
| | | return features, xyz, label_by_index, class_schema, splits, sources, np.asarray(class_codes, dtype=np.int64) |
|---|
| | | |
|---|
| | | |
|---|
| | | 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("--annotation", type=Path, help="One immutable annotation revision (backward-compatible single-source mode).") |
|---|
| | | parser.add_argument("--annotations", type=Path, nargs="+", help="One immutable revision per RGB/XYZ annotation source for concentrated training.") |
|---|
| | | 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) |
|---|
| | |
|---|
| | | 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.") |
|---|
| | | class_schema = annotation_schema(annotation) |
|---|
| | | 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.") |
|---|
| | | if args.annotation and args.annotations: |
|---|
| | | raise SystemExit("Use either --annotation or --annotations, not both.") |
|---|
| | | annotation_paths = args.annotations or ([args.annotation] if args.annotation else []) |
|---|
| | | features, xyz, label_by_index, class_schema, splits, sources, class_codes_array = load_training_dataset(annotation_paths) |
|---|
| | | class_codes = class_codes_array.tolist() |
|---|
| | | per_class = {code: sum(value == code for value in label_by_index.values()) for code in class_codes} |
|---|
| | | train_idx, validation_idx, test_idx = splits["train"], splits["validation"], splits["test"] |
|---|
| | | 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)} |
|---|
| | | label_by_index = {index: code_to_class[code] for index, code in label_by_index.items()} |
|---|
| | | 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.") |
|---|
| | |
|---|
| | | 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) |
|---|
| | | 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_true = np.asarray([class_codes[label_by_index[int(index)]] for index in validation_idx], dtype=np.int64) |
|---|
| | | test_true = np.asarray([class_codes[label_by_index[int(index)]] for index in test_idx], dtype=np.int64) |
|---|
| | | |
|---|
| | | def predict(indices_to_predict: np.ndarray) -> np.ndarray: |
|---|
| | | model.eval(); parts: list[np.ndarray] = [] |
|---|
| | |
|---|
| | | 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) |
|---|
| | | # Sources can use unrelated local coordinate frames. The first selected |
|---|
| | | # source remains the fixed preview artifact; per-source provenance records |
|---|
| | | # the full concentrated-training dataset without pretending it is one scene. |
|---|
| | | preview_count = int(sources[0]["point_count"]) |
|---|
| | | 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) |
|---|
| | | predicted.points = o3d.utility.Vector3dVector(xyz[:preview_count]) |
|---|
| | | predicted.colors = o3d.utility.Vector3dVector(np.asarray([class_schema[int(code)]["color"] for code in all_pred[:preview_count]], 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}, |
|---|
| | | "torch_version": torch.__version__, "annotation": sources[0]["annotation"], "source": sources[0]["source"], |
|---|
| | | "source_sha256": sources[0]["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, class_schema), "test": metrics(test_true, test_pred, class_codes, class_schema), |
|---|
| | | "normalizer": normalizer, "epochs": args.epochs, "best_epoch": best_epoch, "best_validation_macro_f1": best_validation_macro_f1, "batch_size": args.batch_size, "seed": args.seed, |
|---|
| | | "normalizer": sources[0]["normalizer"], "epochs": args.epochs, "best_epoch": best_epoch, "best_validation_macro_f1": best_validation_macro_f1, "batch_size": args.batch_size, "seed": args.seed, |
|---|
| | | "training_dataset": {"schema_version": 1, "kind": "rgb_xyz_annotation_revisions", "annotation_count": len(sources), "annotations": sources, "preview_annotation": sources[0]["annotation"], "preview_scope": "first_selected_source_only"}, |
|---|
| | | "elapsed_seconds": round(time.perf_counter() - started, 3), |
|---|
| | | "limitations": ["Metrics cover only human-confirmed points in this annotation revision.", "The retained checkpoint is selected by validation macro F1; the test split remains separate from that selection.", "Spatial blocks reduce leakage but one small source cannot establish field-wide generalization.", "Rule candidate colours were not used as labels or input features."], |
|---|
| | | "limitations": ["Metrics cover only human-confirmed points in the selected immutable annotation revisions.", "The retained checkpoint is selected by validation macro F1; the test split remains separate from that selection.", "Each source is split into local XY blocks before concentrated training; this reduces local leakage but does not prove field-wide generalization.", "The fixed preview PLY contains only the first selected source because different sources can use unrelated local coordinate frames.", "Rule candidate colours were not used as labels or input features."], |
|---|
| | | } |
|---|
| | | torch.save({"state_dict": model.cpu().state_dict(), "class_codes": class_codes, "classes": {str(code): class_schema[code] for code in class_codes}, "normalizer": normalizer, "schema_version": 2}, args.output / "model.pt") |
|---|
| | | torch.save({"state_dict": model.cpu().state_dict(), "class_codes": class_codes, "classes": {str(code): class_schema[code] for code in class_codes}, "normalizer": sources[0]["normalizer"], "schema_version": 2}, args.output / "model.pt") |
|---|
| | | (args.output / "metrics.json").write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8") |
|---|
| | | (args.output / "training_dataset.json").write_text(json.dumps(payload["training_dataset"], ensure_ascii=False, indent=2), encoding="utf-8") |
|---|
| | | print(json.dumps(payload, ensure_ascii=False)) |
|---|
| | | return 0 |
|---|
| | | |
|---|
