"""Debug and threshold-tuning script for pattern detection pipeline.""" import argparse import os import sys import time import cv2 import numpy as np sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..")) from src.pipeline import PatternDetectionPipeline def find_example_pairs(examples_dir: str) -> list: """Find (pattern, drawing) file pairs in examples_dir.""" pairs = {} for fname in sorted(os.listdir(examples_dir)): if not fname.lower().endswith((".png", ".jpg", ".jpeg", ".tiff")): continue name_no_ext = os.path.splitext(fname)[0] parts = name_no_ext.split("_") if len(parts) < 2: continue key = parts[0] tag = "_".join(parts[1:]) pairs.setdefault(key, {})[tag] = os.path.join(examples_dir, fname) result = [] for key in sorted(pairs): imgs = pairs[key] if "pattern" in imgs and "drawing" in imgs: result.append((key, imgs["pattern"], imgs["drawing"])) return result def save_visualization(output_dir: str, name: str, vis: np.ndarray): os.makedirs(output_dir, exist_ok=True) out_path = os.path.join(output_dir, f"{name}.png") cv2.imwrite(out_path, cv2.cvtColor(vis, cv2.COLOR_RGB2BGR)) def main(): parser = argparse.ArgumentParser(description="Tune thresholds for pattern detection pipeline.") parser.add_argument("--examples_dir", default="examples") parser.add_argument("--output_dir", default="debug_output") parser.add_argument("--max_pairs", type=int, default=3) parser.add_argument("--quick", action="store_true", help="Only test NCC=0.55, cosine=0.75") args = parser.parse_args() if not os.path.isdir(args.examples_dir): print(f"[ERROR] Examples directory not found: {args.examples_dir}") sys.exit(1) pairs = find_example_pairs(args.examples_dir) if not pairs: print(f"[ERROR] No pattern/drawing pairs found in {args.examples_dir}") sys.exit(1) pairs = pairs[: args.max_pairs] if args.quick: ncc_thresholds = [0.55] cosine_thresholds = [0.75] else: ncc_thresholds = [0.45, 0.50, 0.55, 0.60, 0.65] cosine_thresholds = [0.65, 0.70, 0.75, 0.80] # Header header = f"{'Example':<12} {'NCC':>6} {'Cosine':>8} {'#Det':>6} {'Time(s)':>9}" print("\n" + header) print("-" * len(header)) pipeline = PatternDetectionPipeline() total_runs = len(pairs) * len(ncc_thresholds) * len(cosine_thresholds) run_idx = 0 for key, pattern_path, drawing_path in pairs: for ncc_t in ncc_thresholds: for cos_t in cosine_thresholds: run_idx += 1 print(f"\r[{run_idx}/{total_runs}] Running...", end="", flush=True) pipeline.update_thresholds(ncc_threshold=ncc_t, cosine_threshold=cos_t) t0 = time.time() try: result = pipeline.detect(pattern_path, drawing_path, return_visualization=True) except Exception as e: print(f"\r[ERROR] {key} ncc={ncc_t} cos={cos_t}: {e}") continue elapsed = time.time() - t0 n_det = result["total_detections"] row = f"{key:<12} {ncc_t:>6.2f} {cos_t:>8.2f} {n_det:>6} {elapsed:>9.2f}" print(f"\r{row}") if "visualization" in result: tag = f"{key}_ncc{ncc_t:.2f}_cos{cos_t:.2f}" save_visualization(args.output_dir, tag, result["visualization"]) print(f"\n[Done] Visualizations saved to: {args.output_dir}/") if __name__ == "__main__": main()