Files
AI_portal/dwg-counting/vector_match.py
2026-07-07 13:16:26 +02:00

218 lines
7.8 KiB
Python

"""Exact symbol counting by repeated-geometry matching.
The symbol is whatever the user boxes on the drawing — a set of DXF entities
(lines / arcs / circles / polylines). CAD copy-paste preserves geometry
exactly, so other occurrences are the SAME entity set under a rigid transform
(translation + rotation, optionally mirrored). We find them exactly:
1. Describe each template entity by a transform-invariant shape descriptor
plus a representative model-space point.
2. Pick a rare 'anchor' entity and the farthest 'orient' entity. Every
(anchor-candidate, orient-candidate) pair in the drawing whose shapes and
separation match the template fixes a candidate rigid transform.
3. Map all template points through the transform and verify — via a spatial
hash of the drawing's entities — that EVERY template entity is present.
4. Dedupe verified placements; report exact instance centres.
A placement is reported only if every entity of the symbol is found at the
transformed position, so there are no false positives and the count is exact.
"""
from __future__ import annotations
import logging
import math
from collections import defaultdict
from pathlib import Path
import ezdxf
logger = logging.getLogger(__name__)
TOL = 0.05 # absolute model-unit tolerance (exact copies match ~1e-6)
_KINDS = ("LINE", "CIRCLE", "ARC", "LWPOLYLINE", "POLYLINE")
def _pts(e) -> list[tuple[float, float]]:
t = e.dxftype()
if t == "LINE":
s, en = e.dxf.start, e.dxf.end
return [(s.x, s.y), (en.x, en.y)]
if t == "CIRCLE":
c = e.dxf.center
return [(c.x, c.y)]
if t == "ARC":
c, r = e.dxf.center, e.dxf.radius
a0, a1 = math.radians(e.dxf.start_angle), math.radians(e.dxf.end_angle)
return [(c.x, c.y),
(c.x + r * math.cos(a0), c.y + r * math.sin(a0)),
(c.x + r * math.cos(a1), c.y + r * math.sin(a1))]
if t in ("LWPOLYLINE", "POLYLINE"):
try:
return [(float(p[0]), float(p[1])) for p in e.get_points()]
except Exception:
return []
return []
def _descriptor(e) -> tuple | None:
"""Rotation/translation/mirror-invariant shape signature."""
t = e.dxftype()
if t == "LINE":
s, en = e.dxf.start, e.dxf.end
return ("L", round(math.hypot(en.x - s.x, en.y - s.y), 3))
if t == "CIRCLE":
return ("C", round(e.dxf.radius, 3))
if t == "ARC":
span = (e.dxf.end_angle - e.dxf.start_angle) % 360.0
return ("A", round(e.dxf.radius, 3), round(span, 2))
if t in ("LWPOLYLINE", "POLYLINE"):
p = _pts(e)
if len(p) < 2:
return None
segs = tuple(round(math.hypot(p[i + 1][0] - p[i][0],
p[i + 1][1] - p[i][1]), 3)
for i in range(len(p) - 1))
closed = bool(getattr(e, "closed", False) or
(getattr(e.dxf, "flags", 0) & 1))
return ("P", min(segs, tuple(reversed(segs))), closed, len(p))
return None
def _primary(e) -> tuple[float, float]:
"""One representative point per entity (centroid of its key points)."""
p = _pts(e)
if not p:
b = e.dxf.start if e.dxftype() == "LINE" else e.dxf.center
return (b.x, b.y)
if e.dxftype() == "CIRCLE":
return p[0]
return (sum(x for x, _ in p) / len(p), sum(y for _, y in p) / len(p))
def _iter_entities(msp):
"""Yield modelspace entities, expanding INSERTs into their world-space
geometry so symbols built from blocks are visible to the matcher."""
for e in msp:
t = e.dxftype()
if t == "INSERT":
try:
for ve in e.virtual_entities():
if ve.dxftype() in _KINDS:
yield ve
except Exception as exc:
logger.debug("virtual_entities failed on INSERT %s: %s",
getattr(e.dxf, "name", "?"), exc)
continue
if t in _KINDS:
yield e
def _collect(msp, bbox=None):
out = []
for e in _iter_entities(msp):
d = _descriptor(e)
if d is None:
continue
px, py = _primary(e)
if bbox is not None:
x0, y0, x1, y1 = bbox
if not (x0 <= px <= x1 and y0 <= py <= y1):
continue
out.append((d, (px, py)))
return out
def _snap(x, y):
return (round(x / TOL), round(y / TOL))
def _map(anchor_q, a_p, cos, sin, mirror, x, y):
"""Map model point (x,y): express relative to template anchor a_p, mirror
in x if needed, rotate, translate to the drawing anchor anchor_q."""
rx, ry = x - a_p[0], y - a_p[1]
if mirror:
rx = -rx
return (anchor_q[0] + cos * rx - sin * ry,
anchor_q[1] + sin * rx + cos * ry)
def match_region(dxf_path: Path, bbox: tuple) -> dict:
"""bbox = (x0,y0,x1,y1) in MODEL coords. Returns
{"count", "instances":[{"x","y"}], "template_entities"}."""
doc = ezdxf.readfile(str(dxf_path))
msp = doc.modelspace()
template = _collect(msp, bbox)
if not template:
return {"count": 0, "instances": [], "template_entities": 0,
"error": "Ve vybrané oblasti nejsou žádné rozpoznatelné křivky."}
drawing = _collect(msp, None)
by_desc: dict[tuple, list] = defaultdict(list)
grid: dict = defaultdict(list)
for d, p in drawing:
by_desc[d].append(p)
grid[(d, _snap(*p))].append(p)
def present(d, qx, qy) -> bool:
sx, sy = _snap(qx, qy)
for dx in (-1, 0, 1):
for dy in (-1, 0, 1):
for gp in grid.get((d, (sx + dx, sy + dy)), ()):
if abs(gp[0] - qx) <= TOL and abs(gp[1] - qy) <= TOL:
return True
return False
# Anchor = template entity with the fewest drawing candidates.
anchor_i = min(range(len(template)),
key=lambda i: len(by_desc.get(template[i][0], [])) or 1 << 30)
a_d, a_p = template[anchor_i]
def dist(p, q):
return math.hypot(p[0] - q[0], p[1] - q[1])
orient_i = max((i for i in range(len(template)) if i != anchor_i),
key=lambda i: dist(template[i][1], a_p), default=None)
instances: list[tuple[float, float]] = []
seen = set()
cx = sum(p[0] for _, p in template) / len(template)
cy = sum(p[1] for _, p in template) / len(template)
dedupe_cell = max(TOL * 20, dist((cx, cy), a_p) * 0.25 or TOL * 20)
def record(anchor_q, c, s, m):
for d, (px, py) in template:
qx, qy = _map(anchor_q, a_p, c, s, m, px, py)
if not present(d, qx, qy):
return
ix, iy = _map(anchor_q, a_p, c, s, m, cx, cy)
k = (round(ix / dedupe_cell), round(iy / dedupe_cell))
if k not in seen:
seen.add(k)
instances.append((ix, iy))
if orient_i is None:
# Single-entity symbol → translation only.
for q in by_desc.get(a_d, []):
record(q, 1.0, 0.0, False)
else:
o_d, o_p = template[orient_i]
sep = dist(a_p, o_p)
for aq in by_desc.get(a_d, []):
for oq in by_desc.get(o_d, []):
if abs(dist(aq, oq) - sep) > TOL:
continue
for mirror in (False, True):
vx, vy = o_p[0] - a_p[0], o_p[1] - a_p[1]
if mirror:
vx = -vx
th = (math.atan2(oq[1] - aq[1], oq[0] - aq[0]) -
math.atan2(vy, vx))
record(aq, math.cos(th), math.sin(th), mirror)
logger.info("vector_match: %d template entities → %d instances",
len(template), len(instances))
return {"count": len(instances),
"instances": [{"x": x, "y": y} for x, y in instances],
"template_entities": len(template)}