Upload 2 files

batch.py

@@ -0,0 +1,138 @@
+"""Batch CLI: STEP folder -> Excel. Single process, models loaded once."""
+from __future__ import annotations
+import argparse
+import shutil
+import sys
+import tempfile
+from datetime import datetime
+from pathlib import Path
+from typing import Dict, List
+
+from openpyxl import Workbook
+
+from heg_brep import (
+    DEFAULT_PASS1_MODEL, DEFAULT_ELBOW_MODEL, DEFAULT_TEE_MODEL, REPO_ROOT,
+)
+from heg_brep.inference import LoadedModel, TwoPassClassifier
+from heg_brep.extraction import extract_folder_to_npz
+
+
+COLUMNS = [
+    "file_name", "file_path", "stem", "status",
+    "pass1_argmax", "pass1_conf", "route",
+    "pass2_argmax", "pass2_predicted", "pass2_conf",
+    "final_label", "final_conf", "npz_path", "note",
+]
+
+
+def parse_args() -> argparse.Namespace:
+    ap = argparse.ArgumentParser(description=__doc__)
+    ap.add_argument("step_folder")
+    ap.add_argument("output_excel", nargs="?", default=None,
+                    help="Defaults to ./brep_two_pass_<timestamp>.xlsx")
+    ap.add_argument("--pass1_model", default=str(DEFAULT_PASS1_MODEL))
+    ap.add_argument("--elbow_model", default=str(DEFAULT_ELBOW_MODEL))
+    ap.add_argument("--tee_model",   default=str(DEFAULT_TEE_MODEL))
+    ap.add_argument("--device", default="cpu", choices=["cpu", "cuda"])
+    ap.add_argument("--pass2_min_conf", type=float, default=0.85)
+    ap.add_argument("--pass2_tau", type=float, default=0.0)
+    ap.add_argument("--num_workers", type=int, default=1)
+    ap.add_argument("--max_file_mb", type=float, default=None)
+    ap.add_argument("--npz_dir", default=None, help="Persist NPZs here (default: temp).")
+    return ap.parse_args()
+
+
+def write_excel(rows: List[Dict[str, object]], output_excel: Path) -> None:
+    wb = Workbook()
+    ws = wb.active
+    ws.title = "brep_two_pass"
+    ws.append(COLUMNS)
+    for r in rows:
+        ws.append([r.get(c, "") for c in COLUMNS])
+    for col in ws.columns:
+        letter = col[0].column_letter
+        max_len = max((len(str(c.value)) if c.value is not None else 0 for c in col), default=0)
+        ws.column_dimensions[letter].width = min(max(12, max_len + 2), 70)
+    output_excel.parent.mkdir(parents=True, exist_ok=True)
+    wb.save(output_excel)
+
+
+def main() -> int:
+    args = parse_args()
+    step_folder = Path(args.step_folder).expanduser().resolve()
+    if not step_folder.is_dir():
+        raise FileNotFoundError(f"step_folder not found: {step_folder}")
+
+    output_excel = (Path(args.output_excel).expanduser().resolve()
+                    if args.output_excel
+                    else REPO_ROOT / f"brep_two_pass_{datetime.now().strftime('%Y%m%d_%H%M%S')}.xlsx")
+    temp_npz_dir = args.npz_dir is None
+    npz_dir = (Path(args.npz_dir).expanduser().resolve() if args.npz_dir
+               else Path(tempfile.mkdtemp(prefix="heg_brep_npz_")))
+    try:
+        print(f"Extracting STEP files in {step_folder} -> {npz_dir}")
+        t0 = datetime.now()
+        summary = extract_folder_to_npz(
+            step_folder, output_dir=npz_dir,
+            num_workers=args.num_workers, max_file_mb=args.max_file_mb,
+        )
+        print(f"  extraction took {(datetime.now() - t0).total_seconds():.1f}s "
+              f"({len(summary['ok_stems'])} OK, {len(summary['skipped'])} skipped)")
+
+        print(f"Loading models on {args.device}...")
+        t0 = datetime.now()
+        pass1 = LoadedModel(Path(args.pass1_model), device=args.device)
+        elbow = LoadedModel(Path(args.elbow_model), device=args.device) if Path(args.elbow_model).exists() else None
+        tee   = LoadedModel(Path(args.tee_model),   device=args.device) if Path(args.tee_model).exists()   else None
+        clf = TwoPassClassifier(pass1=pass1, elbow=elbow, tee=tee,
+                                pass2_min_conf=args.pass2_min_conf, pass2_tau=args.pass2_tau)
+        print(f"  models loaded in {(datetime.now() - t0).total_seconds():.1f}s")
+
+        print(f"Classifying {len(summary['step_paths'])} files...")
+        t0 = datetime.now()
+        rows: List[Dict[str, object]] = []
+        ok = pass1_failed = pass2_failed = extraction_failed = 0
+        for sp in summary["step_paths"]:
+            stem = sp.stem
+            base = {c: "" for c in COLUMNS}
+            base.update({"file_name": sp.name, "file_path": str(sp), "stem": stem})
+            npz = npz_dir / f"{stem}.npz"
+            if not npz.exists():
+                base["status"] = "EXTRACTION_FAILED"
+                base["note"] = summary["skipped"].get(stem, "NPZ missing after extraction")
+                extraction_failed += 1
+                rows.append(base); continue
+            base["npz_path"] = str(npz)
+            try:
+                result = clf.classify_npz(npz)
+            except Exception as exc:
+                base["status"] = "INFERENCE_FAILED"
+                base["note"] = str(exc)[:500]
+                pass1_failed += 1
+                rows.append(base); continue
+            base.update({k: result.get(k, "") for k in result})
+            # Round confidences for the spreadsheet
+            for k in ("pass1_conf", "pass2_conf", "final_conf"):
+                v = base.get(k)
+                base[k] = round(float(v), 6) if isinstance(v, (int, float)) else v
+            base["status"] = "OK"
+            ok += 1
+            rows.append(base)
+
+        elapsed = (datetime.now() - t0).total_seconds()
+        print(f"  inference took {elapsed:.1f}s ({elapsed / max(1, len(rows)):.2f}s/file)")
+
+        write_excel(rows, output_excel)
+        print(f"Saved Excel: {output_excel}")
+        print(f"Summary -> total: {len(rows)} | OK: {ok} | "
+              f"extraction_failed: {extraction_failed} | inference_failed: {pass1_failed}")
+        if not temp_npz_dir:
+            print(f"NPZ dir: {npz_dir}")
+        return 0
+    finally:
+        if temp_npz_dir:
+            shutil.rmtree(npz_dir, ignore_errors=True)
+
+
+if __name__ == "__main__":
+    sys.exit(main())

