""" testVaseMatcher.py -- batch-test the Vase Matcher pipeline against a labeled test set. Expects a folder of subfolders, each named after an existing template (e.g. SKU_1/, SKU_2/, ...), each containing photos that are expected to match that same-named template: TEST_SKU_S/ SKU_1/ photo1.jpg photo2.jpg ... SKU_2/ ... For every test image, runs the exact same per-image computation the web app uses -- SIFT, ORB, SuperPoint+LightGlue, LoFTR, the weighted verdict (now including Color at weight 0.7, same as config.METHOD_WEIGHTS), the Borda-count overall verdict (still just the original four methods, Color doesn't participate there), the color-space comparison, and the color family-grid area-match (numbers only here -- no grid images are rendered for a bulk run) -- against the live template set (config.TEMPLATE_IMAGES_DIR), then reports whether each method's pick (and the color-space pick) matches the expected template. Writes a detailed per-image CSV plus a per-method accuracy summary. Usage: cd /media/suman/Backup_of_extra_/Sasi/featureTransform python3 testVaseMatcher.py python3 testVaseMatcher.py /path/to/other/test/set --csv my_report.csv python3 testVaseMatcher.py --limit 3 # only first 3 images/folder Resource notes -- this reuses the exact same pipeline modules the Flask app uses, so the same CUDA-OOM safeguards apply automatically: - every image is downscaled to config.MAX_IMAGE_DIM before touching any model (bg removal, SIFT/ORB, SuperPoint, LoFTR all cap to this) - images are processed strictly one at a time, and within each image the four methods run sequentially (not concurrently) -- there is never more than one heavy CPU/GPU operation in flight - torch.cuda.empty_cache() runs after every method AND after every image - a failure on one image or one method (including a CUDA OOM) is caught, logged, and recorded as "no result" for that cell instead of aborting the whole run Caveat: this script loads its own copy of every model onto the GPU. If the Flask app (app.py) is ALSO running at the same time, that's two separate processes both holding GPU memory -- normally fine (each is a few hundred MB), but worth knowing if you ever do see a CUDA OOM here. The script prints current GPU memory usage at startup so you can check before a big run. """ import argparse import csv import gc import logging import os import sys import time import warnings import config # noqa: F401 -- must import first, caps thread envs / sets PYTORCH_CUDA_ALLOC_CONF import cv2 import torch from PIL import Image from pipeline import bg_removal, color, color_grid, engine warnings.filterwarnings("ignore", category=Image.DecompressionBombWarning) logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)-7s [%(name)s] %(message)s") logger = logging.getLogger("testVaseMatcher") DEFAULT_TEST_DIR = "/media/suman/Backup_of_extra_/Sasi/TEST_SKU_S" def log_gpu_memory(): if not torch.cuda.is_available(): return try: import subprocess out = subprocess.check_output( ["nvidia-smi", "--query-compute-apps=pid,used_memory", "--format=csv,noheader"], text=True, timeout=5, ).strip() logger.info("Current GPU memory in use by other processes:\n%s", out or "(none)") except Exception: pass # purely advisory, never fatal def release_gpu(): gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() def list_expected_folders(test_dir): return sorted( d for d in os.listdir(test_dir) if os.path.isdir(os.path.join(test_dir, d)) ) def list_test_images(folder, limit=None): files = sorted( f for f in os.listdir(folder) if f.lower().endswith(config.VALID_EXTS) ) return files[:limit] if limit else files def _timed_with_oom_retry(method, fn, bgr, mask, max_retries=1): """Same contract as engine._timed (returns results, elapsed, error) but retries once after a hard cache-clear if the failure was a CUDA OOM -- that's usually transient allocator fragmentation from back-to-back images of varying sizes, not a real per-image failure, and a tight 56-image loop is far more likely to hit it than the web app's sporadic single requests. Any other kind of error still fails immediately, same as engine._timed.""" start = time.perf_counter() attempt = 0 while True: try: result = fn(bgr, mask) error = None break except Exception as e: is_oom = "out of memory" in str(e).lower() if is_oom and attempt < max_retries: logger.warning("%s hit a CUDA OOM -- clearing cache and retrying once...", method) release_gpu() time.sleep(1) attempt += 1 continue logger.exception("Method %s failed", method) result = [] error = str(e) break elapsed = time.perf_counter() - start return result, elapsed, error def run_single_image(image_bytes, cache_dir): """Mirrors pipeline.engine._process_upload_locked's core computation, minus the request-folder/annotated-image file writes that only matter for the web UI -- this just needs the numbers.""" rgba, _ = bg_removal.remove_background_bytes(image_bytes, cache_dir) bgr, _, mask = bg_removal.split_rgba(rgba) method_results = {} for method in engine.METHODS: results, elapsed, error = _timed_with_oom_retry( method, engine._METHOD_RUNNERS[method], bgr, mask ) method_results[method] = {"results": results, "time_sec": elapsed, "error": error} release_gpu() # Borda-count overall verdict: unchanged, still just the original four # methods -- Color doesn't participate here, same as engine.py. overall_best = engine._overall_best(method_results) try: color_analysis = color.compare_input_to_templates(bgr, mask) color_error = None except Exception: logger.exception("color analysis failed") color_analysis = {"input_dominant_colors": [], "templates": []} color_error = "color analysis failed" # Weighted verdict now folds Color in too (config.METHOD_WEIGHTS["Color"]), # via the exact same reshaping engine._process_upload_locked uses -- kept # separate from method_results so it still doesn't show up as if it were # a 5th method result anywhere else in this script. methods_with_color = dict(method_results) methods_with_color["Color"] = engine._color_as_method_result(color_analysis, 0.0, color_error) weighted_best, _weighted_scores = engine._weighted_scores(methods_with_color) # Color family grid: numbers only here (no grid images rendered/saved # for a bulk run) -- how well the winning template's color regions # actually line up with the upload's, area-wise. family_area_match_pct = None if weighted_best is not None: try: input_families = color_grid.cluster_families(bgr, mask) template_rgba = cv2.imread( engine.templates_meta()[weighted_best]["nobg_path"], cv2.IMREAD_UNCHANGED ) template_bgr, _, template_mask = bg_removal.split_rgba(template_rgba) template_families = color_grid.cluster_families(template_bgr, template_mask) matches = color_grid.match_families(input_families, template_families) family_area_match_pct = color_grid.overall_area_match(matches) except Exception: logger.exception("family grid failed") return { "method_results": method_results, "weighted_best": weighted_best, "overall_best": overall_best, "color_analysis": color_analysis, "family_area_match_pct": family_area_match_pct, } def top_pick(results_list): return results_list[0]["template"] if results_list else None def color_top_pick(color_analysis): templates = color_analysis.get("templates", []) if not templates: return None, None top = max(templates, key=lambda t: t["match_pct"]) return top["template"], top["match_pct"] def color_pct_for(color_analysis, template_name): for t in color_analysis.get("templates", []): if t["template"] == template_name: return t["match_pct"] return None def new_tally(): return {"correct": 0, "total": 0} def record(tally_entry, correct): tally_entry["total"] += 1 tally_entry["correct"] += int(bool(correct)) def main(): parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) parser.add_argument("test_dir", nargs="?", default=DEFAULT_TEST_DIR, help=f"Folder of per-template subfolders (default: {DEFAULT_TEST_DIR})") parser.add_argument("--csv", default="test_results.csv", help="Where to write the detailed per-image CSV report") parser.add_argument("--limit", type=int, default=None, help="Only test the first N images per folder (quick smoke test)") args = parser.parse_args() if not os.path.isdir(args.test_dir): logger.error("Test folder not found: %s", args.test_dir) sys.exit(1) log_gpu_memory() logger.info("Bootstrapping templates from %s ...", config.TEMPLATE_IMAGES_DIR) engine.bootstrap() expected_folders = list_expected_folders(args.test_dir) known_templates = set(engine.templates_meta().keys()) unknown = [f for f in expected_folders if f not in known_templates] if unknown: logger.warning("These test folders don't match any known template " "and will be skipped: %s", unknown) method_names = list(engine.METHODS) tally = {m: new_tally() for m in method_names} tally["Weighted"] = new_tally() tally["Overall"] = new_tally() tally["Color"] = new_tally() rows = [] total_start = time.perf_counter() for expected in expected_folders: if expected not in known_templates: continue folder = os.path.join(args.test_dir, expected) images = list_test_images(folder, args.limit) logger.info("=== %s: %d test image(s), expected match = %s ===", expected, len(images), expected) for fname in images: path = os.path.join(folder, fname) with open(path, "rb") as f: image_bytes = f.read() img_start = time.perf_counter() try: outcome = run_single_image(image_bytes, config.UPLOADS_NOBG_CACHE) except Exception: logger.exception("FAILED on %s/%s -- skipping", expected, fname) release_gpu() continue img_elapsed = time.perf_counter() - img_start row = {"expected": expected, "file": fname, "time_sec": round(img_elapsed, 2)} for method in method_names: mr = outcome["method_results"][method] pick = top_pick(mr["results"]) score = mr["results"][0]["score"] if mr["results"] else None correct = pick == expected row[f"{method}_pick"] = pick row[f"{method}_score"] = score row[f"{method}_correct"] = correct record(tally[method], correct) weighted_correct = outcome["weighted_best"] == expected row["weighted_pick"] = outcome["weighted_best"] row["weighted_correct"] = weighted_correct record(tally["Weighted"], weighted_correct) overall_correct = outcome["overall_best"] == expected row["overall_pick"] = outcome["overall_best"] row["overall_correct"] = overall_correct record(tally["Overall"], overall_correct) c_pick, c_pct = color_top_pick(outcome["color_analysis"]) c_correct = c_pick == expected row["color_pick"] = c_pick row["color_match_pct"] = c_pct row["color_pct_for_expected"] = color_pct_for(outcome["color_analysis"], expected) row["color_correct"] = c_correct record(tally["Color"], c_correct) row["family_area_match_pct"] = outcome.get("family_area_match_pct") rows.append(row) status = "OK " if weighted_correct else "MISS" logger.info(" [%s] %-24s weighted -> %-14s color -> %-14s (%.2fs)", status, fname, outcome["weighted_best"], c_pick, img_elapsed) release_gpu() total_elapsed = time.perf_counter() - total_start if rows: with open(args.csv, "w", newline="") as f: writer = csv.DictWriter(f, fieldnames=list(rows[0].keys())) writer.writeheader() writer.writerows(rows) logger.info("Wrote detailed report: %s", os.path.abspath(args.csv)) print("\n" + "=" * 58) print(f"{'Method':<12}{'Correct':>10}{'Total':>8}{'Accuracy':>13}") print("-" * 58) for name, t in tally.items(): acc = (t["correct"] / t["total"] * 100) if t["total"] else 0.0 print(f"{name:<12}{t['correct']:>10}{t['total']:>8}{acc:>12.1f}%") print("=" * 58) print(f"Images tested: {len(rows)} Total time: {total_elapsed:.1f}s") if __name__ == "__main__": main()