File size: 5,075 Bytes
6464353 fb71420 6464353 fb71420 6464353 fb71420 6464353 fb71420 6464353 fb71420 6464353 fb71420 6464353 fb71420 6464353 fb71420 6464353 fb71420 6464353 fb71420 6464353 fb71420 6464353 fb71420 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 | """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})))
# ---- path A: a single polygon, strong on the training domain
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:
# ---- path B: edge_map + geometry (rectangles on real photographs)
# The model edge map lives in the 512 letterbox frame: the helper Canny must be computed on
# THAT frame, not on the original image (different areas and different scales).
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()
|