server.py

@@ -0,0 +1,178 @@
+"""Local FastAPI service: STEP-file-in, classification-out. One process, models loaded once.
+
+Designed to be launched by a C# host app as a child process. The service binds to
+127.0.0.1 only; never exposes a network-facing port.
+
+Startup protocol:
+  - Bind a free port (default: OS-assigned via --port 0).
+  - Once models are loaded and the server is accepting requests, write
+    `READY port=<PORT>` to stdout (single line). The host reads that line to learn
+    where to connect.
+  - On `POST /shutdown` (or SIGTERM), drain in-flight requests and exit.
+
+Endpoints:
+  GET  /health                 -> service / model state
+  POST /classify {step_path}   -> single STEP classification
+  POST /classify_batch {paths} -> multiple in one call (sequential, single-process)
+  POST /shutdown               -> graceful exit
+"""
+from __future__ import annotations
+import argparse
+import os
+import shutil
+import sys
+import tempfile
+import threading
+import time
+from pathlib import Path
+from typing import List, Optional
+
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+import uvicorn
+
+from heg_brep import (
+    DEFAULT_PASS1_MODEL, DEFAULT_ELBOW_MODEL, DEFAULT_TEE_MODEL,
+)
+from heg_brep.inference import LoadedModel, TwoPassClassifier
+from heg_brep.extraction import extract_step_to_npz
+
+# Eagerly import the BRepExtractor pipeline at startup. Without this, the first
+# /classify call pays ~20-30s of OCC + occwl + igl cold-import cost — terrible
+# UX for an interactive viewer.
+import pipeline.extract_brep_extractor_data_from_step  # noqa: F401
+
+
+class State:
+    classifier: Optional[TwoPassClassifier] = None
+    started_at: float = 0.0
+    device: str = "cpu"
+    pass2_min_conf: float = 0.85
+    pass2_tau: float = 0.0
+
+
+app = FastAPI(title="HEG BRep component identification")
+
+
+class ClassifyRequest(BaseModel):
+    step_path: str
+    # Optional override: persist the NPZ to disk for inspection.
+    npz_keep_dir: Optional[str] = None
+
+
+class ClassifyBatchRequest(BaseModel):
+    step_paths: List[str]
+    npz_keep_dir: Optional[str] = None
+
+
+@app.get("/health")
+def health():
+    return {
+        "status": "ok",
+        "models_loaded": State.classifier is not None,
+        "device": State.device,
+        "uptime_sec": round(time.time() - State.started_at, 2),
+    }
+
+
+def _classify_one(step_path: str, npz_keep_dir: Optional[str]) -> dict:
+    if State.classifier is None:
+        raise HTTPException(status_code=503, detail="models not loaded yet")
+    sp = Path(step_path).expanduser()
+    if not sp.exists():
+        return {"step_path": str(sp), "status": "error", "error": "step_path not found"}
+    temp_npz_dir = npz_keep_dir is None
+    out_dir = Path(npz_keep_dir).expanduser().resolve() if npz_keep_dir \
+              else Path(tempfile.mkdtemp(prefix="heg_brep_npz_"))
+    try:
+        try:
+            npz = extract_step_to_npz(sp, out_dir)
+        except Exception as exc:
+            return {"step_path": str(sp), "status": "extraction_failed", "error": str(exc)[:500]}
+        try:
+            result = State.classifier.classify_npz(npz)
+        except Exception as exc:
+            return {"step_path": str(sp), "status": "inference_failed", "error": str(exc)[:500],
+                    "npz_path": str(npz)}
+        result.update({"step_path": str(sp), "status": "ok", "npz_path": str(npz)})
+        return result
+    finally:
+        if temp_npz_dir:
+            shutil.rmtree(out_dir, ignore_errors=True)
+
+
+@app.post("/classify")
+def classify(req: ClassifyRequest):
+    return _classify_one(req.step_path, req.npz_keep_dir)
+
+
+@app.post("/classify_batch")
+def classify_batch(req: ClassifyBatchRequest):
+    return {"results": [_classify_one(p, req.npz_keep_dir) for p in req.step_paths]}
+
+
+@app.post("/shutdown")
+def shutdown():
+    threading.Thread(target=lambda: (time.sleep(0.2), os._exit(0)), daemon=True).start()
+    return {"status": "shutting_down"}
+
+
+def _load_models(args) -> None:
+    pass1 = LoadedModel(Path(args.pass1_model), device=args.device)
+    elbow = LoadedModel(Path(args.elbow_model), device=args.device) if Path(args.elbow_model).exists() else None
+    tee   = LoadedModel(Path(args.tee_model),   device=args.device) if Path(args.tee_model).exists()   else None
+    State.classifier = TwoPassClassifier(
+        pass1=pass1, elbow=elbow, tee=tee,
+        pass2_min_conf=args.pass2_min_conf, pass2_tau=args.pass2_tau,
+    )
+    State.device = args.device
+    State.pass2_min_conf = args.pass2_min_conf
+    State.pass2_tau = args.pass2_tau
+
+
+def parse_args() -> argparse.Namespace:
+    ap = argparse.ArgumentParser(description="HEG BRep classification service")
+    ap.add_argument("--host", default="127.0.0.1")
+    ap.add_argument("--port", type=int, default=0,
+                    help="Port to bind. 0 = let the OS pick. The chosen port is "
+                         "printed to stdout as a single line `READY port=<N>`.")
+    ap.add_argument("--pass1_model", default=str(DEFAULT_PASS1_MODEL))
+    ap.add_argument("--elbow_model", default=str(DEFAULT_ELBOW_MODEL))
+    ap.add_argument("--tee_model",   default=str(DEFAULT_TEE_MODEL))
+    ap.add_argument("--device", default="cpu", choices=["cpu", "cuda"])
+    ap.add_argument("--pass2_min_conf", type=float, default=0.85)
+    ap.add_argument("--pass2_tau", type=float, default=0.0)
+    return ap.parse_args()
+
+
+def main() -> int:
+    args = parse_args()
+    State.started_at = time.time()
+    print(f"[heg_brep] loading models from:", file=sys.stderr)
+    print(f"  pass1: {args.pass1_model}", file=sys.stderr)
+    print(f"  elbow: {args.elbow_model}", file=sys.stderr)
+    print(f"  tee  : {args.tee_model}",   file=sys.stderr)
+    print(f"  device: {args.device}",     file=sys.stderr)
+    t0 = time.time()
+    _load_models(args)
+    print(f"[heg_brep] models loaded in {time.time() - t0:.1f}s", file=sys.stderr)
+
+    # Bind socket up front so we can tell the parent which port we got.
+    import socket
+    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+    sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+    sock.bind((args.host, int(args.port)))
+    bound_port = sock.getsockname()[1]
+    print(f"READY port={bound_port}", flush=True)
+
+    config = uvicorn.Config(app=app, host=args.host, port=bound_port, log_level="warning")
+    server = uvicorn.Server(config)
+    # uvicorn doesn't accept an already-bound socket in the simple Config API; close
+    # ours and let uvicorn rebind. Race window is fine because we're loopback only.
+    sock.close()
+    server.run()
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())