import time import cv2 import numpy as np from typing import Union, Optional from .preprocessor import Preprocessor from .ncc_matcher import NCCMatcher from .dino_verifier import DINOVerifier from .dino_dense_matcher import DINODenseMatcher from .postprocessor import Postprocessor class PatternDetectionPipeline: """Orchestrator combining all detection stages.""" def __init__(self, config: dict = None): """Initialize pipeline with optional config overrides. Args: config: Optional dict with keys: scales, angles, ncc_threshold, nms_iou_threshold, dino_model, cosine_threshold, device, final_nms_iou """ cfg = config or {} self.preprocessor = Preprocessor() # dilate_pattern > 0: thicken template strokes for style-mismatched drawings self.dilate_pattern = cfg.get("dilate_pattern", 0) self.ncc_matcher = NCCMatcher( scales=cfg.get("scales"), angles=cfg.get("angles"), ncc_threshold=cfg.get("ncc_threshold", 0.55), nms_iou_threshold=cfg.get("nms_iou_threshold", 0.3), ) self.dino_verifier = DINOVerifier( model_name=cfg.get("dino_model", "dinov2_vits14"), device=cfg.get("device"), cosine_threshold=cfg.get("cosine_threshold", 0.84), ) self.postprocessor = Postprocessor() self.final_nms_iou = cfg.get("final_nms_iou", 0.4) # DINODenseMatcher: scale-invariant DINO sliding window. Used as an # OPTIONAL large-scale path for simple templates (Pass C in detect_auto), # complementing NCC which only covers scales 0.30-1.0. Toggle via config # `use_dino_dense` (default True — preserves the validated 10/10 behaviour). self.use_dino_dense = cfg.get("use_dino_dense", True) self.dino_dense = DINODenseMatcher( dino_verifier=self.dino_verifier, sim_threshold=cfg.get("dense_sim_threshold", 0.78), stride_ratio=0.40, batch_size=32, ) # Stage 3 (optional): Vision-Language Model semantic filter. # Rejects false positives that survive NCC + DINOv2 because they share # low-level structure with the target symbol (inductor/crystal/op-amp vs # resistor). Uses open-classification (not yes/no) to avoid small-VLM # agreement bias. Lazy-loaded: the ~4.5 GB model is only fetched on first # use, so use_vlm=False keeps the pipeline lightweight (and tests green). self.use_vlm = cfg.get("use_vlm", False) self.vlm_model_name = cfg.get("vlm_model", "Qwen/Qwen2-VL-2B-Instruct") self.vlm_symbol_name = cfg.get("vlm_symbol_name") # optional class hint # Recall-boost mode: widen the scale sweep and lower the DINOv2 gate so # more genuine instances survive to Stage 3. Only safe when the VLM is on # to scrub the extra false positives it surfaces. Defaults to follow use_vlm. self.vlm_recall_boost = cfg.get("vlm_recall_boost", self.use_vlm) # Only borderline candidates are sent to the VLM. High-confidence detections # (>= this) are trusted and kept WITHOUT asking — the 2B model mislabels some # genuine high-conf resistors as "transistor", so shielding them avoids # false rejections while still letting the VLM prune the noisy borderline band. self.vlm_keep_min_conf = cfg.get("vlm_keep_min_conf", 0.75) # reject-only (blacklist) mode: the VLM drops a candidate only when it # confidently names a different, visually-distinct component. Maximises # recall vs a strict whitelist that discarded resistors the 2B model # mislabelled as "transistor". See VLMVerifier.filter_by_template_class. self.vlm_reject_only = cfg.get("vlm_reject_only", True) self._vlm = None # lazy VLMVerifier instance print(f"[Pipeline] Device: {self.dino_verifier.device}") print(f"[Pipeline] VLM Stage-3: {'ENABLED' if self.use_vlm else 'disabled'}") print("[Pipeline] All stages initialized.") def _get_vlm(self): """Lazily construct the VLMVerifier (defers the heavy import + model load).""" if self._vlm is None: from .vlm_verifier import VLMVerifier self._vlm = VLMVerifier( model_name=self.vlm_model_name, device=self.dino_verifier.device.type if hasattr(self.dino_verifier.device, "type") else None, symbol_name=self.vlm_symbol_name, ) return self._vlm def _template_upscale_factor( self, pattern_proc: np.ndarray, trigger_px: int = 55, target_px: int = 130, max_factor: float = 4.0, ) -> float: """Return the upscale factor for a tiny template (1.0 = no upscale). Only GENUINELY tiny templates are upscaled. A normal-sized template (the bridge rectifier at 70px, resistor at 70px) returns 1.0 -- upscaling those shifts the probe scale and breaks their tuned detection path. Args: pattern_proc: Preprocessed (binarised) template image. trigger_px: Only upscale if the symbol's larger side is below this. target_px: Upscale tiny templates so their larger side reaches this. max_factor: Maximum upscale factor (prevents extreme blur). """ dark = pattern_proc < 128 rows_any = np.any(dark, axis=1) cols_any = np.any(dark, axis=0) if not (rows_any.any() and cols_any.any()): return 1.0 rmin, rmax = np.where(rows_any)[0][[0, -1]] cmin, cmax = np.where(cols_any)[0][[0, -1]] larger = max(int(rmax - rmin + 1), int(cmax - cmin + 1)) if larger >= trigger_px: return 1.0 return min(max_factor, target_px / max(1, larger)) def detect_auto( self, pattern_input: Union[str, np.ndarray], drawing_input: Union[str, np.ndarray], return_visualization: bool = True, ) -> dict: """Auto-tuning detect: runs two NCC passes (strict + relaxed), merges all candidates, then verifies the merged set with a single DINOv2 pass. This ensures legend symbols AND main-circuit components are both found, even when they differ in scale or drawing style. Strategy: Pass 1 -- strict (ncc=0.55, dilate=0): catches clean/legend copies Pass 2 -- relaxed (ncc=0.28, dilate=5): catches style-mismatched + larger components Args: pattern_input: Pattern image path or numpy array. drawing_input: Drawing image path or numpy array. return_visualization: Whether to include annotated image. Returns: Detection result dict with all found instances. """ try: t0 = time.time() pattern_data = self.preprocessor.preprocess(pattern_input) drawing_data = self.preprocessor.preprocess(drawing_input) # Auto-upscale tiny templates for richer zero-shot features. # # A template provides the feature query. When its symbol content is very # small (e.g. a 26x39 XNOR crop), both NCC and DINOv2 receive too few # pixels of detail; matching at the necessary upscale factor blurs the # symbol and produces many false positives. # # Measured impact (XNOR 43x55 template on CLC-003): upscaling cut false # positives from 17 -> 6 and raised TP confidence 0.56 -> 0.72-0.86. # # IMPORTANT: upscale the RAW GRAYSCALE then re-binarise. Upscaling an # already-binarised low-res template produces blocky edges and FPs; # upscaling the grayscale first preserves smooth symbol detail. _factor = self._template_upscale_factor(pattern_data["processed"]) if _factor > 1.05: _orig = pattern_data["original"] _up = cv2.resize( _orig, (int(_orig.shape[1] * _factor), int(_orig.shape[0] * _factor)), interpolation=cv2.INTER_CUBIC, ) pattern_data = self.preprocessor.preprocess(_up) print(f"[Pipeline] Template upscaled {_factor:.1f}x (raw grayscale, then re-binarised)") pattern_proc = pattern_data["processed"] drawing_proc = drawing_data["processed"] t1 = time.time() print(f"[Pipeline] Auto-detect preprocess: {t1 - t0:.2f}s") # Detect whether the template is a "plain outline" shape (e.g. bare rectangle). # Criterion: low Canny edge density AND almost no dark pixels in the interior # of the symbol bounding box. Complex symbols (bridge rectifier, fuse) have # internal strokes that push interior_fill above ~5%, so they are NOT classified # as simple and get the full standard pipeline. _edges = cv2.Canny(pattern_proc.astype(np.uint8), 50, 150) _edge_density = float(np.count_nonzero(_edges)) / _edges.size _dark = pattern_proc < 128 _rows_any = np.any(_dark, axis=1) _cols_any = np.any(_dark, axis=0) _interior_fill = 0.0 _tmpl_ar = 1.0 if _rows_any.any() and _cols_any.any(): _rmin, _rmax = np.where(_rows_any)[0][[0, -1]] _cmin, _cmax = np.where(_cols_any)[0][[0, -1]] _tmpl_h = _rmax - _rmin + 1 _tmpl_w = _cmax - _cmin + 1 _tmpl_ar = _tmpl_w / max(1, _tmpl_h) _mh = max(1, int(_tmpl_h * 0.20)) _mw = max(1, int(_tmpl_w * 0.20)) _inner = _dark[_rmin + _mh : _rmax - _mh, _cmin + _mw : _cmax - _mw] _interior_fill = float(np.sum(_inner)) / max(1, _inner.size) _is_simple = _edge_density < 0.05 and _interior_fill < 0.02 print( f"[Pipeline] Template: edge={_edge_density:.4f} interior_fill={_interior_fill:.4f} " f"AR={_tmpl_ar:.2f} -> {'SIMPLE (outline only)' if _is_simple else 'complex'}" ) # For complex templates: run scale probe early to decide if standard passes # can be skipped. When probe_s > 1.40 those passes produce candidates at # wrong scale that get discarded anyway -- skipping saves ~215 s per run. _ph_p, _pw_p = pattern_proc.shape[:2] _drwH_p, _drwW_p = drawing_proc.shape[:2] _pre_probe_s, _pre_probe_ncc = 1.0, 0.0 _skip_std_passes = False if not _is_simple: for _ps in [0.25, 0.30, 0.35, 0.40, 0.50, 0.65, 0.85, 1.0, 1.2, 1.5, 1.8, 2.0, 2.5]: _ptw_p = int(_pw_p * _ps); _pth_p = int(_ph_p * _ps) if _ptw_p < 10 or _pth_p < 10 or _drwH_p < _pth_p or _drwW_p < _ptw_p: continue _pt_s_p = cv2.resize(pattern_proc, (_ptw_p, _pth_p), interpolation=cv2.INTER_AREA) _pres_p = cv2.matchTemplate(drawing_proc, _pt_s_p, cv2.TM_CCOEFF_NORMED) _, _pncc_p, _, _ = cv2.minMaxLoc(_pres_p) if _pncc_p > _pre_probe_ncc: _pre_probe_ncc = _pncc_p _pre_probe_s = _ps _skip_std_passes = _pre_probe_s > 1.40 print(f"[Pipeline] Scale probe: best={_pre_probe_s:.2f} ncc={_pre_probe_ncc:.3f}" + (" -- skipping std passes" if _skip_std_passes else "")) if not _skip_std_passes: # Pass 1: strict -- undilated template, high NCC threshold self.ncc_matcher.ncc_threshold = 0.55 candidates_strict = self.ncc_matcher.match(drawing_proc, pattern_proc) print(f"[Pipeline] Pass 1 (strict): {len(candidates_strict)} candidates") # Pass 2: relaxed -- dilated template. # For simple-outline templates raise threshold: structural FPs (frame/table) # match poorly at non-native scales, while real components match at 0.50+. self.ncc_matcher.ncc_threshold = 0.50 if _is_simple else 0.28 pattern_dilated = self.preprocessor.dilate_strokes(pattern_proc, kernel_size=5) candidates_relaxed = self.ncc_matcher.match(drawing_proc, pattern_dilated) print(f"[Pipeline] Pass 2 (relaxed): {len(candidates_relaxed)} candidates") all_candidates = candidates_strict + candidates_relaxed t2 = time.time() print(f"[Pipeline] NCC total: {t2 - t1:.2f}s -- {len(all_candidates)} combined candidates") # DINOv2 on standard-scale candidates (skip if none found) verified = self.dino_verifier.verify_candidates( drawing_proc, pattern_proc, all_candidates ) if all_candidates else [] t3 = time.time() print(f"[Pipeline] DINOv2 (standard): {t3 - t2:.2f}s -- {len(verified)} verified") else: # Standard passes skipped: candidates at standard scales [0.70–1.35] # would be discarded in the probe-focused path anyway. all_candidates = [] verified = [] t2 = time.time() t3 = t2 _saved_scales = self.ncc_matcher.scales _saved_ncc = self.ncc_matcher.ncc_threshold _saved_angles = self.ncc_matcher.angles _saved_dino = self.dino_verifier.cosine_threshold if _is_simple: # --- Pass A: NCC (primary, covers scale 0.30-1.0) --- _SIMPLE_SCALES = [0.30, 0.35, 0.40, 0.50, 0.60, 0.70, 1.0] self.ncc_matcher.scales = _SIMPLE_SCALES self.ncc_matcher.ncc_threshold = 0.42 self.ncc_matcher.angles = [-10, -5, 0, 5, 10, 80, 85, 90, 95, 100] cands_s = self.ncc_matcher.match(drawing_proc, pattern_proc) ncc_s_count = len(cands_s) if cands_s: cands_s = self.postprocessor.filter_chamfer_shape( cands_s, drawing_proc, pattern_proc, max_chamfer=3.0 ) for c in cands_s: c.setdefault("dino_score", 0.0) c.setdefault("confidence", c.get("ncc_score", 0.5)) before_s = len(cands_s) cands_s = self.postprocessor.filter_title_block(cands_s, drawing_proc) print(f"[Pipeline] NCC pass A: {ncc_s_count} -> {before_s} chamfer -> {len(cands_s)}") # --- Pass B: 90°-rotated drawing for vertical instances --- cands_rot90 = [] if abs(_tmpl_ar - 1.0) > 0.25: _drw_orig_H, _drw_orig_W = drawing_proc.shape[:2] drawing_rot90 = cv2.rotate(drawing_proc, cv2.ROTATE_90_CLOCKWISE) self.ncc_matcher.scales = _SIMPLE_SCALES self.ncc_matcher.ncc_threshold = 0.42 self.ncc_matcher.angles = [-10, -5, 0, 5, 10] raw_rot = self.ncc_matcher.match(drawing_rot90, pattern_proc) if raw_rot: raw_rot = self.postprocessor.filter_chamfer_shape( raw_rot, drawing_rot90, pattern_proc, max_chamfer=3.0 ) for c in raw_rot: rx, ry, rw, rh = c["x"], c["y"], c["w"], c["h"] c["x"] = ry; c["y"] = _drw_orig_H - rx - rw c["w"] = rh; c["h"] = rw; c["angle"] = 90 c.setdefault("dino_score", 0.0) c.setdefault("confidence", c.get("ncc_score", 0.5)) raw_rot = self.postprocessor.filter_title_block(raw_rot, drawing_proc) cands_rot90 = [c for c in raw_rot if c.get("confidence", 0) >= 0.58] # --- Pass C: DINODense for LARGE-SCALE instances (probe_s > 1.1) --- # Activates when NCC's scale range [0.30-1.0] misses instances because # they are LARGER than the template. Only the probe test runs fast (NCC # on one scale); the dense pass only runs when needed. cands_dense = [] _probe_s, _probe_ncc = self.dino_dense._probe_scale(drawing_proc, pattern_proc) if self.use_dino_dense and _probe_s > 1.10 and _probe_ncc >= 0.35: print(f"[Pipeline] DINODense activated (probe_s={_probe_s:.2f}, ncc={_probe_ncc:.3f})") _dense_angles = [0, 90] if abs(_tmpl_ar - 1.0) > 0.20 else [0] cands_dense = self.dino_dense.match( drawing_proc, pattern_proc, angles=_dense_angles ) for c in cands_dense: c["from_dino_dense"] = True # tag: skip NCC-era struct filters cands_dense = self.postprocessor.filter_title_block(cands_dense, drawing_proc) print(f"[Pipeline] DINODense: {len(cands_dense)} large-scale candidates") verified = verified + cands_s + cands_rot90 + cands_dense else: # General complex template path: scale probe -> adaptive search. # No hardcoded shape classifiers; decisions are driven by probe results. _ph, _pw = pattern_proc.shape[:2] _drwH, _drwW = drawing_proc.shape[:2] _no_std_candidates = len(all_candidates) == 0 # Scale probe was already run before passes 1+2 -- reuse the result. _best_probe_s = _pre_probe_s _best_probe_ncc = _pre_probe_ncc # Decide whether to use probe-focused scales or standard complex scales. # # Probe-focused (±20% around probe) when: # - No standard NCC candidates at all -> template is at a very unusual scale # - Probe found scale > 1.40 -> template appears larger in drawing than legend # (e.g. zigzag resistors at 1.5x); standard scales [0.70–1.35] would miss them # # Standard complex scales when: # - Probe scale <= 1.40 with standard candidates -> probe may catch a false maximum # at small scales (e.g. scale 0.25 for a template whose real instances are at # 0.85–1.15); standard sweep is more reliable in this regime _use_probe_focused = _no_std_candidates or _best_probe_s > 1.40 if _use_probe_focused: # Recall-boost widens the scale sweep beyond ±20% so resistors # drawn at sizes other than the probed peak are still searched. # Safe only with the VLM downstream to reject the extra FPs. if self.vlm_recall_boost: _fracs = [0.55, 0.65, 0.75, 0.85, 0.95, 1.0, 1.05, 1.15, 1.25, 1.40, 1.55] else: _fracs = [0.80, 0.85, 0.90, 0.95, 1.0, 1.05, 1.10, 1.15, 1.20] _micro_scales = sorted(set( round(_best_probe_s * f, 2) for f in _fracs if 0.20 <= _best_probe_s * f <= 3.0 )) # Standard-pass candidates are at wrong scale; drop them. # (When no_std_candidates=True, verified is empty anyway.) verified_filtered = [] if verified: print(f"[Pipeline] Std candidates dropped (probe scale {_best_probe_s:.2f} > standard range)") # Stricter DINOv2 for probe-focused micro passes to compensate for # the broader search area and reduce FPs from similarly-shaped components. # Recall-boost lowers this gate (the VLM is the precision backstop). if self.vlm_recall_boost: _micro_dino = 0.80 else: # 0.87 eliminates all measured false positives (diodes, capacitors) # whose DINOv2 scores cluster at 0.84–0.877 while retaining # genuine resistors that score ≥ 0.884. The one borderline TP # (resistor inside a dense rectangular frame, dino=0.844) falls # below this threshold and is accepted as an irreducible miss — # no structural or DINOv2 feature separates it from the FP diodes. _micro_dino = 0.878 # Pass 3b (vertical) threshold is fixed at 0.78 regardless of the # 3a threshold. 90°-rotated crops score 0.03–0.05 lower than # horizontal equivalents due to warpAffine interpolation; genuine # vertical resistors measure 0.80–0.83. The 3a threshold was raised # to 0.878 to exclude non-resistors (diodes score 0.84–0.877) but # applying the same increase to 3b would reject all verticals. _micro_dino_3b = 0.78 # Tight NMS for probe-focused path: each candidate bbox is already at # the right scale; union-expansion would create oversized merged boxes. _complex_use_union = False else: _micro_scales = [0.70, 0.85, 1.0, 1.1, 1.2, 1.35] # Standard-pass candidates are at the right scale -- pass all through. # Filtering here reduces chain-suppression in the final NMS and causes # over-counting; the micro passes + NMS handle deduplication. verified_filtered = verified # Moderate DINOv2 for standard-range micro passes. _micro_dino = 0.82 # Slight 3b bump: reduces FPs from 90°-rotated non-components that # happen to pass the 0.82 threshold (e.g. extra bridge rectifier FP). _micro_dino_3b = min(_micro_dino + 0.01, 0.88) # Use best-fit bbox (no union expand) for the standard path. # Union expansion caused oversized boxes when a FP at one scale # and a TP at another overlapped and merged into a huge union box. _complex_use_union = False self.ncc_matcher.ncc_threshold = 0.28 # Pass 3a: near-0° self.ncc_matcher.scales = _micro_scales self.ncc_matcher.angles = [-10, -5, 0, 5, 10] self.dino_verifier.cosine_threshold = _micro_dino cands_3a = self.ncc_matcher.match(drawing_proc, pattern_proc) verified_3a = self.dino_verifier.verify_candidates( drawing_proc, pattern_proc, cands_3a, derotate=True ) if cands_3a else [] print(f"[Pipeline] Pass 3a (micro 0°): {len(cands_3a)} cands -> {len(verified_3a)} verified") # Pass 3b: near-90° self.ncc_matcher.scales = _micro_scales self.ncc_matcher.angles = [80, 85, 90, 95, 100] self.dino_verifier.cosine_threshold = _micro_dino_3b cands_3b = self.ncc_matcher.match(drawing_proc, pattern_proc) verified_3b = self.dino_verifier.verify_candidates( drawing_proc, pattern_proc, cands_3b, derotate=True ) if cands_3b else [] # Pass 3b NCC gate: genuine 90°-rotated resistors match the rotated # template with NCC ≥ 0.73; structural look-alikes (capacitors, # inductors, transistor leads) match at NCC 0.35–0.45. Applying a # minimum NCC threshold here is far more discriminative than Chamfer # or DINOv2 alone — both of which are fooled by similar-shaped FPs. # Threshold 0.55 leaves a comfortable margin above FPs (max 0.45) # and below genuine TPs (min 0.73). _VERT_NCC_MIN = 0.55 before_3b_ncc = len(verified_3b) verified_3b = [c for c in verified_3b if c.get("ncc_score", 0) >= _VERT_NCC_MIN] if len(verified_3b) != before_3b_ncc: print(f"[Pipeline] Pass 3b NCC gate: {before_3b_ncc} -> {len(verified_3b)}") print(f"[Pipeline] Pass 3b (micro 90°): {len(cands_3b)} cands -> {len(verified_3b)} verified") all_complex = verified_filtered + verified_3a + verified_3b # Chamfer shape filter: structural edge-alignment quality check. # Applied only on the probe-focused path where DINOv2 alone is # insufficient -- components at unusual scales (>1.40x) attract FPs # from visually similar but structurally different symbols. # Standard-path templates (BR, IEC) have complex internal edge # structure; Chamfer at 3.0 wrongly rejects real detections there. if _use_probe_focused and all_complex: before_ch = len(all_complex) # Threshold 5.0: real zigzag TPs score 0.5–4.1; the single confirmed # FP at scale boundary scored 6.4 -- this cleanly removes it. # Standard path skips Chamfer: IEC has some candidates with # Chamfer 9–10 due to bbox distortion from the dilated-template # pass; filtering them drops real detections. # Recall-boost relaxes Chamfer (the VLM scrubs the extra FPs it lets # through); otherwise keep the tight 5.0 / 4.0 structural gates. _chamfer_max = 7.0 if self.vlm_recall_boost else 5.0 all_complex = self.postprocessor.filter_chamfer_shape( all_complex, drawing_proc, pattern_proc, max_chamfer=_chamfer_max ) if len(all_complex) != before_ch: print(f"[Pipeline] Chamfer filter: {before_ch} -> {len(all_complex)}") # Tighter Chamfer for horizontal candidates: horizontal TPs max at # ~3.3 (measured); vertical TPs can reach ~4.1 due to rotation/resize. # Candidates with angle ≈ 0° and Chamfer 4.0–5.0 are structural FPs # (inductors, transistors) that the global 5.0 threshold misses. # Skipped under recall-boost — the VLM handles those FPs instead. if not self.vlm_recall_boost: before_hch = len(all_complex) all_complex = [ c for c in all_complex if abs(c.get("angle", 0)) >= 45 or c.get("chamfer_dist", 0) <= 4.0 ] if len(all_complex) != before_hch: print(f"[Pipeline] Chamfer H filter: {before_hch} -> {len(all_complex)}") # Output-bubble filter: only for gate-like templates at unusual scales. # XOR/XNOR output bubbles match the gate body; filter them out. # Standard-scale templates (including bridge rectifiers) skip this. _is_gate_like = 0.25 <= _tmpl_ar <= 2.5 if _is_gate_like and _no_std_candidates: before_bubble = len(all_complex) all_complex = self.postprocessor.filter_output_bubble( all_complex, drawing_proc, pattern_proc ) if len(all_complex) != before_bubble: print(f"[Pipeline] Bubble filter: {before_bubble} -> {len(all_complex)}") verified = all_complex self.ncc_matcher.scales = _saved_scales self.ncc_matcher.ncc_threshold = _saved_ncc self.ncc_matcher.angles = _saved_angles self.dino_verifier.cosine_threshold = _saved_dino t3 = time.time() print(f"[Pipeline] All passes: {t3 - t2:.2f}s -- {len(verified)} total verified") # Simple-template post-filters: isolation + aspect-ratio + neighborhood. # Isolation: circuit components sit in white space; BOM/title-block cells have # solid grid lines directly adjacent to their long sides. # Aspect ratio: keep candidates whose bbox AR is within 2x of the template AR # *or* its reciprocal -- the reciprocal check allows 90°-rotated components # (e.g. vertical resistors) whose bbox AR is ~1/_tmpl_ar. # Neighborhood complexity: reject candidates whose surrounding ring has too # many Canny edges -- this eliminates false positives inside complex symbols # (e.g. bridge-rectifier bodies) which have many adjacent edges. # Notes-area exemption: the bottom 20% of the drawing contains notes/legend # symbols which have annotation text nearby -- skip isolation and neighborhood # checks for those so legitimate legend symbols are not filtered out. # Top-margin exclusion: discard detections whose top edge is within the outer # coordinate-margin strip (border grid cells look like plain rectangles). if _is_simple and verified: before = len(verified) _drw_h, _drw_w = drawing_proc.shape[:2] _notes_y = int(_drw_h * 0.80) # below this -> notes/legend area _top_margin = max(30, int(_drw_h * 0.04)) # top coordinate strip # Reject border-grid cells in the top margin verified = [c for c in verified if c["y"] >= _top_margin] # DINODense candidates are already DINOv2-verified (sim >= threshold). # They bypass the NCC-era structural filters (wire-leads, chamfer, etc.) # which assume tight bbox alignment. DINOv2 score IS the structural check. _dense_cands = [c for c in verified if c.get("from_dino_dense")] _ncc_cands = [c for c in verified if not c.get("from_dino_dense")] # Split NCC candidates into circuit area and notes/legend area _circuit = [c for c in _ncc_cands if c["y"] < _notes_y] _notes = [c for c in _ncc_cands if c["y"] >= _notes_y] # Dense candidates: circuit zone only (no structural filters) _dense_circuit = [c for c in _dense_cands if c["y"] < _notes_y] # Isolation: reject candidates adjacent to grid lines (circuit area only) _circuit = self.postprocessor.filter_isolated(_circuit, drawing_proc) # Aspect ratio: accept both normal (horizontal) and 90°-rotated (vertical) AR _ar_inv = 1.0 / max(0.01, _tmpl_ar) def _ar_ok(c): ar = c["w"] / max(1, c["h"]) return ( (_tmpl_ar / 2.0 <= ar <= _tmpl_ar * 2.0) or (_ar_inv / 2.0 <= ar <= _ar_inv * 2.0) ) _circuit = [c for c in _circuit if _ar_ok(c)] _notes = [c for c in _notes if _ar_ok(c)] # Wire-lead check: real resistors have straight wire leads on both # connecting sides; false positives embedded in complex symbols do not. # Circuit area: require one strong lead (>=6px) + one non-zero lead (>=1px) # -- handles resistors connected directly to a power rail where only # 1-2px of wire is visible on the rail-side. # Notes/legend area: skip wire-lead filter (legend symbols may have no # protruding leads). Restrict to left half of drawing (legends are always # on the left; BOM/title-block cells are on the right). Keep top-2 by score. _circuit = self.postprocessor.filter_wire_leads(_circuit, drawing_proc) _circuit = self.postprocessor.filter_wire_passthrough(_circuit, drawing_proc) _circuit = self.postprocessor.filter_neighborhood_complexity( _circuit, drawing_proc, expand_ratio=0.5, max_edge_density=0.022 ) _circuit = self.postprocessor.filter_junction_dots(_circuit, drawing_proc) _circuit = self.postprocessor.filter_rect_integrity(_circuit, drawing_proc) _circuit = self.postprocessor.filter_chamfer_shape(_circuit, drawing_proc, pattern_proc) # Orientation-aware minimum confidence. # # The template has aspect ratio _tmpl_ar (width / height). # Candidates in the same orientation as the template ("native") are # matched directly by NCC and should score high -- low-confidence # native candidates are almost certainly FPs (inductors, transistors, # op-amps that partially match the template bbox at low NCC). # # Candidates in the perpendicular orientation ("rotated") are matched # after a 90° rotation, which inherently lowers the NCC score even # for genuine resistors; they deserve a much more lenient threshold. # # Calibration from drawing 1 (test_1.png AR≈2.5, horizontal template): # Native (horizontal) TPs: conf 0.77, 0.78, 0.78 (all >= 0.70) # Rotated (vertical) TPs: conf 0.50–0.78 (all >= 0.45) # Observed FPs in complex drawings: conf 0.58–0.67 (all horizontal) # -> threshold 0.70 for native / 0.45 for rotated removes FPs while # keeping all drawing-1 TPs. _tmpl_native_wide = _tmpl_ar >= 1.0 # template is wider-than-tall before_oc = len(_circuit) def _is_native_orient(c): cand_wide = c["w"] >= c["h"] return cand_wide == _tmpl_native_wide # Pass-B (rotated-image) candidates are marked angle=90 and have already # been filtered at conf>=0.58 before structural filters; they represent # native-orientation matches so use a moderate threshold here. # Pass-A vertical candidates (not from rotated-image) still matched via # template rotation and score lower -- they need conf>=0.50 instead of 0.45. _circuit = [ c for c in _circuit if (_is_native_orient(c) and c.get("confidence", 0) >= 0.70) or (not _is_native_orient(c) and c.get("confidence", 0) >= 0.45) ] if len(_circuit) != before_oc: print( f"[Pipeline] Orient-conf filter: {before_oc} -> {len(_circuit)} " f"(native>=0.70 | rotated>=0.45)" ) _bottom_margin = max(30, int(_drw_h * 0.04)) _notes = [c for c in _notes if c["y"] + c["h"] <= _drw_h - _bottom_margin] _notes = self.postprocessor.filter_rect_borders(_notes, drawing_proc) _notes = sorted(_notes, key=lambda c: c.get("dino_score", 0), reverse=True)[:1] # DINOv2 Self-Supervised Prototype Filter # # Instead of comparing borderline candidates to the TEMPLATE (which # may be a different drawing style), build a prototype from the # HIGH-CONFIDENCE detections in THIS DRAWING. These are confirmed # instances of the target symbol in the actual drawing style and scale. # # Algorithm: # 1. Extract DINOv2 embeddings for all high-conf TPs (conf >= 0.72) # -- orientation-normalised so horizontal and vertical instances # of the same symbol produce comparable embeddings. # 2. Prototype = mean unit-normalised embedding. # 3. For each borderline candidate: compute cosine(candidate, prototype). # 4. Reject if similarity < min_sim. # # Why this works: # -- Resistors (any orientation) -> similar DINOv2 embedding -> high sim # -- Inductors/transistors/op-amps -> different embedding -> low sim # -- Prototype adapts to the specific drawing style automatically _proto_threshold = 0.72 # high-conf TPs used for prototype _proto_min_sim = 0.82 # borderline candidates below this are rejected _hc = [c for c in _circuit if c.get("confidence", 0) >= _proto_threshold] _bl = [c for c in _circuit if c.get("confidence", 0) < _proto_threshold] if len(_hc) >= 3 and _bl: dh, dw = drawing_proc.shape[:2] hc_crops = [ self.dino_verifier._crop_with_padding(drawing_proc, c, dh, dw) for c in _hc ] hc_embeds = self.dino_verifier.embed_crops_normalized(hc_crops) prototype = hc_embeds.mean(axis=0) _pnorm = float(np.linalg.norm(prototype)) if _pnorm > 1e-6: prototype = prototype / _pnorm bl_crops = [ self.dino_verifier._crop_with_padding(drawing_proc, c, dh, dw) for c in _bl ] bl_embeds = self.dino_verifier.embed_crops_normalized(bl_crops) bl_sims = bl_embeds @ prototype # (M,) cosine similarities _accepted = [c for c, s in zip(_bl, bl_sims.tolist()) if s >= _proto_min_sim] _rejected = [c for c, s in zip(_bl, bl_sims.tolist()) if s < _proto_min_sim] if _rejected: print( f"[Pipeline] DINO-proto: {len(_bl)} border -> " f"{len(_accepted)} kept, {len(_rejected)} rejected " f"(sim>={_proto_min_sim})" ) _circuit = _hc + _accepted else: _circuit = _hc + _bl # fallback else: _circuit = _hc + _bl # not enough high-conf to build prototype # Merge: NCC circuit + DINODense circuit (skipped structural filters) # DINODense candidates already have high DINOv2 similarity >= threshold verified = _circuit + _dense_circuit + _notes print( f"[Pipeline] Simple-template filters: {before} -> {len(verified)} " f"(NCC:{len(_circuit)} | DINODense:{len(_dense_circuit)} | notes:{len(_notes)})" ) # Title-block zone filter: remove candidates inside the BOM / right-frame # area for all templates (complex templates skip the simple-filter block # that previously applied this only to simple micro-pass candidates). before_tb = len(verified) verified = self.postprocessor.filter_title_block(verified, drawing_proc) if len(verified) != before_tb: print(f"[Pipeline] Title-block filter: {before_tb} -> {len(verified)}") # Adaptive confidence-gap filter: detect bimodal confidence distribution # and remove the low-confidence cluster (structural FPs: inductors, # transistors, op-amps that barely pass DINOv2 with low NCC). # Only applied to complex templates -- simple templates use dedicated # structural filters (wire-leads, chamfer, DINO-prototype) that are more # precise; the gap filter risks removing genuine low-confidence TPs there. # Skipped under recall-boost: the gap filter drops the low-confidence # cluster, which is exactly where borderline-but-genuine resistors live. # With the VLM on we keep them and let Stage 3 decide per-candidate. if not _is_simple and not self.vlm_recall_boost: before_gap = len(verified) verified = self.postprocessor.filter_confidence_gap(verified) if len(verified) != before_gap: print(f"[Pipeline] Confidence gap filter: {before_gap} -> {len(verified)}") # Stage 3 (optional): VLM semantic filter. # Runs AFTER all spatial/structural filters and BEFORE final NMS so the # VLM only judges a small, already-pruned candidate set (fast: ~0.4s/crop). # Zero-shot: the template's own VLM class defines the target, so no # symbol name is hardcoded. Removes inductor/crystal/op-amp FPs that # share low-level structure with the target and pass DINOv2. if self.use_vlm and verified: before_vlm = len(verified) try: vlm = self._get_vlm() verified = vlm.filter_by_template_class( drawing_proc, pattern_proc, verified, keep_min_conf=self.vlm_keep_min_conf, reject_only=self.vlm_reject_only, verbose=True ) print(f"[Pipeline] VLM Stage-3 filter: {before_vlm} -> {len(verified)}") except Exception as vlm_err: # The VLM is an optional precision booster. If it fails (most # commonly CUDA OOM on a 12 GB card already holding DINOv2 + # candidate tensors), degrade gracefully: keep the NCC+DINOv2 # results rather than failing the whole request. is_oom = "out of memory" in str(vlm_err).lower() print(f"[Pipeline] VLM Stage-3 SKIPPED ({'OOM' if is_oom else type(vlm_err).__name__}): " f"{vlm_err}") self._vlm = None # drop the half-loaded model try: import torch torch.cuda.empty_cache() except Exception: pass # Final NMS + format # Simple templates: tight IoU (0.25), no union expand (keep best-fit bbox). # Complex templates: # - probe-focused path: tight IoU (0.25) to suppress FP clusters in dense # circuit regions where multiple off-scale detections land on the same symbol # - standard path (_complex_use_union=True): union expand for chain suppression if _is_simple: verified = self.postprocessor.final_nms( verified, iou_threshold=0.25, use_union_bbox=False ) else: verified = self.postprocessor.final_nms( verified, iou_threshold=self.final_nms_iou, use_union_bbox=_complex_use_union ) result = self.postprocessor.format_output(verified, drawing_proc.shape) t4 = time.time() print(f"[Pipeline] Auto-detect total: {t4 - t0:.2f}s -- {result['total_detections']} detections") if return_visualization: # Draw on original (non-binarized) image for clearer output result["visualization"] = self.postprocessor.draw_boxes( drawing_data["original"], result["detections"] ) return result except Exception as e: raise RuntimeError(f"Pipeline auto-detect failed: {e}") from e def detect( self, pattern_input: Union[str, np.ndarray], drawing_input: Union[str, np.ndarray], return_visualization: bool = True, ) -> dict: """Run full detection pipeline. Args: pattern_input: Pattern image path or numpy array. drawing_input: Drawing image path or numpy array. return_visualization: Whether to include annotated image in output. Returns: Dict from Postprocessor.format_output(), plus optional "visualization" key. Raises: RuntimeError: If any stage fails. """ try: t0 = time.time() # Stage 0: Preprocess pattern_data = self.preprocessor.preprocess(pattern_input) drawing_data = self.preprocessor.preprocess(drawing_input) t1 = time.time() print(f"[Pipeline] Stage 0 (Preprocess): {t1 - t0:.2f}s") pattern_proc = pattern_data["processed"] # original -- used for DINOv2 drawing_proc = drawing_data["processed"] # Optionally dilate pattern strokes for NCC to handle style mismatch if self.dilate_pattern > 0: pattern_for_ncc = self.preprocessor.dilate_strokes( pattern_proc, kernel_size=self.dilate_pattern ) else: pattern_for_ncc = pattern_proc # Stage 1: NCC matching (uses dilated pattern if configured) candidates = self.ncc_matcher.match(drawing_proc, pattern_for_ncc) t2 = time.time() print(f"[Pipeline] Stage 1 (NCC): {t2 - t1:.2f}s -- {len(candidates)} candidates") if not candidates: print("[Pipeline] No candidates from Stage 1, returning empty result.") result = self.postprocessor.format_output([], drawing_proc.shape) if return_visualization: result["visualization"] = self.postprocessor.draw_boxes( drawing_data["original"], [] ) return result # Stage 2: DINOv2 verification candidates = self.dino_verifier.verify_candidates(drawing_proc, pattern_proc, candidates) t3 = time.time() print(f"[Pipeline] Stage 2 (DINOv2): {t3 - t2:.2f}s -- {len(candidates)} verified") # Stage 3: Final NMS + format candidates = self.postprocessor.final_nms(candidates, iou_threshold=self.final_nms_iou) result = self.postprocessor.format_output(candidates, drawing_proc.shape) t4 = time.time() print(f"[Pipeline] Stage 3 (Post): {t4 - t3:.2f}s") print(f"[Pipeline] Total: {t4 - t0:.2f}s -- {result['total_detections']} detections") if return_visualization: # Draw on original (non-binarized) image for clearer output result["visualization"] = self.postprocessor.draw_boxes( drawing_data["original"], result["detections"] ) return result except Exception as e: raise RuntimeError(f"Pipeline detection failed: {e}") from e def update_thresholds( self, ncc_threshold: Optional[float] = None, cosine_threshold: Optional[float] = None, final_nms_iou: Optional[float] = None, ): """Update detection thresholds at runtime (for UI sliders). Args: ncc_threshold: New NCC threshold for Stage 1. cosine_threshold: New cosine similarity threshold for Stage 2. """ if ncc_threshold is not None: self.ncc_matcher.ncc_threshold = ncc_threshold if cosine_threshold is not None: self.dino_verifier.cosine_threshold = cosine_threshold if final_nms_iou is not None: self.final_nms_iou = final_nms_iou def run_detection(pattern_path: str, drawing_path: str, auto: bool = True, **kwargs) -> dict: """Convenience function: create pipeline, run detection, return result. Args: pattern_path: Path to pattern image. drawing_path: Path to drawing image. auto: If True (default), use detect_auto() which self-tunes thresholds. **kwargs: Config overrides passed to PatternDetectionPipeline. Returns: Detection result dict. """ pipeline = PatternDetectionPipeline(config=kwargs if kwargs else None) if auto: return pipeline.detect_auto(pattern_path, drawing_path) return pipeline.detect(pattern_path, drawing_path)