"""Shape-based symbol counting via contour matching. Why not cv2.matchTemplate: CAD symbols are usually thin-line drawings on a white background (~10-20% ink). Normalized cross-correlation gives spuriously high scores for any sparse-ink region (e.g. wall edges), producing false positives everywhere. Approach used here: 1. Binarize template + drawing to "ink" maps. 2. Find external contours in both. 3. Use the template's main contour as a shape reference. 4. For every contour in the drawing, compare via cv2.matchShapes (Hu moments — invariant to scale, rotation, translation). 5. Filter by area ratio (similar size to template, allowing ±factor). 6. Keep contours below a shape-distance threshold. This is the classic approach to CAD symbol takeoff. """ import logging from pathlib import Path from typing import Iterable import cv2 import numpy as np logger = logging.getLogger(__name__) # Template matching: higher = stricter. 0.65 is permissive, 0.85 strict. DEFAULT_THRESHOLD = 0.7 import threading import time # GPU dispatch — if torch + CUDA are available we delegate to the GPU # implementation, which is mathematically equivalent but ~20-50× faster. _GPU_AVAILABLE = False try: import torch if torch.cuda.is_available(): from counting_gpu import count_template as _count_template_gpu _GPU_AVAILABLE = True logger.info("CUDA detected: %s — GPU template matching enabled", torch.cuda.get_device_name(0)) else: logger.info("CUDA not available — using CPU template matching") except Exception as exc: logger.info("torch unavailable (%s) — using CPU template matching", exc) # Serialize ALL GPU work in this process. Two concurrent count requests each # allocate large FFT tensors for an ~8000px drawing; running them at once # exhausts VRAM and the second one dies with a CUDA error. The lock makes # GPU counts run one-at-a-time (each is only a few seconds anyway). _GPU_LOCK = threading.Lock() def _try_gpu_count(template_path, drawing_path, threshold, exclude_box, mirror) -> dict | None: """Run the GPU matcher under the global lock, with one retry. Returns the result dict on success, or None if the GPU stays unavailable (the caller then falls back to the CPU implementation — slower but it always works). Transient `cudaErrorDevicesUnavailable` happens when another process (e.g. Ollama loading an embedding model for the Firemní asistent) is mid-allocation; a short wait + retry usually clears it. """ with _GPU_LOCK: for attempt in (1, 2): try: return _count_template_gpu( template_path=template_path, drawing_path=drawing_path, threshold=threshold, exclude_box=exclude_box, mirror=mirror, ) except Exception as exc: # noqa: BLE001 — any GPU failure → CPU try: torch.cuda.empty_cache() except Exception: pass if attempt == 1: logger.warning( "GPU count failed (%s) — retrying in 2s", exc) time.sleep(2.0) continue logger.error( "GPU still unavailable after retry (%s) — CPU fallback", exc) return None def _prep(img_path: Path) -> np.ndarray: """Binarize to 'ink vs not-ink'. No dilation — for contour-based matching we need lines to stay narrow so distinct symbols don't merge into one mega-contour through CAD's dense linework. """ arr = cv2.imread(str(img_path), cv2.IMREAD_GRAYSCALE) if arr is None: raise RuntimeError(f"Could not load {img_path}") # 245 — keep in sync with INK_THRESHOLD in counting_gpu.py (faint # templates lose all ink below ~245; see note there). _, ink = cv2.threshold(arr, 245, 255, cv2.THRESH_BINARY_INV) return ink def _crop_to_content(template: np.ndarray, bg_threshold: int = 240) -> np.ndarray: """Crop a template to its non-background bounding box. After _prep the template is BINARY (ink=255, bg=0). So 'content' is everything with value > 0. The bg_threshold param is kept for the grayscale case for backward compat. """ if template.max() <= 1 or template.min() == 0 and template.max() == 255: # Binary map: any non-zero pixel is ink mask = template > 0 else: mask = template < bg_threshold if not mask.any(): return template ys, xs = np.where(mask) y0, y1 = ys.min(), ys.max() + 1 x0, x1 = xs.min(), xs.max() + 1 pad = 2 y0 = max(0, y0 - pad) x0 = max(0, x0 - pad) y1 = min(template.shape[0], y1 + pad) x1 = min(template.shape[1], x1 + pad) return template[y0:y1, x0:x1] def _nms(boxes: list[tuple], scores: list[float], overlap_thresh: float = 0.3) -> list[int]: """Non-max suppression. Returns indices of kept boxes.""" if not boxes: return [] boxes_arr = np.array(boxes, dtype=np.float32) x1 = boxes_arr[:, 0] y1 = boxes_arr[:, 1] x2 = boxes_arr[:, 2] y2 = boxes_arr[:, 3] areas = (x2 - x1) * (y2 - y1) order = np.argsort(scores)[::-1] keep = [] while order.size > 0: i = order[0] keep.append(int(i)) xx1 = np.maximum(x1[i], x1[order[1:]]) yy1 = np.maximum(y1[i], y1[order[1:]]) xx2 = np.minimum(x2[i], x2[order[1:]]) yy2 = np.minimum(y2[i], y2[order[1:]]) w = np.maximum(0.0, xx2 - xx1) h = np.maximum(0.0, yy2 - yy1) inter = w * h union = areas[i] + areas[order[1:]] - inter iou = inter / np.maximum(union, 1e-6) order = order[1:][iou <= overlap_thresh] return keep MAX_DRAWING_PX = 9000 # effectively no downscale for typical A1/A0 def debug_template(template_path: Path, drawing_path: Path) -> dict: """Return diagnostics for tuning a symbol template.""" template = _prep(template_path) template_cropped = _crop_to_content(template) drawing = _prep(drawing_path) info = { "template_size": list(template.shape), "template_cropped_size": list(template_cropped.shape), "template_ink_pixels": int((template_cropped > 0).sum() if template_cropped.max() == 255 and template_cropped.min() == 0 else (template_cropped < 240).sum()), "template_total_pixels": int(template_cropped.size), "drawing_size": list(drawing.shape), } if min(template_cropped.shape) < 8: info["error"] = "template too small after content crop" return info # Scale sweep — finds the size at which template matches best scale_scan = [] best_overall = -1.0 for scale in (0.15, 0.25, 0.35, 0.5, 0.7, 0.85, 1.0, 1.2, 1.5, 2.0): nw = max(8, int(template_cropped.shape[1] * scale)) nh = max(8, int(template_cropped.shape[0] * scale)) if nh > drawing.shape[0] or nw > drawing.shape[1]: continue tmpl = cv2.resize(template_cropped, (nw, nh), interpolation=cv2.INTER_AREA) res = cv2.matchTemplate(drawing, tmpl, cv2.TM_CCOEFF_NORMED) m = float(res.max()) scale_scan.append({ "scale": scale, "template_px": [nh, nw], "max_score": round(m, 3), "count_at_0.7": int((res >= 0.7).sum()), "count_at_0.6": int((res >= 0.6).sum()), }) if m > best_overall: best_overall = m info["max_score"] = best_overall info["scale_scan"] = scale_scan # Threshold counts at scale=1.0 for reference result = cv2.matchTemplate(drawing, template_cropped, cv2.TM_CCOEFF_NORMED) info["matches_at_threshold"] = { "0.60": int((result >= 0.60).sum()), "0.70": int((result >= 0.70).sum()), "0.75": int((result >= 0.75).sum()), "0.80": int((result >= 0.80).sum()), "0.85": int((result >= 0.85).sum()), "0.90": int((result >= 0.90).sum()), } return info def _merge_template_contour(ink: np.ndarray) -> np.ndarray | None: """Combine all strokes of a template into one contour via closing. Returns the biggest connected component. The template is small enough that closing safely merges the line + arc of symbols like '-C' into one shape without affecting matching against the drawing. """ closed = cv2.morphologyEx(ink, cv2.MORPH_CLOSE, np.ones((5, 5), np.uint8)) contours, _ = cv2.findContours(closed, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE) if not contours: return None return max(contours, key=cv2.contourArea) def count_template( template_path: Path, drawing_path: Path, threshold: float = DEFAULT_THRESHOLD, rotations: Iterable[int] = (0, 90, 180, 270), scales: Iterable[float] = (0.6, 0.8, 1.0, 1.25, 1.5), exclude_box: tuple | None = None, area_tolerance: float = 200.0, mirror: bool = True, ) -> dict: """Multi-scale, multi-rotation template matching (TM_SQDIFF_NORMED). With mirror=True the template is also tested flipped horizontally — at each rotation — so we catch mirrored instances (e.g. a socket facing left vs facing right). """ if _GPU_AVAILABLE: # GPU path defaults to a denser sweep (8 rotations × 7 scales) since # the per-variant cost is negligible. On persistent GPU unavailability # this returns None and we fall through to the CPU path below. gpu_result = _try_gpu_count( template_path, drawing_path, threshold, exclude_box, mirror) if gpu_result is not None: return gpu_result logger.warning("Falling back to CPU template matching for this job") template = _prep(template_path) template = _crop_to_content(template) drawing = _prep(drawing_path) coord_scale = 1.0 if max(drawing.shape) > MAX_DRAWING_PX: coord_scale = max(drawing.shape) / MAX_DRAWING_PX new_w = int(drawing.shape[1] / coord_scale) new_h = int(drawing.shape[0] / coord_scale) drawing = cv2.resize(drawing, (new_w, new_h), interpolation=cv2.INTER_AREA) if exclude_box is not None: ex_x, ex_y, ex_w, ex_h = exclude_box ex_x = int(ex_x / coord_scale) ex_y = int(ex_y / coord_scale) ex_w = int(ex_w / coord_scale) ex_h = int(ex_h / coord_scale) h, w = drawing.shape ex_x = max(0, min(w, ex_x)) ex_y = max(0, min(h, ex_y)) ex_x2 = max(0, min(w, ex_x + ex_w)) ex_y2 = max(0, min(h, ex_y + ex_h)) drawing[ex_y:ex_y2, ex_x:ex_x2] = 0 logger.info("Masked legend region (%d,%d,%d,%d)", ex_x, ex_y, ex_w, ex_h) if min(template.shape) < 8 or int((template > 0).sum()) < 5: logger.warning("Template too small / empty after preprocessing") return {"count": 0, "matches": [], "threshold_used": threshold} all_boxes: list[tuple] = [] all_scores: list[float] = [] # Build the variants we'll search with: rotations × (original, mirrored) variants = [(template, False)] if mirror: variants.append((cv2.flip(template, 1), True)) def _rotate(img, angle_deg): if angle_deg == 0: return img if angle_deg in (90, 180, 270): return np.rot90(img, k=angle_deg // 90) # Arbitrary angle — rotate with bbox expansion + white fill h, w = img.shape[:2] center = (w / 2, h / 2) M = cv2.getRotationMatrix2D(center, angle_deg, 1.0) cos = abs(M[0, 0]); sin = abs(M[0, 1]) new_w = int(h * sin + w * cos) new_h = int(h * cos + w * sin) M[0, 2] += new_w / 2 - center[0] M[1, 2] += new_h / 2 - center[1] return cv2.warpAffine(img, M, (new_w, new_h), flags=cv2.INTER_NEAREST, borderValue=0) # Build the full job list (one entry per rot×scale×mirror) jobs_list = [] for variant, _ in variants: for rot in rotations: base = _rotate(variant, rot) for scale in scales: new_w = max(8, int(base.shape[1] * scale)) new_h = max(8, int(base.shape[0] * scale)) if new_h > drawing.shape[0] or new_w > drawing.shape[1]: continue tmpl = cv2.resize(base, (new_w, new_h), interpolation=cv2.INTER_AREA) jobs_list.append((tmpl, new_w, new_h)) def _run(args): tmpl, new_w, new_h = args sq = cv2.matchTemplate(drawing, tmpl, cv2.TM_SQDIFF_NORMED) cc = cv2.matchTemplate(drawing, tmpl, cv2.TM_CCOEFF_NORMED) sim = np.maximum(1.0 - sq, cc) ys, xs = np.where(sim >= threshold) out = [] for y, x in zip(ys, xs): out.append((float(x), float(y), float(x + new_w), float(y + new_h), float(sim[y, x]))) return out # Threading: cv2.matchTemplate releases the GIL, so threads give real speedup import concurrent.futures with concurrent.futures.ThreadPoolExecutor(max_workers=4) as ex: for result in ex.map(_run, jobs_list): for x0, y0, x1, y1, score in result: all_boxes.append((x0, y0, x1, y1)) all_scores.append(score) if not all_boxes: return {"count": 0, "matches": [], "threshold_used": threshold} keep = _nms(all_boxes, all_scores, overlap_thresh=0.3) matches = [] for i in keep: x0, y0, x1, y1 = all_boxes[i] matches.append({ "x": int(x0 * coord_scale), "y": int(y0 * coord_scale), "w": int((x1 - x0) * coord_scale), "h": int((y1 - y0) * coord_scale), "score": round(all_scores[i], 3), }) return {"count": len(matches), "matches": matches, "threshold_used": threshold}