| |
| """ |
| Classical CV detector for target unit symbols in construction blueprints. |
| |
| Default target shape: a diamond (rotated square) outline containing exactly |
| one capital letter followed by one digit (e.g. A1, B2, C3) β see shapes/diamond.py |
| for that pipeline. |
| |
| If a user supplies an example crop of a different target symbol, detection |
| switches to exemplar-driven matching: a mask/template is built from the |
| example, the page is searched for it via multi-scale normalised cross- |
| correlation, and each match is run through a shape-specific secondary |
| validator (see shapes/) chosen by the exemplar's classified shape type |
| before OCR. |
| |
| Usage: |
| python3 classical_cv_detector.py |
| python3 classical_cv_detector.py --pdf "Data/Testing/Input/Window Test Images.pdf" |
| python3 classical_cv_detector.py --dpi 300 --tile-size 768 |
| python3 classical_cv_detector.py --exemplar path/to/symbol_crop.png |
| python3 classical_cv_detector.py --diagnose # prints per-filter drop counts |
| python3 classical_cv_detector.py --debug # saves binary tile image |
| """ |
|
|
| import argparse |
| import os |
| from collections import Counter |
|
|
| import cv2 |
| import numpy as np |
|
|
| import shapes |
| from cv_core import ( |
| DPI, TILE_SIZE, TILE_OVERLAP, MIN_AREA, OCR_UPSCALE, MASK_ERODE, NMS_DIST, |
| binarize, remove_hv_lines, deduplicate, iter_tiles, load_pages, _ocr_with_mask, |
| find_marked_boxes, match_marked_boxes, report_progress, SpatialIndex, |
| ) |
| from shapes import _geometry |
| from shapes.diamond import TMPL_PAD |
|
|
|
|
| |
|
|
| def classify_shape(cnt: np.ndarray) -> str: |
| """ |
| Classify a contour's rough shape type from vertex count and circularity. |
| |
| Used to label the shape extracted from a user-supplied example image, |
| and to choose which shapes/ validator (if any) runs as a secondary check |
| on exemplar-driven matches. |
| """ |
| hull = cv2.convexHull(cnt) |
| area = cv2.contourArea(hull) |
| peri = cv2.arcLength(hull, True) |
| if area <= 0 or peri <= 0: |
| return "unknown" |
|
|
| circularity = 4 * np.pi * area / (peri ** 2) |
| x, y, w, h = cv2.boundingRect(hull) |
| aspect = min(w, h) / max(w, h) if max(w, h) else 0.0 |
| approx = cv2.approxPolyDP(hull, 0.02 * peri, True) |
| n = len(approx) |
|
|
| |
| |
| |
| if n == 3: |
| return "triangle" |
| if n == 4: |
| return "diamond" if _geometry.is_diamond_vertex_layout(approx, x, y, w, h) else "square" |
| if n == 5: |
| return "pentagon" |
| if n == 6: |
| return "hexagon" |
| if circularity > 0.85 and aspect > 0.85: |
| return "circle" |
| return "polygon" |
|
|
|
|
| def extract_shape_template( |
| exemplar_bgr: np.ndarray, |
| pad: int = TMPL_PAD, |
| ) -> dict | None: |
| """ |
| Build a binary outline template + interior mask from a tightly-cropped |
| example image of the target symbol, and classify its rough shape type. |
| |
| Mirrors the diamond pipeline's binarisation so the resulting template is |
| directly comparable via cv2.matchTemplate against page tiles binarised |
| the same way. Returns None if no usable contour is found (e.g. a blank |
| or near-empty crop). |
| """ |
| binary = binarize(exemplar_bgr) |
| contours, _ = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) |
| if not contours: |
| return None |
| cnt = max(contours, key=cv2.contourArea) |
| if cv2.contourArea(cnt) < MIN_AREA: |
| return None |
|
|
| hull = cv2.convexHull(cnt) |
| x, y, w, h = cv2.boundingRect(hull) |
| shape_type = classify_shape(cnt) |
|
|
| tw, th = w + 2 * pad, h + 2 * pad |
| origin = np.array([x - pad, y - pad]) |
|
|
| mask = np.zeros((th, tw), np.uint8) |
| cv2.drawContours(mask, [hull - origin], -1, 255, thickness=-1) |
| kernel = np.ones((5, 5), np.uint8) |
| mask = cv2.erode(mask, kernel, iterations=MASK_ERODE) |
|
|
| template = np.zeros((th, tw), np.uint8) |
| cv2.drawContours(template, [cnt - origin], -1, 255, thickness=2) |
|
|
| |
| |
| |
| |
| radius = max(w, h) / 2.0 |
|
|
| return { |
| "template": template, |
| "mask": mask, |
| "size": (tw, th), |
| "shape_type": shape_type, |
| "radius": radius, |
| } |
|
|
|
|
| CONFIRMED_TMPL_MIN_N = 3 |
|
|
|
|
| def build_confirmed_template( |
| page_bgr: np.ndarray, |
| confirmed_dets: list[dict], |
| pad: int = TMPL_PAD, |
| ) -> tuple[np.ndarray | None, tuple[int, int]]: |
| """ |
| Build a binary outline template by averaging binarised crops of on-page |
| detections that already passed OCR (code not in ("", "?")). |
| |
| Same averaging technique as shapes.diamond.build_exemplar_template |
| (outline pixels are spatially consistent across crops and reinforce |
| toward 1.0; interior text/codes vary by position and average out), but |
| shape-agnostic so it applies to any exemplar shape_type, not just |
| diamond. The result is calibrated to this page's own real pixel size, so |
| exemplar_template_match_pass can match against it at a single scale |
| instead of sweeping 0.5x-1.5x β unlike exemplar_data['template'], built |
| from a user-supplied crop whose scale relative to the page render is |
| unknown. |
| |
| Only OCR-confirmed hits are used deliberately: a shape-only ("?") hit's |
| bbox size is unverified (its geometry passed but nothing confirms it's |
| really the target shape), so anchoring the averaged size/outline on one |
| would risk baking that uncertainty into every match this template makes. |
| |
| Returns (None, (0, 0)) if fewer than CONFIRMED_TMPL_MIN_N detections are |
| given. |
| """ |
| if len(confirmed_dets) < CONFIRMED_TMPL_MIN_N: |
| return None, (0, 0) |
|
|
| H, W = page_bgr.shape[:2] |
| med_w = int(np.median([d["bbox"][2] for d in confirmed_dets])) |
| med_h = int(np.median([d["bbox"][3] for d in confirmed_dets])) |
| if med_w < 10 or med_h < 10: |
| return None, (0, 0) |
|
|
| tw, th = med_w + 2 * pad, med_h + 2 * pad |
|
|
| crops: list[np.ndarray] = [] |
| for d in confirmed_dets: |
| x, y, w, h = d["bbox"] |
| cx, cy = x + w // 2, y + h // 2 |
| x1 = max(0, cx - tw // 2) |
| y1 = max(0, cy - th // 2) |
| x2 = min(W, x1 + tw) |
| y2 = min(H, y1 + th) |
| if (x2 - x1) < tw // 2 or (y2 - y1) < th // 2: |
| continue |
| region = page_bgr[y1:y2, x1:x2] |
| binary = binarize(region) |
| resized = cv2.resize(binary, (tw, th), interpolation=cv2.INTER_NEAREST) |
| crops.append(resized.astype(np.float32) / 255.0) |
|
|
| if len(crops) < CONFIRMED_TMPL_MIN_N: |
| return None, (0, 0) |
|
|
| avg = np.mean(crops, axis=0) |
| _, template = cv2.threshold((avg * 255).astype(np.uint8), 127, 255, cv2.THRESH_BINARY) |
| return template, (tw, th) |
|
|
|
|
| def exemplar_template_match_pass( |
| page_bgr: np.ndarray, |
| exemplar_data: dict, |
| ocr_upscale: int = OCR_UPSCALE, |
| tile_size: int = TILE_SIZE, |
| overlap: float = TILE_OVERLAP, |
| scales: tuple = (0.5, 0.6, 0.7, 0.8, 0.9, 1.0, 1.1, 1.2, 1.3, 1.4, 1.5), |
| threshold: float = 0.30, |
| existing_dets: list[dict] | None = None, |
| shape_only_min_score: float = 0.55, |
| progress_dir: str | None = None, |
| progress_key=None, |
| ) -> list[dict]: |
| """ |
| Slide a user-supplied exemplar's outline template over every page tile |
| using normalised cross-correlation, sweeping a wide range of scales |
| (0.5x-1.5x) since the exemplar crop's DPI/zoom/style commonly won't match |
| the page render closely β a clean vector-icon exemplar matched against |
| noisier real CAD line work peaks at a modest correlation score even at |
| the right scale, so threshold is intentionally permissive: the |
| shape-specific geometric validator below (and OCR) carry most of the |
| false-positive rejection burden, not raw correlation strength. |
| |
| existing_dets (optional): detections already found by a prior, cheaper |
| pass (e.g. the shape-specific contour pipeline in detect_page() below). |
| When given, peaks within NMS_DIST of an existing detection are skipped, |
| so this pass only spends correlation/OCR budget on gaps the prior pass |
| missed (mirrors shapes.diamond.exemplar_match_pass's own first_pass_dets |
| NMS, applied here to the generic exemplar path). The scale sweep is also |
| adjusted, in order of preference: |
| 1. If >= CONFIRMED_TMPL_MIN_N of existing_dets are OCR-confirmed (code |
| not in ("", "?")), build_confirmed_template() averages their crops |
| into a page-native mask already at the page's real pixel size β |
| matching then runs at a single scale (no sweep at all), since the |
| template's scale relative to the page is known exactly rather than |
| guessed. |
| 2. Otherwise, if existing_dets is non-empty (e.g. too few are |
| OCR-confirmed yet, or the shape has no build_confirmed_template |
| population to draw from), the sweep narrows to a band around the |
| empirically observed size ratio (median existing-detection bbox |
| size / exemplar bbox size) instead of blindly sweeping the full |
| 0.5x-1.5x range. |
| 3. Otherwise (no existing_dets at all β e.g. a shape_type with no |
| registered contour pipeline) the full 0.5x-1.5x sweep runs, since |
| nothing on this page yet tells us the true scale. |
| |
| Each surviving peak (after self-NMS) is run through the shape-specific |
| secondary validator for exemplar_data['shape_type'] (see shapes/), if |
| one is registered β this is the geometric "is this really a <shape>" |
| check, mirroring the gates the diamond pipeline already has. Candidates |
| that fail are rejected before paying the OCR cost. Shape types with no |
| registered validator (square/triangle/pentagon/polygon/unknown) skip |
| this gate and rely on the correlation score + OCR alone. |
| |
| OCR uses the exemplar's own interior mask, not a synthetic diamond |
| contour, so it works for any shape. |
| """ |
| base_template = exemplar_data["template"] |
| base_mask = exemplar_data["mask"] |
| shape_type = exemplar_data["shape_type"] |
| base_radius = exemplar_data["radius"] |
| th_t, tw_t = base_template.shape[:2] |
| H, W = page_bgr.shape[:2] |
| validator = shapes.get_validator(shape_type) |
|
|
| confirmed_dets = [d for d in (existing_dets or []) if d.get("code") not in (None, "", "?")] |
| confirmed_template, confirmed_size = build_confirmed_template(page_bgr, confirmed_dets) |
|
|
| if confirmed_template is not None: |
| base_template = confirmed_template |
| base_mask = cv2.resize(base_mask, confirmed_size, interpolation=cv2.INTER_NEAREST) |
| th_t, tw_t = base_template.shape[:2] |
| base_radius = max(confirmed_size) / 2.0 |
| scales = (1.0,) |
| elif existing_dets: |
| observed_size = np.median([max(d["bbox"][2], d["bbox"][3]) for d in existing_dets]) |
| exemplar_size = max(tw_t, th_t) |
| if exemplar_size > 0 and observed_size > 0: |
| observed_ratio = observed_size / exemplar_size |
| scales = tuple(sorted({ |
| round(observed_ratio * f, 3) for f in (0.85, 1.0, 1.15) |
| })) |
|
|
| existing_centres = [ |
| (d["bbox"][0] + d["bbox"][2] // 2, d["bbox"][1] + d["bbox"][3] // 2) |
| for d in (existing_dets or []) |
| ] |
|
|
| raw_peaks: list[tuple[float, int, int, float]] = [] |
|
|
| tiles = list(iter_tiles(page_bgr, tile_size, overlap)) |
| n_tiles = len(tiles) |
| report_progress( |
| progress_dir, progress_key, |
| f"mask/template pass: starting ({n_tiles} tile(s) x {len(scales)} scale(s))", |
| ) |
| for _ti, (tile, x0, y0) in enumerate(tiles): |
| binary = binarize(tile) |
| cleaned = remove_hv_lines(binary) |
|
|
| for scale in scales: |
| sw, sh = max(1, int(tw_t * scale)), max(1, int(th_t * scale)) |
| tmpl_f32 = cv2.resize( |
| base_template, (sw, sh), interpolation=cv2.INTER_NEAREST |
| ).astype(np.float32) |
|
|
| for search_img in (binary, cleaned): |
| if search_img.shape[0] < sh or search_img.shape[1] < sw: |
| continue |
| result = cv2.matchTemplate( |
| search_img.astype(np.float32), tmpl_f32, cv2.TM_CCOEFF_NORMED |
| ) |
| ys, xs = np.where(result >= threshold) |
| for r, c in zip(ys, xs): |
| score = float(result[r, c]) |
| page_cx = c + x0 + sw // 2 |
| page_cy = r + y0 + sh // 2 |
| raw_peaks.append((score, page_cx, page_cy, scale)) |
| if (_ti + 1) % 5 == 0 or (_ti + 1) == n_tiles: |
| report_progress( |
| progress_dir, progress_key, |
| f"mask/template pass: tile {_ti + 1}/{n_tiles} " |
| f"({len(raw_peaks)} raw peak(s) so far)", |
| ) |
|
|
| |
| raw_peaks.sort(reverse=True) |
| report_progress( |
| progress_dir, progress_key, |
| f"mask/template pass: correlation done, {len(raw_peaks)} raw peak(s) -- starting NMS", |
| ) |
|
|
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| existing_index = SpatialIndex(NMS_DIST) |
| for ex, ey in existing_centres: |
| existing_index.add(ex, ey) |
| evaluated_index = SpatialIndex(NMS_DIST) |
|
|
| new_dets: list[dict] = [] |
| _n_raw_peaks = len(raw_peaks) |
|
|
| for _peak_i, (score, cx, cy, scale) in enumerate(raw_peaks): |
| if _peak_i % 200 == 0 or _peak_i + 1 == _n_raw_peaks: |
| report_progress( |
| progress_dir, progress_key, |
| f"mask/template pass: NMS+OCR peak {_peak_i + 1}/{_n_raw_peaks} " |
| f"({len(new_dets)} accepted so far)", |
| ) |
| |
| if existing_index.near_any(cx, cy): |
| continue |
| if evaluated_index.near_any(cx, cy): |
| continue |
| evaluated_index.add(cx, cy) |
|
|
| sw, sh = max(1, int(tw_t * scale)), max(1, int(th_t * scale)) |
| scaled_mask = cv2.resize(base_mask, (sw, sh), interpolation=cv2.INTER_NEAREST) |
|
|
| x1 = max(0, cx - sw // 2) |
| y1 = max(0, cy - sh // 2) |
| x2 = min(W, x1 + sw) |
| y2 = min(H, y1 + sh) |
| local_tile = page_bgr[y1:y2, x1:x2] |
|
|
| code = "" |
| if local_tile.size > 0 and local_tile.shape[:2] == scaled_mask.shape[:2]: |
| lx, ly = cx - x1, cy - y1 |
| r = max(1, int(base_radius * scale)) |
|
|
| if validator is not None: |
| local_binary = binarize(local_tile) |
| local_cleaned = remove_hv_lines(local_binary) |
| if not validator(local_binary, local_cleaned, lx, ly, r): |
| continue |
|
|
| code = _ocr_with_mask( |
| local_tile, scaled_mask, (0, 0, sw, sh), upscale=ocr_upscale, |
| ) |
|
|
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| if not code and score < shape_only_min_score: |
| continue |
|
|
| new_dets.append({ |
| "bbox": (cx - sw // 2, cy - sh // 2, sw, sh), |
| "code": code if code else "?", |
| "score": score, |
| }) |
|
|
| report_progress(progress_dir, progress_key, |
| f"mask/template pass done: {len(new_dets)} accepted det(s)") |
| return new_dets |
|
|
|
|
| |
|
|
| def detect_page( |
| bgr: np.ndarray, |
| tile_size: int = TILE_SIZE, |
| overlap: float = TILE_OVERLAP, |
| ocr_upscale: int = OCR_UPSCALE, |
| exemplar_bgr: np.ndarray | None = None, |
| progress_dir: str | None = None, |
| progress_key=None, |
| ) -> list[dict]: |
| """ |
| Full page detection pipeline. |
| |
| If `exemplar_bgr` is supplied (a tightly-cropped example image of the |
| target symbol), detection switches to exemplar-driven matching, run in |
| shape-first order: |
| |
| 1. Shape pass: extract_shape_template() classifies the exemplar's |
| shape_type. If a tuned contour/text-first pipeline is registered |
| for that shape_type (see shapes.get_pipeline() β currently diamond |
| and hexagon), it runs first. This is cheap relative to correlation |
| (no sliding multi-scale search) and carries all of that shape |
| module's recall tuning (tile-density gates, text-first character |
| detection, etc.) β none of which the generic correlation path below |
| can reproduce on its own. |
| 2. Mask/template pass: exemplar_template_match_pass() then searches |
| for the exemplar's outline via multi-scale normalised cross- |
| correlation, but only to fill gaps the shape pass missed β peaks |
| near an existing shape-pass detection are skipped, and the scale |
| sweep is narrowed around the size actually observed in the shape |
| pass's own hits (falls back to the full 0.5x-1.5x sweep when the |
| shape pass found nothing, or shape_type has no registered |
| pipeline). |
| |
| This mirrors the ordering shapes.diamond.detect_page() already uses |
| internally (contour passes first, exemplar template built from and run |
| after those hits) β applied here to the user-supplied exemplar path too, |
| instead of running correlation-first across the whole page. |
| |
| If no usable shape can be extracted from the exemplar, falls back to the |
| diamond pipeline. Default (no exemplar): delegates to |
| shapes.diamond.detect_page(), the legacy passes-1-4 diamond pipeline. |
| """ |
| if exemplar_bgr is not None: |
| exemplar_data = extract_shape_template(exemplar_bgr) |
| if exemplar_data is not None: |
| shape_type = exemplar_data["shape_type"] |
| pipeline = shapes.get_pipeline(shape_type) |
|
|
| shape_dets: list[dict] = [] |
| if pipeline is not None: |
| shape_dets = pipeline( |
| bgr, tile_size=tile_size, overlap=overlap, ocr_upscale=ocr_upscale, |
| progress_dir=progress_dir, progress_key=progress_key, |
| ) |
|
|
| mask_dets = exemplar_template_match_pass( |
| bgr, exemplar_data, |
| ocr_upscale=ocr_upscale, |
| tile_size=tile_size, |
| overlap=overlap, |
| existing_dets=shape_dets, |
| progress_dir=progress_dir, |
| progress_key=progress_key, |
| ) |
| return deduplicate(shape_dets + mask_dets) |
| print("Warning: no usable shape found in exemplar image β " |
| "falling back to the diamond pipeline.") |
|
|
| return shapes.diamond.detect_page( |
| bgr, tile_size=tile_size, overlap=overlap, ocr_upscale=ocr_upscale, |
| progress_dir=progress_dir, progress_key=progress_key, |
| ) |
|
|
|
|
| |
|
|
| def annotate( |
| bgr: np.ndarray, |
| detections: list[dict], |
| valid: set[str] | None = None, |
| missed_boxes: list[tuple[int, int, int, int]] | None = None, |
| ) -> np.ndarray: |
| """ |
| Draw bounding boxes on a copy of the image. |
| |
| Colour scheme when a legend (valid code set) is supplied: |
| Green β confirmed (code is in the legend) |
| Red β flagged (code has a value but is NOT in the legend) |
| Blue β shape-only (code is '?', outline detected but OCR failed) |
| Magenta β missed (see missed_boxes below); drawn thicker so it's |
| unambiguous against the source PDF's own pre-existing red |
| highlight boxes, which this colour scheme otherwise looks |
| nothing like on purpose. |
| |
| Without a legend every box is drawn green (original behaviour). |
| |
| missed_boxes: ground-truth marked boxes (see cv_core.find_marked_boxes / |
| match_marked_boxes) that had no matching detection at all -- i.e. a |
| symbol the source "Marked" PDF already highlighted that this pipeline |
| failed to find. Drawn in a distinct colour/thickness so a reviewer can |
| immediately spot real misses on the annotated page, separate from the |
| confirmed/flagged/shape-only outcomes (which only apply to symbols the |
| pipeline actually detected). |
| """ |
| |
| CLR_CONFIRMED = (0, 180, 0) |
| CLR_FLAGGED = (0, 0, 220) |
| CLR_SHAPE_ONLY = (200, 120, 0) |
| CLR_MISSED = (255, 0, 255) |
|
|
| out = bgr.copy() |
| for d in detections: |
| code = d["code"] |
| if valid is None: |
| colour = CLR_SHAPE_ONLY if code == "?" else CLR_CONFIRMED |
| elif code == "?": |
| colour = CLR_SHAPE_ONLY |
| elif code in valid: |
| colour = CLR_CONFIRMED |
| else: |
| colour = CLR_FLAGGED |
|
|
| x, y, w, h = d["bbox"] |
| cv2.rectangle(out, (x, y), (x + w, y + h), colour, 2) |
| cv2.putText( |
| out, code, |
| (x, max(0, y - 6)), |
| cv2.FONT_HERSHEY_SIMPLEX, 0.6, colour, 2, |
| ) |
|
|
| for (x, y, w, h) in (missed_boxes or []): |
| cv2.rectangle(out, (x, y), (x + w, y + h), CLR_MISSED, 4) |
| cv2.putText( |
| out, "MISSED", |
| (x, max(0, y - 6)), |
| cv2.FONT_HERSHEY_SIMPLEX, 0.6, CLR_MISSED, 2, |
| ) |
|
|
| return out |
|
|
|
|
| |
|
|
| def lower_worker_priority() -> None: |
| """ |
| Pool(initializer=...) target, run once per worker process at startup. |
| Workers are CPU-bound (EasyOCR) and on a 2-vCPU host can fully starve the |
| main Streamlit process of the CPU time it needs to answer WebSocket |
| pings, causing the browser to disconnect/reconnect mid-analysis and |
| restart the whole page loop from scratch. Raising worker niceness gives |
| the OS scheduler a reason to prefer the main process under contention. |
| """ |
| try: |
| os.nice(10) |
| except (AttributeError, OSError): |
| pass |
|
|
|
|
| def detect_page_worker(args: tuple) -> list[dict]: |
| """ |
| Top-level worker for multiprocessing.Pool β must live here (not in app.py) |
| so that spawned worker processes only import this module, never the Streamlit |
| app module. Importing app.py in a worker causes all top-level Streamlit code |
| to run inside the subprocess, which hangs or raises ScriptRunContext errors. |
| """ |
| import time, traceback, os |
| page_bytes, shape, dtype_str, overlap, ocr_upscale, exemplar_bytes, progress_dir, page_idx = args |
| pid = os.getpid() |
| h, w = shape[:2] |
| _ts = lambda: time.strftime("%H:%M:%S") |
| print(f"[worker pid={pid}][{_ts()}] start shape=({h},{w}) dtype={dtype_str}", flush=True) |
|
|
| |
| |
| |
| |
| |
| |
| |
| if progress_dir is not None: |
| try: |
| _canary_path = os.path.join(progress_dir, f"page_{page_idx}.txt") |
| with open(_canary_path, "w") as _cf: |
| _cf.write("worker started") |
| with open(_canary_path) as _cf: |
| _readback = _cf.read() |
| print(f"[worker pid={pid}] progress_dir canary OK: wrote+read " |
| f"'{_readback}' at {_canary_path}", flush=True) |
| except Exception as _e: |
| print(f"[worker pid={pid}] progress_dir canary FAILED at " |
| f"{progress_dir!r}: {_e!r}", flush=True) |
|
|
| t0 = time.time() |
| try: |
| page = np.frombuffer(page_bytes, dtype=np.dtype(dtype_str)).reshape(shape) |
| exemplar_bgr = None |
| if exemplar_bytes is not None: |
| exemplar_bgr = cv2.imdecode(np.frombuffer(exemplar_bytes, np.uint8), cv2.IMREAD_COLOR) |
| dets = detect_page(page, overlap=overlap, ocr_upscale=ocr_upscale, exemplar_bgr=exemplar_bgr, |
| progress_dir=progress_dir, progress_key=page_idx) |
| result = [{"bbox": d["bbox"], "code": d["code"]} for d in dets] |
| print(f"[worker pid={pid}][{_ts()}] done {len(result)} det(s) in {time.time()-t0:.1f}s", flush=True) |
| return result |
| except Exception: |
| print(f"[worker pid={pid}] ERROR after {time.time()-t0:.1f}s:\n{traceback.format_exc()}", flush=True) |
| raise |
|
|
|
|
| |
|
|
| def run( |
| pdf_path: str, |
| output_dir: str, |
| dpi: int = DPI, |
| tile_size: int = TILE_SIZE, |
| ocr_upscale: int = OCR_UPSCALE, |
| legend_path: str | None = None, |
| manual_codes: list[str] | None = None, |
| debug: bool = False, |
| diagnose: bool = False, |
| exemplar_path: str | None = None, |
| dry_run_legend: bool = False, |
| ) -> None: |
| """ |
| legend_path Optional path to a legend/schedule screenshot. When provided |
| the parser extracts valid type-mark codes and each page's |
| detections are split into confirmed / flagged / shape-only. |
| If omitted all detections are reported without validation. |
| manual_codes Additional codes to treat as valid (supplements the legend). |
| Useful when OCR misses a row in the legend image. |
| exemplar_path Optional path to a tightly-cropped example image of the |
| target symbol. When provided, detection switches entirely |
| to exemplar-driven matching instead of the diamond pipeline |
| (see detect_page()). |
| dry_run_legend When True, load the exemplar (to pick the right legend |
| parser) and parse the legend, print the resulting valid |
| codes and any skipped-row diagnostics, then return WITHOUT |
| loading the PDF or running any detection. Legend parsing |
| takes seconds; a full detection run can take hours, so |
| this lets a bad/incomplete legend parse (missing codes, |
| unrecognised header, etc.) be caught and fixed with |
| --manual-codes up front instead of discovered only after |
| the full run finishes. |
| """ |
| os.makedirs(output_dir, exist_ok=True) |
|
|
| exemplar_bgr = cv2.imread(exemplar_path) if exemplar_path else None |
| shape_type: str | None = None |
| if exemplar_path and exemplar_bgr is None: |
| print(f"Warning: could not read exemplar image '{exemplar_path}' β " |
| "running the diamond pipeline instead.\n") |
| elif exemplar_bgr is not None: |
| exemplar_data = extract_shape_template(exemplar_bgr) |
| if exemplar_data is None: |
| print(f"Warning: no usable shape found in '{exemplar_path}' β " |
| "running the diamond pipeline instead.\n") |
| exemplar_bgr = None |
| else: |
| shape_type = exemplar_data["shape_type"] |
| print(f"Exemplar loaded from '{exemplar_path}' " |
| f"(shape_type={shape_type})\n") |
|
|
| |
| |
| |
| |
| |
| |
| |
| |
| def _manual_codes_only() -> set[str]: |
| import re as _re |
| code_re = _re.compile(r"^[A-Z][0-9]{0,2}$") if is_hexagon else _re.compile(r"^[A-Z]\d$") |
| return {c for c in manual_codes if code_re.match(c)} |
|
|
| valid: set[str] | None = None |
| is_hexagon = shape_type == "hexagon" |
| if legend_path: |
| try: |
| if is_hexagon: |
| from legend_parser import parse_hexagon_legend_image, valid_codes_hexagon |
| legend_df = parse_hexagon_legend_image(legend_path, manual_codes=manual_codes) |
| valid = valid_codes_hexagon(legend_df) |
| else: |
| from legend_parser import parse_legend_image, valid_codes |
| legend_df = parse_legend_image(legend_path, manual_codes=manual_codes) |
| valid = valid_codes(legend_df) |
| print(f"Legend loaded from '{legend_path}'") |
| print(f" Valid codes ({len(valid)}): {sorted(valid)}\n") |
| except Exception as exc: |
| |
| |
| |
| |
| if manual_codes: |
| valid = _manual_codes_only() |
| print(f"Warning: could not load legend ({exc}) β " |
| f"using manual code list only: {sorted(valid)}\n") |
| else: |
| print(f"Warning: could not load legend ({exc}) β running without validation.\n") |
| elif manual_codes: |
| valid = _manual_codes_only() |
| print(f"Using manual code list: {sorted(valid)}\n") |
|
|
| if dry_run_legend: |
| if valid is None: |
| print("Dry run: no legend or manual codes loaded β nothing to validate.") |
| else: |
| print(f"Dry run: legend parsing complete, {len(valid)} valid code(s) " |
| "recognised (see any skipped-row warnings above). Skipping detection.") |
| return |
|
|
| pages = load_pages(pdf_path, dpi=dpi) |
| print(f"Loaded {len(pages)} page(s) from '{pdf_path}' at {dpi} DPI") |
| print(f"Tile size: {tile_size}px Overlap: {int(TILE_OVERLAP*100)}%\n") |
|
|
| total_dets: int = 0 |
| global_confirmed: Counter = Counter() |
| global_flagged: Counter = Counter() |
| global_marked_total = 0 |
| global_marked_missed = 0 |
| pages_with_misses: list[tuple[int, int]] = [] |
|
|
| for i, page in enumerate(pages, 1): |
| if diagnose: |
| print(f"Page {i} diagnostics:") |
| shapes.diamond.diagnose_page(page, tile_size=tile_size) |
|
|
| |
| |
| |
| |
| marked_boxes = find_marked_boxes(page) |
|
|
| dets = detect_page(page, tile_size=tile_size, ocr_upscale=ocr_upscale, |
| exemplar_bgr=exemplar_bgr) |
| total_dets += len(dets) |
|
|
| missed_boxes: list[tuple[int, int, int, int]] = [] |
| if marked_boxes: |
| recall = match_marked_boxes(marked_boxes, dets) |
| missed_boxes = recall["missed"] |
| global_marked_total += recall["total_marked"] |
| global_marked_missed += len(missed_boxes) |
| if missed_boxes: |
| pages_with_misses.append((i, len(missed_boxes))) |
| print(f" Page {i:>2} recall: {recall['matched']}/{recall['total_marked']} " |
| f"pre-marked symbol(s) detected" |
| + (f" β {len(missed_boxes)} missed at {missed_boxes}" if missed_boxes else "")) |
|
|
| named = [d["code"] for d in dets if d["code"] != "?"] |
| unknown = sum(1 for d in dets if d["code"] == "?") |
|
|
| if valid is not None: |
| confirmed = [c for c in named if c in valid] |
| flagged = [c for c in named if c not in valid] |
| global_confirmed += Counter(confirmed) |
| global_flagged += Counter(flagged) |
|
|
| print(f" Page {i:>2}: {len(dets):3d} detection(s)" |
| f" ({len(confirmed)} confirmed, {len(flagged)} flagged," |
| f" {unknown} shape-only)") |
| if confirmed: |
| print(f" Confirmed:") |
| for code, n in sorted(Counter(confirmed).items()): |
| print(f" {code}: {n}") |
| if flagged: |
| print(f" Flagged (not in legend):") |
| for code, n in sorted(Counter(flagged).items()): |
| print(f" {code}: {n}") |
| if unknown: |
| print(f" Shape-only (?): {unknown}") |
| else: |
| global_confirmed += Counter(named) |
| print(f" Page {i:>2}: {len(dets):3d} detection(s)" |
| f" ({len(named)} coded, {unknown} shape-only)") |
| for code, n in sorted(Counter(named).items()): |
| print(f" {code}: {n}") |
| if unknown: |
| print(f" ?: {unknown}") |
|
|
| cv2.imwrite(os.path.join(output_dir, f"page_{i:02d}.png"), |
| annotate(page, dets, valid=valid, missed_boxes=missed_boxes)) |
|
|
| if debug: |
| first_tile = next(iter_tiles(page, tile_size, TILE_OVERLAP))[0] |
| cv2.imwrite( |
| os.path.join(output_dir, f"page_{i:02d}_binary_tile0.png"), |
| binarize(first_tile), |
| ) |
|
|
| |
| W = 46 |
| print(f"\n{chr(9552)*W}") |
| print(f" Total detections : {total_dets}") |
|
|
| if valid is not None: |
| total_confirmed = sum(global_confirmed.values()) |
| total_flagged = sum(global_flagged.values()) |
| print(f" Confirmed : {total_confirmed}") |
| print(f" Flagged : {total_flagged}") |
| print(f"{chr(9472)*W}") |
| print(f" {'Code':<8} {'Confirmed':>9} {'Flagged':>7}") |
| print(f"{chr(9472)*W}") |
| all_codes = sorted(set(global_confirmed) | set(global_flagged)) |
| for code in all_codes: |
| marker = "" if code in valid else " !" |
| print(f" {code:<8} {global_confirmed[code]:>9} " |
| f"{global_flagged[code]:>7}{marker}") |
| else: |
| print(f" Unique codes : {len(global_confirmed)}") |
| print(f"{chr(9472)*W}") |
| print(f" {'Code':<8} {'Count':>6}") |
| print(f"{chr(9472)*W}") |
| for code, n in sorted(global_confirmed.items()): |
| print(f" {code:<8} {n:>6}") |
|
|
| if global_marked_total > 0: |
| matched = global_marked_total - global_marked_missed |
| pct = 100.0 * matched / global_marked_total |
| print(f"{chr(9472)*W}") |
| print(f" Ground-truth recall (pre-marked source boxes):") |
| print(f" {matched}/{global_marked_total} detected ({pct:.1f}%), " |
| f"{global_marked_missed} missed") |
| if pages_with_misses: |
| print(f" Missed on page(s): " |
| + ", ".join(f"{p} ({n})" for p, n in pages_with_misses)) |
| print(f" See magenta \"MISSED\" boxes in the annotated output for exact locations.") |
|
|
| print(f"{chr(9552)*W}") |
| print(f"\nAnnotated pages saved to '{output_dir}'") |
|
|
|
|
| if __name__ == "__main__": |
| parser = argparse.ArgumentParser(description="Classical CV target-symbol detector for blueprints") |
| parser.add_argument("--pdf", default="Data/Example/Subset.pdf") |
| parser.add_argument("--out", default="Data/Testing/Output/classical_cv") |
| parser.add_argument("--dpi", type=int, default=DPI, |
| help=f"PDF render DPI (default {DPI})") |
| parser.add_argument("--tile-size", type=int, default=TILE_SIZE, |
| help=f"Sliding-window tile size in px (default {TILE_SIZE})") |
| parser.add_argument("--ocr-upscale", type=int, default=OCR_UPSCALE, |
| help=f"Upscale factor before EasyOCR (default {OCR_UPSCALE})") |
| parser.add_argument("--legend", default=None, |
| help="Path to legend/schedule screenshot for validation (optional)") |
| parser.add_argument("--manual-codes", default=None, |
| help="Comma-separated codes to add to the legend (e.g. E4,F1,F2,J1)") |
| parser.add_argument("--debug", action="store_true", |
| help="Save binary tile image alongside annotated output") |
| parser.add_argument("--diagnose", action="store_true", |
| help="Print per-filter drop counts to diagnose zero detections") |
| parser.add_argument("--exemplar", default=None, |
| help="Path to a tightly-cropped example image of the target symbol. " |
| "When given, detection uses this shape instead of the diamond pipeline.") |
| parser.add_argument("--dry-run-legend", action="store_true", |
| help="Parse the exemplar + legend, print the recognised valid codes and " |
| "any skipped-row warnings, then exit without running detection. " |
| "Use this to sanity-check --legend/--manual-codes before committing " |
| "to a full (potentially hours-long) run.") |
| args = parser.parse_args() |
|
|
| manual = [c.strip() for c in args.manual_codes.split(",")] if args.manual_codes else None |
|
|
| run( |
| args.pdf, args.out, |
| dpi=args.dpi, |
| tile_size=args.tile_size, |
| ocr_upscale=args.ocr_upscale, |
| legend_path=args.legend, |
| manual_codes=manual, |
| debug=args.debug, |
| diagnose=args.diagnose, |
| exemplar_path=args.exemplar, |
| dry_run_legend=args.dry_run_legend, |
| ) |
|
|