from __future__ import annotations
import argparse
import base64
import json
import logging
import os
import sys
from pathlib import Path
LOG = logging.getLogger("captchaforge.auto_label")
def parse_args(argv: list[str]) -> argparse.Namespace:
p = argparse.ArgumentParser(
prog="auto_label.py",
description="Auto-label a captchaforge TrainingCorpus with a teacher VLM.",
formatter_class=argparse.ArgumentDefaultsHelpFormatter,
)
p.add_argument(
"--corpus",
type=Path,
required=True,
help="Path to a captchaforge TrainingCorpus directory.",
)
p.add_argument(
"--vendors",
type=lambda s: [v.strip() for v in s.split(",") if v.strip()],
required=True,
help="Comma-separated list of vendor JSONL files to label "
"(e.g. 'hcaptcha-mock,recaptcha-v2-test'). Per-vendor opt-in "
"is mandatory — see TOS notes in the README.",
)
p.add_argument(
"--out",
type=Path,
required=True,
help="Output VlmDataset directory (manifest.jsonl + images/).",
)
p.add_argument(
"--teacher",
default="openai:gpt-4o",
help="`<provider>:<model>`. Supported: openai:* (env OPENAI_API_KEY), "
"anthropic:* (env ANTHROPIC_API_KEY).",
)
p.add_argument(
"--max-samples",
type=int,
default=None,
help="Cap per-vendor sample count.",
)
p.add_argument(
"--dry-run",
action="store_true",
help="Parse corpus + report counts; skip teacher calls.",
)
p.add_argument(
"--log-level",
default="INFO",
choices=["DEBUG", "INFO", "WARNING", "ERROR"],
)
return p.parse_args(argv)
def load_corpus_vendor(corpus_dir: Path, vendor: str) -> list[dict]:
path = corpus_dir / f"{vendor}.jsonl"
if not path.exists():
return []
out: list[dict] = []
with path.open("r", encoding="utf-8") as f:
for line_no, raw in enumerate(f, start=1):
raw = raw.strip()
if not raw:
continue
try:
out.append(json.loads(raw))
except json.JSONDecodeError as e:
LOG.warning("vendor %s line %d malformed: %s", vendor, line_no, e)
return out
def write_vlm_sample(
out_dir: Path, kind: str, image_bytes: bytes, prompt: str, answer: str
) -> str:
import hashlib
images_dir = out_dir / "images"
images_dir.mkdir(parents=True, exist_ok=True)
img_id = hashlib.sha256(image_bytes).hexdigest()
img_path = images_dir / f"{img_id}.png"
if not img_path.exists():
img_path.write_bytes(image_bytes)
sample = {
"id": img_id,
"kind": kind,
"image_path": f"images/{img_id}.png",
"prompt": prompt,
"answer": answer,
"is_negative": False,
}
manifest = out_dir / "manifest.jsonl"
with manifest.open("a", encoding="utf-8") as f:
f.write(json.dumps(sample))
f.write("\n")
return img_id
def call_teacher(teacher_spec: str, image_bytes: bytes, prompt: str) -> str:
provider, _, model = teacher_spec.partition(":")
if not model:
raise ValueError(f"--teacher must be 'provider:model' (got {teacher_spec!r})")
image_b64 = base64.b64encode(image_bytes).decode("ascii")
if provider == "openai":
from openai import OpenAI
client = OpenAI()
resp = client.chat.completions.create(
model=model,
messages=[
{
"role": "user",
"content": [
{"type": "text", "text": prompt},
{
"type": "image_url",
"image_url": {"url": f"data:image/png;base64,{image_b64}"},
},
],
}
],
max_tokens=256,
temperature=0.0,
)
return resp.choices[0].message.content or ""
if provider == "anthropic":
from anthropic import Anthropic
client = Anthropic()
resp = client.messages.create(
model=model,
max_tokens=256,
messages=[
{
"role": "user",
"content": [
{
"type": "image",
"source": {
"type": "base64",
"media_type": "image/png",
"data": image_b64,
},
},
{"type": "text", "text": prompt},
],
}
],
)
return "".join(
block.text for block in resp.content if hasattr(block, "text")
)
raise ValueError(f"unknown teacher provider: {provider!r}")
PROMPT_TEMPLATE = (
"You are labelling a captcha challenge for a training dataset. "
"Inspect the image and respond ONLY with a JSON object describing "
"what action would solve this captcha. For image-grid captchas "
"respond with `{{\"cells\": [<indices>]}}`. For text/transcription "
"captchas respond with `{{\"transcription\": \"<text>\"}}`. For "
"rotation captchas respond with `{{\"degrees\": <0-360>}}`. For "
"click captchas respond with `{{\"x\": <px>, \"y\": <px>}}`. "
"Vendor: {vendor}. Detected kind: {kind}."
)
def main(argv: list[str]) -> int:
args = parse_args(argv)
logging.basicConfig(
level=getattr(logging, args.log_level),
format="%(asctime)s %(levelname)s %(name)s %(message)s",
)
args.out.mkdir(parents=True, exist_ok=True)
written_total = 0
for vendor in args.vendors:
samples = load_corpus_vendor(args.corpus, vendor)
LOG.info("vendor %s: %d samples in corpus", vendor, len(samples))
if args.max_samples is not None:
samples = samples[: args.max_samples]
for i, sample in enumerate(samples):
screenshot_b64 = sample.get("screenshot_b64")
if not screenshot_b64:
LOG.debug("vendor %s sample %d: no screenshot, skipping", vendor, i)
continue
try:
image_bytes = base64.b64decode(screenshot_b64)
except Exception as e:
LOG.warning(
"vendor %s sample %d: bad base64 (%s), skipping",
vendor,
i,
e,
)
continue
prompt = PROMPT_TEMPLATE.format(
vendor=vendor, kind=sample.get("detected_kind", "unknown")
)
if args.dry_run:
LOG.info(
"[dry-run] would call teacher for vendor=%s sample=%d", vendor, i
)
continue
try:
answer = call_teacher(args.teacher, image_bytes, prompt)
except Exception as e:
LOG.warning(
"vendor %s sample %d: teacher call failed: %s",
vendor,
i,
e,
)
continue
write_vlm_sample(
args.out, kind=vendor, image_bytes=image_bytes, prompt=prompt, answer=answer
)
written_total += 1
if written_total % 10 == 0:
LOG.info("written %d total samples", written_total)
LOG.info("auto-label done: %d samples in %s", written_total, args.out)
return 0
if __name__ == "__main__":
sys.exit(main(sys.argv[1:]))