| """Inference with the hlt-geometric-vision release (ONNXRuntime + OpenCV, no PyTorch).
|
|
|
| Path A: the vertices come out of the model (1 polygon, training domain = clean rectangles).
|
| Path B: model edge_map + geometric extractor -> N rectangles on real photographs.
|
|
|
| python inference_example.py --model ../model.onnx --image photo.jpg [--nested]
|
|
|
| The extractor is a copy of hlt/postproc/quads.py from the source repository (numpy + OpenCV,
|
| no weights): see documentation.md section 6 for the specification, section 8 for the parameters.
|
| """
|
| from __future__ import annotations
|
|
|
| import argparse
|
| import json
|
| from pathlib import Path
|
|
|
| import cv2
|
| import numpy as np
|
| import onnxruntime as ort
|
|
|
| from hlt_postproc_quads import QuadConfig, extract_quads, hybrid_edge_map
|
|
|
|
|
| def preprocess(bgr, size=512):
|
| """Letterbox to a square, aspect preserved, with a REPLICATED border (a constant gray fill
|
| creates artificial lines in the edge map). Also returns the real-image-area mask."""
|
| h0, w0 = bgr.shape[:2]
|
| scale = float(size) / float(max(h0, w0))
|
| nw, nh = max(1, round(w0 * scale)), max(1, round(h0 * scale))
|
| small = cv2.resize(bgr, (nw, nh), interpolation=cv2.INTER_AREA)
|
| ox, oy = (size - nw) // 2, (size - nh) // 2
|
| disp = cv2.copyMakeBorder(small, oy, size - nh - oy, ox, size - nw - ox, cv2.BORDER_REPLICATE)
|
| valid = np.zeros((size, size), bool)
|
| valid[oy:oy + nh, ox:ox + nw] = True
|
| x = cv2.cvtColor(disp, cv2.COLOR_BGR2RGB).astype(np.float32).transpose(2, 0, 1) / 255.0
|
| return np.ascontiguousarray(x)[None], disp, valid, scale, ox, oy
|
|
|
|
|
| def to_original(corners, scale, ox, oy):
|
| """512 frame -> pixels of the original image. The half-pixel correction is not cosmetic: on
|
| large photographs the mapping error reaches ~1.6 px."""
|
| return (corners + 0.5 - np.array([ox, oy], np.float64)) / scale - 0.5
|
|
|
|
|
| def main():
|
| ap = argparse.ArgumentParser()
|
| ap.add_argument("--model", default="model.onnx")
|
| ap.add_argument("--image", required=True)
|
| ap.add_argument("--nested", action="store_true", help="also keep nested rectangles")
|
| ap.add_argument("--only-net", action="store_true", help="path A only (model vertices)")
|
| ap.add_argument("--no-canny", action="store_true", help="network edge map only (no Canny)")
|
| ap.add_argument("--threads", type=int, default=0, help="ONNXRuntime threads (0 = default)")
|
| ap.add_argument("--max-quads", type=int, default=64,
|
| help="cap on kept rectangles (QuadConfig default: 64; the release examples use 12)")
|
| args = ap.parse_args()
|
|
|
| so = ort.SessionOptions()
|
| if args.threads:
|
| so.intra_op_num_threads = args.threads
|
| sess = ort.InferenceSession(args.model, sess_options=so, providers=["CPUExecutionProvider"])
|
| in_name = sess.get_inputs()[0].name
|
| out_names = [o.name for o in sess.get_outputs()]
|
|
|
| img = cv2.imread(args.image, cv2.IMREAD_COLOR)
|
| if img is None:
|
| raise SystemExit("image not decoded: " + args.image)
|
| x, disp, valid, scale, ox, oy = preprocess(img)
|
| out = dict(zip(out_names, sess.run(out_names, {in_name: x})))
|
|
|
|
|
| mask = out["vertex_mask"][0, 0] > 0.5
|
| net_corners = to_original(out["vertices"][0, 0][mask], scale, ox, oy)
|
|
|
| result = {"image": args.image, "size_wh": [img.shape[1], img.shape[0]],
|
| "path_A_model_polygon": np.round(net_corners, 2).tolist(),
|
| "path_A_active_vertices": int(mask.sum()),
|
| "note": "the model confidence is not calibrated: do not use it as a threshold"}
|
|
|
| if not args.only_net:
|
|
|
|
|
|
|
| g512 = cv2.cvtColor(disp, cv2.COLOR_BGR2GRAY)
|
| er = cv2.erode(valid.astype(np.uint8), np.ones((5, 5), np.uint8)).astype(bool)
|
| if args.no_canny:
|
| edge = out["edge_map"][0, 0] * er
|
| cfg = QuadConfig(nested=args.nested, max_quads=args.max_quads)
|
| else:
|
| edge = hybrid_edge_map(out["edge_map"][0, 0], g512, valid=er)
|
| cfg = QuadConfig.hybrid(nested=args.nested, max_quads=args.max_quads)
|
| quads = extract_quads(edge, cfg, gray=g512.astype(np.float32))
|
| result["path_B_rectangles"] = [
|
| {"corners_xy": to_original(q["corners"], scale, ox, oy).round(2).tolist(),
|
| "score": round(float(q["score"]), 3),
|
| "side_edge_coverage": [round(float(v), 2) for v in q["edge_cov"]]}
|
| for q in quads]
|
| result["path_B_count"] = len(quads)
|
|
|
| print(json.dumps(result, ensure_ascii=False, indent=2))
|
| Path(args.image).with_suffix(".quads.json").write_text(
|
| json.dumps(result, ensure_ascii=False, indent=2), encoding="utf-8")
|
|
|
|
|
| if __name__ == "__main__":
|
| main()
|
|
|
