Add project files
This commit is contained in:
Suman
353
testVaseMatcher.py
Normal file
353
testVaseMatcher.py
Normal file
@@ -0,0 +1,353 @@
|
||||
"""
|
||||
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()
|
||||
Reference in New Issue
Block a user