Add project files
This commit is contained in:
Suman
284
tester.py
Normal file
284
tester.py
Normal file
@@ -0,0 +1,284 @@
|
||||
"""
|
||||
tester.py -- simple batch accuracy check for the Vase Matcher pipeline.
|
||||
|
||||
Expects a folder of subfolders, each named after an existing template (e.g.
|
||||
SKU_1/, SKU_2/, ...), each containing photos expected to match that
|
||||
same-named template:
|
||||
|
||||
flowers/
|
||||
SKU_1/ photo1.jpg photo2.jpg ...
|
||||
SKU_2/ ...
|
||||
|
||||
For every image, runs the exact same code path the web app's /api/match
|
||||
uses (pipeline.engine.process_upload -- no reimplementation of the scoring
|
||||
logic here), checks whether the weighted final match picked the expected
|
||||
template, then (unless --skip-flowers) also runs the same SAM3-based flower
|
||||
count + vase comparison the "Count flowers" button runs
|
||||
(pipeline.engine.count_flowers) and saves one combined report image per
|
||||
test photo to --report-dir, showing side by side:
|
||||
|
||||
- your photo's flower segmentation vs the matched template's
|
||||
- flower counts on both sides (+ the same mismatch note the web UI shows)
|
||||
- the vase-identity comparison (DINOv2 + CLIP crops + verdict)
|
||||
|
||||
Prints a pass/fail line per image, then a per-SKU and overall accuracy
|
||||
summary. Note: the flower/vase step adds real time per image (SAM3 loads
|
||||
in a subprocess, plus YOLO-World and DINOv2/CLIP) -- expect roughly
|
||||
15-25s/image on top of the ~2-3s the match itself takes; use --skip-flowers
|
||||
for a fast accuracy-only pass.
|
||||
|
||||
Usage:
|
||||
python3 tester.py # uses ./flowers
|
||||
python3 tester.py /path/to/flowers
|
||||
python3 tester.py --limit 3 # only first 3 images per SKU folder
|
||||
python3 tester.py --skip-flowers # fast pass, no report images
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import gc
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
|
||||
import config # noqa: F401 -- must import first, caps thread envs
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from pipeline import engine
|
||||
|
||||
DEFAULT_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "flowers")
|
||||
DEFAULT_REPORT_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "test_reports")
|
||||
|
||||
REPORT_W = 900
|
||||
IMG_ROW_H = 340
|
||||
VASE_ROW_H = 200
|
||||
BAR_H = 34
|
||||
BG = (24, 22, 19)
|
||||
PANEL_BG = (250, 250, 250)
|
||||
|
||||
|
||||
def release_gpu():
|
||||
gc.collect()
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
|
||||
def _fit_into(img, w, h):
|
||||
canvas = np.full((h, w, 3), BG, dtype=np.uint8)
|
||||
if img is None:
|
||||
cv2.putText(canvas, "unavailable", (10, h // 2), cv2.FONT_HERSHEY_SIMPLEX,
|
||||
0.6, (140, 140, 140), 1, cv2.LINE_AA)
|
||||
return canvas
|
||||
ih, iw = img.shape[:2]
|
||||
scale = min(w / iw, h / ih)
|
||||
nw, nh = max(1, int(iw * scale)), max(1, int(ih * scale))
|
||||
resized = cv2.resize(img, (nw, nh), interpolation=cv2.INTER_AREA)
|
||||
x_off, y_off = (w - nw) // 2, (h - nh) // 2
|
||||
canvas[y_off:y_off + nh, x_off:x_off + nw] = resized
|
||||
return canvas
|
||||
|
||||
|
||||
def _wrap_text(text, max_width, font=cv2.FONT_HERSHEY_SIMPLEX, scale=0.52, thickness=1):
|
||||
words = text.split()
|
||||
lines, current = [], ""
|
||||
for word in words:
|
||||
trial = f"{current} {word}".strip()
|
||||
(tw, _), _ = cv2.getTextSize(trial, font, scale, thickness)
|
||||
if tw > max_width and current:
|
||||
lines.append(current)
|
||||
current = word
|
||||
else:
|
||||
current = trial
|
||||
if current:
|
||||
lines.append(current)
|
||||
return lines
|
||||
|
||||
|
||||
def _text_panel(lines, w, line_h=22, pad=10, color=(30, 30, 30)):
|
||||
h = pad * 2 + max(1, len(lines)) * line_h
|
||||
canvas = np.full((h, w, 3), PANEL_BG, dtype=np.uint8)
|
||||
for i, line in enumerate(lines):
|
||||
y = pad + (i + 1) * line_h - 6
|
||||
cv2.putText(canvas, line, (16, y), cv2.FONT_HERSHEY_SIMPLEX, 0.52, color, 1, cv2.LINE_AA)
|
||||
return canvas
|
||||
|
||||
|
||||
def _bar(text, w, color=(30, 30, 30), bg=PANEL_BG, h=BAR_H):
|
||||
canvas = np.full((h, w, 3), bg, dtype=np.uint8)
|
||||
cv2.putText(canvas, text, (16, h - 12), cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 1, cv2.LINE_AA)
|
||||
return canvas
|
||||
|
||||
|
||||
def build_report_image(request_dir, fname, expected, picked, correct, score,
|
||||
flower_result, flower_count_comparison):
|
||||
half_w = REPORT_W // 2
|
||||
|
||||
def load(rel):
|
||||
if not rel:
|
||||
return None
|
||||
p = os.path.join(request_dir, rel)
|
||||
return cv2.imread(p) if os.path.isfile(p) else None
|
||||
|
||||
sam = flower_result.get("sam", {})
|
||||
sam_t = flower_result.get("sam_template", {})
|
||||
vase = flower_result.get("vase_comparison") or {}
|
||||
|
||||
status = "MATCH" if correct else "MISMATCH"
|
||||
status_color = (90, 140, 90) if correct else (60, 60, 200)
|
||||
top_bar = _bar(
|
||||
f"{fname} expected={expected} picked={picked} [{status}] weighted_score={score}",
|
||||
REPORT_W, color=status_color,
|
||||
)
|
||||
|
||||
label_bar = np.full((BAR_H, REPORT_W, 3), PANEL_BG, dtype=np.uint8)
|
||||
cv2.putText(label_bar, f"Your upload -- {sam.get('total_count', 0)} flowers",
|
||||
(16, BAR_H - 12), cv2.FONT_HERSHEY_SIMPLEX, 0.56, (30, 30, 30), 1, cv2.LINE_AA)
|
||||
cv2.putText(label_bar, f"{picked} -- {sam_t.get('total_count', 0)} flowers",
|
||||
(half_w + 16, BAR_H - 12), cv2.FONT_HERSHEY_SIMPLEX, 0.56, (30, 30, 30), 1, cv2.LINE_AA)
|
||||
|
||||
seg_row = cv2.hconcat([
|
||||
_fit_into(load(sam.get("visual_file")), half_w, IMG_ROW_H),
|
||||
_fit_into(load(sam_t.get("visual_file")), half_w, IMG_ROW_H),
|
||||
])
|
||||
|
||||
if flower_count_comparison and not flower_count_comparison.get("error"):
|
||||
fc_lines = _wrap_text(flower_count_comparison["message"], REPORT_W - 32)
|
||||
elif flower_count_comparison:
|
||||
fc_lines = [f"Flower-count comparison unavailable: {flower_count_comparison['error']}"]
|
||||
else:
|
||||
fc_lines = ["Flower-count comparison not run (no matched template)."]
|
||||
fc_panel = _text_panel(fc_lines, REPORT_W)
|
||||
|
||||
vase_row = cv2.hconcat([
|
||||
_fit_into(load(vase.get("input_crop_file")), half_w, VASE_ROW_H),
|
||||
_fit_into(load(vase.get("template_crop_file")), half_w, VASE_ROW_H),
|
||||
])
|
||||
if vase.get("error"):
|
||||
vase_lines = [f"Vase comparison unavailable: {vase['error']}"]
|
||||
elif vase:
|
||||
vase_lines = [
|
||||
f"Vase match: {vase.get('verdict', '?').upper()} (DINOv2 "
|
||||
f"{vase.get('dino_similarity_pct')}%, CLIP {vase.get('clip_similarity_pct')}%, "
|
||||
f"combined {vase.get('combined_pct')}%)"
|
||||
]
|
||||
else:
|
||||
vase_lines = ["Vase comparison not run (no matched template)."]
|
||||
vase_bar = _text_panel(vase_lines, REPORT_W)
|
||||
|
||||
return cv2.vconcat([top_bar, label_bar, seg_row, fc_panel, vase_bar, vase_row])
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description=__doc__,
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter)
|
||||
parser.add_argument("test_dir", nargs="?", default=DEFAULT_DIR,
|
||||
help=f"Folder of per-template subfolders (default: {DEFAULT_DIR})")
|
||||
parser.add_argument("--limit", type=int, default=None,
|
||||
help="Only test the first N images per folder")
|
||||
parser.add_argument("--report-dir", default=DEFAULT_REPORT_DIR,
|
||||
help=f"Where to save combined report images (default: {DEFAULT_REPORT_DIR})")
|
||||
parser.add_argument("--skip-flowers", action="store_true",
|
||||
help="Skip the SAM3 flower-count/vase-comparison step and report images "
|
||||
"(much faster -- accuracy summary only)")
|
||||
args = parser.parse_args()
|
||||
|
||||
if not os.path.isdir(args.test_dir):
|
||||
print(f"Test folder not found: {args.test_dir}")
|
||||
sys.exit(1)
|
||||
|
||||
if not args.skip_flowers:
|
||||
os.makedirs(args.report_dir, exist_ok=True)
|
||||
|
||||
print(f"Bootstrapping templates from {config.TEMPLATE_IMAGES_DIR} ...")
|
||||
engine.bootstrap()
|
||||
known_templates = set(engine.templates_meta().keys())
|
||||
|
||||
per_sku = {}
|
||||
total_correct = 0
|
||||
total_count = 0
|
||||
start = time.perf_counter()
|
||||
|
||||
for expected in sorted(os.listdir(args.test_dir)):
|
||||
folder = os.path.join(args.test_dir, expected)
|
||||
if not os.path.isdir(folder):
|
||||
continue
|
||||
if expected not in known_templates:
|
||||
print(f"Skipping {expected}/ -- no template with that name")
|
||||
continue
|
||||
|
||||
images = sorted(f for f in os.listdir(folder) if f.lower().endswith(config.VALID_EXTS))
|
||||
if args.limit:
|
||||
images = images[:args.limit]
|
||||
|
||||
sku_correct = 0
|
||||
print(f"\n=== {expected}: {len(images)} image(s) ===")
|
||||
|
||||
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:
|
||||
result = engine.process_upload(image_bytes, fname)
|
||||
except Exception as e:
|
||||
print(f" [ERR ] {fname}: {e}")
|
||||
release_gpu()
|
||||
continue
|
||||
|
||||
picked = result.get("weighted_best")
|
||||
score = next(
|
||||
(row["weighted_score"] for row in result.get("weighted_scores", [])
|
||||
if row["template"] == picked),
|
||||
None,
|
||||
)
|
||||
correct = picked == expected
|
||||
sku_correct += int(correct)
|
||||
total_correct += int(correct)
|
||||
total_count += 1
|
||||
|
||||
report_note = ""
|
||||
if not args.skip_flowers:
|
||||
try:
|
||||
flower_result = engine.count_flowers(result["request_id"], template=picked)
|
||||
request_dir = os.path.join(config.UPLOADS_DIR, result["request_id"])
|
||||
report = build_report_image(
|
||||
request_dir, fname, expected, picked, correct, score,
|
||||
flower_result, flower_result.get("flower_count_comparison"),
|
||||
)
|
||||
stem = os.path.splitext(fname)[0]
|
||||
report_path = os.path.join(args.report_dir, f"{expected}_{stem}.png")
|
||||
cv2.imwrite(report_path, report)
|
||||
report_note = f" [report: {os.path.basename(report_path)}]"
|
||||
except Exception as e:
|
||||
report_note = f" [report failed: {e}]"
|
||||
release_gpu()
|
||||
|
||||
img_elapsed = time.perf_counter() - img_start
|
||||
status = "OK " if correct else "MISS"
|
||||
print(f" [{status}] {fname:<40} -> {picked!s:<14} "
|
||||
f"(score {score}, {img_elapsed:.2f}s){report_note}")
|
||||
|
||||
release_gpu()
|
||||
|
||||
per_sku[expected] = (sku_correct, len(images))
|
||||
|
||||
elapsed = time.perf_counter() - start
|
||||
|
||||
print("\n" + "=" * 46)
|
||||
print(f"{'SKU':<20}{'Correct':>10}{'Total':>8}{'Accuracy':>8}")
|
||||
print("-" * 46)
|
||||
for sku, (correct, total) in per_sku.items():
|
||||
acc = (correct / total * 100) if total else 0.0
|
||||
print(f"{sku:<20}{correct:>10}{total:>8}{acc:>7.1f}%")
|
||||
print("=" * 46)
|
||||
overall_acc = (total_correct / total_count * 100) if total_count else 0.0
|
||||
print(f"Overall: {total_correct}/{total_count} correct ({overall_acc:.1f}%) in {elapsed:.1f}s")
|
||||
if not args.skip_flowers:
|
||||
print(f"Report images: {os.path.abspath(args.report_dir)}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user