"""Stage 3 (optional): Vision-Language Model verifier using Qwen2-VL-2B. Why a VLM after NCC + DINOv2: NCC matches by pixel correlation and DINOv2 by patch-feature cosine. Both encode *how a symbol looks* (line density, spatial frequency). Two schematic symbols with similar low-level structure -- a zigzag/rectangle resistor vs a coiled inductor vs a crystal -- land close in both spaces, so the worst false positives (inductor, crystal, op-amp fragments) survive every spatial filter. A VLM reasons about *what the symbol is*. Asked "is this a resistor?", Qwen2-VL can use the learned semantic concept to reject an inductor even when its crop correlates highly with the resistor template. This is the discriminability CLIP image-to-image lacked (see design_spec/model_survey.md). Design: * Drop-in Stage-3 filter: same signature shape as DINOVerifier.verify_candidates. * Lazy load: the 2B model (~4.5 GB bf16) is only loaded on first use, so the rest of the pipeline runs without it when `use_vlm` is off. * Each candidate crop is taken with generous context padding and upscaled so the small (30-80 px) line-art symbol is legible to the VLM, then the model is asked a constrained yes/no question. Only "yes" candidates survive. * The template image is shown to the VLM as a visual reference so the check stays zero-shot (no hardcoded symbol class name needed); an optional `symbol_name` sharpens the prompt when the class is known. Tested on: RTX 3060 12 GB, transformers >= 4.45 (Qwen2-VL support). """ from __future__ import annotations import re from typing import List, Optional import cv2 import numpy as np class VLMVerifier: """Verify candidate crops with a local Qwen2-VL model (yes/no semantic check). Args: model_name: HuggingFace model ID. Default Qwen2-VL-2B-Instruct. device: "cuda" | "cpu" | None (auto-detect). symbol_name: Optional human name of the target symbol ("resistor"). When provided it is woven into the prompt for a sharper decision; when None the model compares each crop against the shown template image only. min_keep_conf: Candidates with confidence >= this are kept WITHOUT asking the VLM (trusted high-confidence detections — saves inference time). max_ask_conf: Candidates with confidence above this are auto-kept; only those in (min_keep_conf upper region downward) borderline band are asked. (See verify_candidates for the exact banding logic.) context_pad_ratio: Fraction of bbox size added on each side as context. upscale_to: Target longer-side pixel size for each crop shown to the VLM. """ def __init__( self, model_name: str = "Qwen/Qwen2-VL-2B-Instruct", device: Optional[str] = None, symbol_name: Optional[str] = None, context_pad_ratio: float = 0.6, upscale_to: int = 224, max_new_tokens: int = 8, ): self.model_name = model_name self.symbol_name = symbol_name self.context_pad_ratio = context_pad_ratio self.upscale_to = upscale_to self.max_new_tokens = max_new_tokens self._device_arg = device self._model = None # lazy self._processor = None # lazy self.device = device or "cuda" # ------------------------------------------------------------------ # Lazy model loading # ------------------------------------------------------------------ def _ensure_loaded(self): if self._model is not None: return import torch from transformers import AutoProcessor, Qwen2VLForConditionalGeneration device = self._device_arg if device is None: device = "cuda" if torch.cuda.is_available() else "cpu" self.device = device # Free any cached allocations from earlier stages (DINOv2, NCC tensors) # before pulling the ~4.5 GB VLM onto the GPU. On a 12 GB card shared with # the OS display this reclaimed headroom is the difference between fitting # and an OOM. if device == "cuda": torch.cuda.empty_cache() dtype = torch.bfloat16 if device == "cuda" else torch.float32 print(f"[VLMVerifier] Loading {self.model_name} on {device} ({dtype})...") self._model = Qwen2VLForConditionalGeneration.from_pretrained( self.model_name, torch_dtype=dtype, device_map=device, low_cpu_mem_usage=True, ) self._model.eval() # min_pixels/max_pixels keep token counts bounded for small crops self._processor = AutoProcessor.from_pretrained( self.model_name, min_pixels=128 * 28 * 28, max_pixels=512 * 28 * 28, ) self._torch = torch print("[VLMVerifier] Model ready.") # ------------------------------------------------------------------ # Crop preparation # ------------------------------------------------------------------ def _crop_for_vlm(self, drawing: np.ndarray, c: dict) -> np.ndarray: """Extract a candidate region with context padding, de-rotate, upscale. Small schematic symbols are illegible to a VLM at native resolution; we add surrounding context (helps the model see connecting wires) and upscale the longer side to `upscale_to` px with cubic interpolation. """ H, W = drawing.shape[:2] x, y, w, h = c["x"], c["y"], c["w"], c["h"] pad_x = int(w * self.context_pad_ratio) pad_y = int(h * self.context_pad_ratio) x1 = max(0, x - pad_x) y1 = max(0, y - pad_y) x2 = min(W, x + w + pad_x) y2 = min(H, y + h + pad_y) crop = drawing[y1:y2, x1:x2] if crop.size == 0: crop = np.full((self.upscale_to, self.upscale_to), 255, np.uint8) # De-rotate vertical candidates to a canonical horizontal orientation angle = c.get("angle", 0) if 70 <= abs(angle) <= 110 and crop.shape[0] > crop.shape[1]: crop = cv2.rotate(crop, cv2.ROTATE_90_COUNTERCLOCKWISE) # Upscale longer side to upscale_to ch, cw = crop.shape[:2] longer = max(ch, cw) if longer > 0 and longer < self.upscale_to: f = self.upscale_to / longer crop = cv2.resize(crop, (max(1, int(cw * f)), max(1, int(ch * f))), interpolation=cv2.INTER_CUBIC) # To 3-channel for the VLM if crop.ndim == 2: crop = cv2.cvtColor(crop, cv2.COLOR_GRAY2RGB) return crop def _prep_template(self, template: np.ndarray) -> np.ndarray: t = template if t.ndim == 2: t = cv2.cvtColor(t, cv2.COLOR_GRAY2RGB) ch, cw = t.shape[:2] longer = max(ch, cw) if longer > 0 and longer < self.upscale_to: f = self.upscale_to / longer t = cv2.resize(t, (max(1, int(cw * f)), max(1, int(ch * f))), interpolation=cv2.INTER_CUBIC) return t # ------------------------------------------------------------------ # Prompting # ------------------------------------------------------------------ # Component vocabulary for classification mode _CLASSES = [ "resistor", "inductor", "capacitor", "diode", "crystal", "transistor", "op-amp", "logic-gate", "wire-junction", "other", ] # Labels the 2B model assigns unreliably: it defaults ambiguous line-art to # "transistor", and "other"/"wire-junction" are catch-alls. A candidate with # one of these is NOT confidently a non-target, so in reject-only mode we keep # it rather than risk discarding a genuine target the model just mislabelled. _UNTRUSTED_LABELS = {"transistor", "wire-junction", "other"} def _build_prompt(self) -> str: sym = self.symbol_name or "the reference component symbol shown in the first image" return ( "You are an expert at reading electronic schematic diagrams.\n" "The FIRST image is a reference symbol. The SECOND image is a candidate " "region cropped from a larger schematic.\n" f"Question: Does the SECOND image contain the SAME type of component as " f"{sym}?\n" "Pay attention to the exact symbol shape. A resistor is a zigzag or a " "plain rectangle. Do NOT confuse it with: an inductor (series of loops/" "humps), a crystal (rectangle between two capacitor plates), a capacitor " "(two parallel lines), a diode (triangle+bar), or an op-amp (large " "triangle). Answer with exactly one word: 'yes' or 'no'." ) def _build_classify_prompt(self) -> str: """Open-classification prompt — avoids the yes/no agreement bias of small VLMs. Instead of confirming "is this X?" (which a 2B model tends to answer 'yes' to regardless), we force it to NAME the component from a closed vocabulary. The caller keeps only candidates classified as the target class. """ classes = ", ".join(self._CLASSES) return ( "You are an expert at reading electronic schematic diagrams. The image " "is a small region cropped from a schematic, possibly with connecting " "wires around the central component.\n" "Identify the SINGLE central electronic component. Shape guide:\n" "- resistor: a zigzag (sawtooth) line, OR a plain rectangle in series " "with a wire.\n" "- inductor: a series of rounded loops / humps / coils.\n" "- capacitor: two short parallel lines (or one curved) with a gap.\n" "- diode: a triangle pointing into a bar.\n" "- crystal: a small rectangle drawn between two capacitor plates.\n" "- transistor: a circle/junction with three leads.\n" "- op-amp: a large triangle.\n" "- logic-gate: AND/OR/XOR gate outline.\n" "- wire-junction: only wires/dots, no component body.\n" f"Answer with EXACTLY ONE word from this list: {classes}." ) @staticmethod def _np_to_pil(arr: np.ndarray): from PIL import Image return Image.fromarray(arr.astype(np.uint8)) def _generate(self, content: list) -> str: """Run one chat turn with the given content list, return decoded text.""" from qwen_vl_utils import process_vision_info messages = [{"role": "user", "content": content}] text = self._processor.apply_chat_template( messages, tokenize=False, add_generation_prompt=True ) image_inputs, video_inputs = process_vision_info(messages) inputs = self._processor( text=[text], images=image_inputs, videos=video_inputs, padding=True, return_tensors="pt", ).to(self.device) with self._torch.no_grad(): gen = self._model.generate(**inputs, max_new_tokens=self.max_new_tokens, do_sample=False) trimmed = gen[:, inputs.input_ids.shape[1]:] out = self._processor.batch_decode( trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False )[0] # Release the per-crop activation/KV-cache spike so memory does not creep # up across a long candidate list on a tight (12 GB) card. del inputs, gen, trimmed if self.device == "cuda": self._torch.cuda.empty_cache() return out.strip().lower() def _ask_one(self, template_rgb: np.ndarray, crop_rgb: np.ndarray): """Run a single yes/no query. Returns (decided_bool, raw_answer).""" ans = self._generate([ {"type": "image", "image": self._np_to_pil(template_rgb)}, {"type": "image", "image": self._np_to_pil(crop_rgb)}, {"type": "text", "text": self._build_prompt()}, ]) m = re.search(r"\b(yes|no)\b", ans) decided = (m.group(1) == "yes") if m else ans.startswith("y") return decided, ans def _classify_one(self, crop_rgb: np.ndarray): """Open-classification query. Returns (class_label, raw_answer).""" ans = self._generate([ {"type": "image", "image": self._np_to_pil(crop_rgb)}, {"type": "text", "text": self._build_classify_prompt()}, ]) # Match the first known class token that appears in the answer label = "other" for cls in self._CLASSES: if re.search(rf"\b{re.escape(cls)}\b", ans): label = cls break return label, ans # ------------------------------------------------------------------ # Public interface # ------------------------------------------------------------------ def classify_template(self, template: np.ndarray) -> str: """Classify the template itself to learn the target component class. Keeps the VLM stage zero-shot: we never hardcode "resistor"; instead the template's own classification defines what candidates must match. """ self._ensure_loaded() t = self._prep_template(template) label, raw = self._classify_one(t) print(f"[VLMVerifier] Template classified as: {label!r} (raw {raw!r})") return label def filter_by_template_class( self, drawing: np.ndarray, template: np.ndarray, candidates: List[dict], target_class: Optional[str] = None, keep_min_conf: float = 1.01, reject_only: bool = True, verbose: bool = False, ) -> List[dict]: """Open-classification Stage-3 filter (robust to small-VLM yes-bias). Two modes: * reject_only=True (default): blacklist. Drop a candidate ONLY when the VLM confidently names it a DIFFERENT, visually-distinct component (inductor/capacitor/diode/op-amp/crystal/logic-gate). Candidates the VLM calls the target — or gives an unreliable label (transistor/other/ wire-junction) — are KEPT. Maximises recall: the 2B model mislabels many genuine resistors as "transistor", so a strict whitelist wrongly discards them. * reject_only=False: whitelist. Keep only candidates the VLM labels as the target class. Higher precision, lower recall. Args: target_class: Class candidates must match / not contradict. If None, derived from the template via classify_template(). keep_min_conf: Candidates with confidence >= this are kept WITHOUT asking the VLM (trusted high-conf detections; saves inference). reject_only: see above. """ if not candidates: return [] self._ensure_loaded() if target_class is None: target_class = self.classify_template(template) # In reject-only mode a label triggers a drop only if it is a trusted, # concretely-different component class (not the unreliable catch-alls). reject_labels = { cls for cls in self._CLASSES if cls != target_class and cls not in self._UNTRUSTED_LABELS } kept = [] for c in candidates: if c.get("confidence", 0.0) >= keep_min_conf: c["vlm_class"] = "auto-keep" kept.append(c) continue crop = self._crop_for_vlm(drawing, c) label, raw = self._classify_one(crop) c["vlm_class"] = label if reject_only: keep = label not in reject_labels else: keep = label == target_class if verbose: print(f" [VLM] ({c['x']},{c['y']}) conf={c.get('confidence',0):.2f} " f"-> {label!r} {'KEEP' if keep else 'DROP'}") if keep: kept.append(c) mode = "reject-only" if reject_only else "whitelist" print(f"[VLMVerifier] class-filter ({target_class}, {mode}): " f"{len(candidates)} -> {len(kept)} kept") return kept def verify_candidates( self, drawing: np.ndarray, template: np.ndarray, candidates: List[dict], ask_band: tuple = (0.0, 1.01), verbose: bool = False, ) -> List[dict]: """Filter candidates with the VLM. Args: drawing: Preprocessed (binarised) drawing, grayscale. template: Preprocessed template image, grayscale. candidates: Detection dicts (need x,y,w,h, optional angle/confidence). ask_band: (low, high) confidence band. Candidates whose confidence is INSIDE this band are sent to the VLM; candidates above `high` are auto-kept (trusted), candidates below `low` are auto-rejected. Default (0.0, 1.01) asks the VLM about every candidate. verbose: Print per-candidate VLM answers. Returns: Filtered list (VLM 'no' candidates removed). Each surviving candidate gains a 'vlm_pass' bool and 'vlm_answer' string when it was asked. """ if not candidates: return [] self._ensure_loaded() template_rgb = self._prep_template(template) low, high = ask_band kept = [] for c in candidates: conf = c.get("confidence", 1.0) if conf >= high: c["vlm_pass"] = True c["vlm_answer"] = "auto-keep" kept.append(c) continue if conf < low: c["vlm_pass"] = False c["vlm_answer"] = "auto-reject" continue crop_rgb = self._crop_for_vlm(drawing, c) decided, raw = self._ask_one(template_rgb, crop_rgb) c["vlm_pass"] = decided c["vlm_answer"] = raw if verbose: print(f" [VLM] ({c['x']},{c['y']}) conf={conf:.2f} -> {raw!r} " f"=> {'KEEP' if decided else 'DROP'}") if decided: kept.append(c) print(f"[VLMVerifier] {len(candidates)} candidates -> {len(kept)} kept") return kept