354 lines
13 KiB
Python
354 lines
13 KiB
Python
"""
|
|
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()
|