Upload test_all_prompts.py with huggingface_hub

test_all_prompts.py

@@ -0,0 +1,139 @@
+"""Run CountGD++ on every prompt combination and save a results table + detection images.
+
+Positive prompt:  text | exemplars | text + exemplars
+Negative prompt:  none | text | exemplars | text + exemplars
+Exemplars from:   the counted image itself ("internal") or a separate image ("external")
+Targets:          strawberry (negative: blueberry) and blueberry (negative: strawberry)
+
+    python test_all_prompts.py [--repo nikigoli/CountGDPlusPlus] [--device cuda] [--out prompt_tests]
+
+Writes <out>/results.csv, <out>/summary.png (all cases on one sheet) and one
+image per case. img/336.jpg has about 16 strawberries and well over 100 blueberries.
+"""
+import argparse
+import csv
+import itertools
+import os
+
+import torch
+from PIL import Image, ImageDraw
+
+from countgdpp import CountGDPlusPlus, count
+
+HERE = os.path.dirname(os.path.abspath(__file__))
+IMG = os.path.join(HERE, "img")
+IMAGE = os.path.join(IMG, "336.jpg")
+
+# Visual exemplars as [x1, y1, x2, y2] pixel boxes.
+EXEMPLARS = {
+    "strawberry": {
+        "internal": (IMAGE, [[158, 62, 206, 134], [95, 133, 145, 184]]),
+        "external": (os.path.join(IMG, "synthetic_strawberry.jpg"), [[205, 95, 822, 915]]),
+    },
+    "blueberry": {
+        "internal": (IMAGE, [[235, 38, 259, 62], [330, 165, 356, 192]]),
+        "external": (os.path.join(IMG, "synthetic_blueberry.jpg"), [[143, 140, 878, 850]]),
+    },
+}
+OTHER = {"strawberry": "blueberry", "blueberry": "strawberry"}
+POS_MODES = ["text", "exemplars", "text+exemplars"]
+NEG_MODES = ["none", "text", "exemplars", "text+exemplars"]
+SOURCES = ["internal", "external"]
+
+
+def build_cases():
+    cases, seen = [], set()
+    for target, pos, neg, src in itertools.product(OTHER, POS_MODES, NEG_MODES, SOURCES):
+        uses_exemplars = "exemplars" in pos or "exemplars" in neg
+        key = (target, pos, neg, src if uses_exemplars else "-")
+        if key in seen:
+            continue
+        seen.add(key)
+        cases.append(key)
+    # Two negatives at once (text + exemplars for the other fruit, text only for "plate").
+    for target in OTHER:
+        cases.append((target, "text+exemplars", "two negatives", "internal"))
+    return cases
+
+
+def run_case(model, target, pos, neg, src):
+    kwargs = {"text": target if "text" in pos else ""}
+    if "exemplars" in pos:
+        img, boxes = EXEMPLARS[target][src]
+        kwargs.update(exemplar_boxes=boxes, exemplar_image=img)
+
+    other = OTHER[target]
+    if neg == "two negatives":
+        img, boxes = EXEMPLARS[other][src]
+        kwargs.update(negative_texts=[other, "plate"],
+                      negative_exemplar_boxes=[boxes, None],
+                      negative_exemplar_images=[img, None])
+    elif neg != "none":
+        kwargs["negative_texts"] = [other if "text" in neg else ""]
+        if "exemplars" in neg:
+            img, boxes = EXEMPLARS[other][src]
+            kwargs.update(negative_exemplar_boxes=[boxes], negative_exemplar_images=[img])
+    return count(model, IMAGE, **kwargs)
+
+
+def draw(result, title):
+    im = Image.open(IMAGE).convert("RGB")
+    w, h = im.size
+    d = ImageDraw.Draw(im)
+    for cx, cy, bw, bh in result["boxes"]:
+        d.rectangle([(cx - bw / 2) * w, (cy - bh / 2) * h, (cx + bw / 2) * w, (cy + bh / 2) * h],
+                    outline="cyan", width=2)
+    canvas = Image.new("RGB", (w, h + 22), "white")
+    canvas.paste(im, (0, 22))
+    ImageDraw.Draw(canvas).text((4, 5), title, fill="black")
+    return canvas
+
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument("--repo", default="nikigoli/CountGDPlusPlus")
+    p.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
+    p.add_argument("--out", default="prompt_tests")
+    args = p.parse_args()
+    os.makedirs(args.out, exist_ok=True)
+
+    model = CountGDPlusPlus.from_pretrained(args.repo).to(args.device).eval()
+
+    rows, tiles = [], []
+    cases = build_cases()
+    print(f"{'#':>3}  {'target':<10} {'positive':<15} {'negative':<15} {'exemplars':<9} {'count':>5}")
+    for i, (target, pos, neg, src) in enumerate(cases, 1):
+        try:
+            res = run_case(model, target, pos, neg, src)
+            n, err = res["count"], ""
+        except Exception as e:  # keep going so one failure doesn't hide the rest
+            res, n, err = None, "", f"{type(e).__name__}: {e}"
+        print(f"{i:>3}  {target:<10} {pos:<15} {neg:<15} {src:<9} {n!s:>5}  {err}")
+        rows.append({"case": i, "target": target, "positive": pos, "negative": neg,
+                     "exemplar_source": src, "count": n, "error": err})
+        if res is not None:
+            tile = draw(res, f"#{i} {target} | pos: {pos} | neg: {neg} | {src} -> {n}")
+            tile.save(os.path.join(args.out, f"{i:02d}_{target}_{pos}_{neg}_{src}.png".replace("+", "-").replace(" ", "-")))
+            tiles.append(tile)
+
+    with open(os.path.join(args.out, "results.csv"), "w", newline="") as f:
+        writer = csv.DictWriter(f, fieldnames=list(rows[0]))
+        writer.writeheader()
+        writer.writerows(rows)
+
+    if tiles:
+        cols = 4
+        tw, th = tiles[0].size
+        sheet = Image.new("RGB", (cols * tw, -(-len(tiles) // cols) * th), "white")
+        for k, t in enumerate(tiles):
+            sheet.paste(t, ((k % cols) * tw, (k // cols) * th))
+        sheet.save(os.path.join(args.out, "summary.png"))
+
+    failed = [r for r in rows if r["error"]]
+    print(f"\n{len(rows) - len(failed)}/{len(rows)} cases ran. Results in {args.out}/")
+    if failed:
+        raise SystemExit(f"{len(failed)} case(s) failed, see the error column.")
+
+
+if __name__ == "__main__":
+    main()