File size: 3,682 Bytes
7d835bb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 | """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()
|