judge.py

"""
Qwen-Image-Bench Judge Model Inference Tool

Evaluate text-to-image generated images using a fine-tuned Qwen3.6-27B judge model.
Uses ms-swift PtEngine for batch inference.

Per-row output preserves all original input fields plus:
  - judge_model_output: combined raw scores JSON across all L1 dimensions
  - <dim>_judge_output: raw judge model text for each L1 dimension

Bench-level scores (L1 / L2 / Total) are aggregated following the
compute_scores.py methodology and saved alongside the per-row output.
"""

import argparse
import json
import sys
from collections import defaultdict
from pathlib import Path

import pandas as pd
from PIL import Image
from tqdm import tqdm

from checklists import (
    DIM_TO_CHECKLIST,
    SYSTEM_PROMPT,
    USER_PROMPT_TEMPLATE,
    parse_dims_by_level1,
)
from score_utils import (
    aggregate_total_score,
    compute_dimension_score,
    extract_json_from_response,
    fix_score_json,
)


DIM_OUTPUT_MAP = {
    "Quality": "quality_judge_output",
    "Aesthetics": "aesthetics_judge_output",
    "Alignment": "alignment_judge_output",
    "Real-world Fidelity": "real_world_fidelity_judge_output",
    "Creative Generation": "creative_generation_judge_output",
}


def load_and_resize_image(path):
    """Load image and resize to 1024x1024 if any dimension > 1024."""
    img = Image.open(path)
    if img.mode != "RGB":
        img = img.convert("RGB")
    if max(img.size) > 1024:
        img = img.resize((1024, 1024), Image.LANCZOS)
    img.load()
    return img


def load_input_file(file_path):
    """Load CSV or JSON/JSONL input file."""
    ext = Path(file_path).suffix.lower()
    if ext == ".csv":
        return pd.read_csv(file_path)
    elif ext == ".json":
        with open(file_path, "r", encoding="utf-8") as f:
            content = f.read().strip()
        if content.startswith("["):
            return pd.DataFrame(json.loads(content))
        else:
            records = [json.loads(line) for line in content.splitlines() if line.strip()]
            return pd.DataFrame(records)
    elif ext == ".jsonl":
        records = []
        with open(file_path, "r", encoding="utf-8") as f:
            for line in f:
                if line.strip():
                    records.append(json.loads(line))
        return pd.DataFrame(records)
    else:
        raise ValueError(f"Unsupported file format: {ext}. Use .csv, .json, or .jsonl")


def load_bench_metadata(hf_bench_repo=None, local_metadata=None):
    """Load bench metadata containing dims_en per ID."""
    if local_metadata:
        return load_input_file(local_metadata)

    if hf_bench_repo:
        from huggingface_hub import hf_hub_download

        local_file = hf_hub_download(
            repo_id=hf_bench_repo,
            filename="qwen_image_bench_hf_v0518.jsonl",
            repo_type="dataset",
        )
        records = []
        with open(local_file, "r", encoding="utf-8") as f:
            for line in f:
                if line.strip():
                    records.append(json.loads(line))
        return pd.DataFrame(records)

    default_path = Path(__file__).parent / "metadata" / "bench_metadata.json"
    if default_path.exists():
        return load_input_file(str(default_path))

    print("ERROR: No metadata source found. Provide --hf-bench-repo, --local-metadata,")
    print(f"       or place metadata at {default_path}")
    sys.exit(1)


def _parse_output_to_scores(output_text, level1_dim):
    """Parse raw judge model output → fixed score_json. Returns None on failure."""
    score_json = extract_json_from_response(output_text)
    if score_json is None:
        return None
    return fix_score_json(score_json, level1_dim)


def _run_batch_inference(judge, args, input_df, metadata_df, desc):
    """Batch inference over all rows and dimensions."""
    tasks = []
    task_meta = []  # (row_idx, level1_dim)
    skipped_rows = set()
    image_failures = 0

    for row_idx, (_, row) in enumerate(input_df.iterrows()):
        row_id = row["ID"]
        prompt = row["prompt"]
        image_path = row["image_path"]

        meta_row = metadata_df[metadata_df["ID"] == row_id]
        if meta_row.empty:
            continue

        dims_en = meta_row.iloc[0]["dims_en"]
        dims_by_level1 = parse_dims_by_level1(dims_en)

        try:
            img = load_and_resize_image(image_path)
        except Exception as e:
            image_failures += 1
            skipped_rows.add(row_idx)
            print(f"WARNING: Failed to load image for ID={row_id}, path={image_path}: {e}")
            continue

        for level1_dim in dims_by_level1:
            if level1_dim not in DIM_TO_CHECKLIST:
                continue
            checklist = DIM_TO_CHECKLIST[level1_dim]
            user_text = USER_PROMPT_TEMPLATE.format(
                prompt=prompt,
                level1_dim=level1_dim,
                format_checklist=checklist,
            )
            tasks.append({
                "system_prompt": SYSTEM_PROMPT,
                "user_text": user_text,
                "image": img,
            })
            task_meta.append((row_idx, level1_dim))

    print(f"Total inference tasks: {len(tasks)}")
    all_outputs = []
    batch_size = args.batch_size

    for i in tqdm(range(0, len(tasks), batch_size), desc=desc):
        batch = tasks[i:i + batch_size]
        outputs = judge.generate_batch(batch)
        all_outputs.extend(outputs)

    row_dim_raw_scores = {}
    row_raw_outputs = {}
    parse_failures = 0

    for (row_idx, level1_dim), output_text in zip(task_meta, all_outputs):
        if row_idx not in row_dim_raw_scores:
            row_dim_raw_scores[row_idx] = {}
            row_raw_outputs[row_idx] = {}
        row_raw_outputs[row_idx][level1_dim] = output_text
        score_json = _parse_output_to_scores(output_text, level1_dim)
        if score_json is None:
            parse_failures += 1
        row_dim_raw_scores[row_idx][level1_dim] = score_json

    results = []
    all_dim_raw_scores = []
    for row_idx, (_, row) in enumerate(input_df.iterrows()):
        if row_idx in skipped_rows:
            results.append(_empty_result(row))
            all_dim_raw_scores.append({})
            continue
        dim_raw_scores = row_dim_raw_scores.get(row_idx, {})
        dim_raw_outputs = row_raw_outputs.get(row_idx, {})
        results.append(_build_row_result(row, dim_raw_scores, dim_raw_outputs))
        all_dim_raw_scores.append(dim_raw_scores)

    return results, parse_failures, all_dim_raw_scores, image_failures


def run_ms_swift_inference(args, input_df, metadata_df):
    """Run inference using ms-swift PtEngine."""
    from backends.ms_swift_backend import MsSwiftJudge

    print(f"Loading model from: {args.model}")
    judge = MsSwiftJudge(
        model_path=args.model,
        max_batch_size=args.max_batch_size,
        max_new_tokens=args.max_new_tokens,
    )
    print("Model loaded successfully.")
    return _run_batch_inference(judge, args, input_df, metadata_df, desc="Batch inference")


def _empty_result(row):
    """Build an empty result row for skipped entries."""
    result = dict(row)
    result["judge_model_output"] = None
    for col in DIM_OUTPUT_MAP.values():
        result[col] = None
    return result


def _build_row_result(row, dim_raw_scores, dim_raw_outputs):
    """
    Build per-row JSONL record.

    Schema (in order):
      - all original row fields (transparent pass-through)
      - judge_model_output: JSON-serialized {L1_dim: fixed_score_json} for all parsed dims
      - <dim>_judge_output: raw judge text for each L1 dim
    """
    result = dict(row)

    raw_output = {
        dim_name: score_json
        for dim_name, score_json in dim_raw_scores.items()
        if score_json is not None
    }
    result["judge_model_output"] = (
        json.dumps(raw_output, ensure_ascii=False) if raw_output else None
    )

    for dim_name, col_name in DIM_OUTPUT_MAP.items():
        result[col_name] = dim_raw_outputs.get(dim_name)

    return result


def save_output(results, input_path):
    """Save per-row results to file in same directory as input."""
    input_p = Path(input_path)
    ext = input_p.suffix.lower()
    output_name = f"{input_p.stem}_judged{ext}"
    output_path = input_p.parent / output_name

    df = pd.DataFrame(results)

    if ext == ".csv":
        df.to_csv(output_path, index=False, encoding="utf-8")
    else:
        with open(output_path, "w", encoding="utf-8") as f:
            for _, row in df.iterrows():
                f.write(json.dumps(row.to_dict(), ensure_ascii=False) + "\n")

    return str(output_path)


def _safe_mean(xs):
    return sum(xs) / len(xs) if xs else None


def compute_bench_scores(all_dim_raw_scores):
    """
    Bench-level aggregation following compute_scores.py methodology:
      per-row L3→L2→L1→Total nested averaging,
      then arithmetic mean across rows (None values skipped).
    """
    l1_accum = defaultdict(list)
    l2_accum = defaultdict(lambda: defaultdict(list))
    total_accum = []

    for row_scores in all_dim_raw_scores:
        dim_results = {}
        for l1_dim, score_json in row_scores.items():
            if score_json is None:
                continue
            dim_results[l1_dim] = compute_dimension_score(score_json)

        row_total = aggregate_total_score(dim_results)
        if row_total is not None:
            total_accum.append(row_total)

        for l1_dim, dim_data in dim_results.items():
            if dim_data["level1_score"] is not None:
                l1_accum[l1_dim].append(dim_data["level1_score"])
            for l2_name, l2_score in dim_data["level2_scores"].items():
                if l2_score is not None:
                    l2_accum[l1_dim][l2_name].append(l2_score)

    return {
        "level1": {d: _safe_mean(v) for d, v in l1_accum.items()},
        "level2": {
            d: {l2: _safe_mean(v) for l2, v in l2d.items()}
            for d, l2d in l2_accum.items()
        },
        "total": _safe_mean(total_accum),
    }


def save_bench_scores(bench, input_path):
    """Save bench-level scores as JSON + Excel beside the input file."""
    input_p = Path(input_path)
    base = input_p.parent / f"{input_p.stem}_bench_scores"
    json_path = f"{base}.json"
    xlsx_path = f"{base}.xlsx"

    with open(json_path, "w", encoding="utf-8") as f:
        json.dump(bench, f, ensure_ascii=False, indent=2)

    with pd.ExcelWriter(xlsx_path, engine="openpyxl") as writer:
        l1_rows = [{"Dimension": d, "Score": s} for d, s in bench["level1"].items()]
        l1_rows.append({"Dimension": "Total", "Score": bench["total"]})
        pd.DataFrame(l1_rows).to_excel(writer, sheet_name="Level-1 Summary", index=False)

        for dim, l2_dict in bench["level2"].items():
            if not l2_dict:
                continue
            df = pd.DataFrame(
                [{"Sub-dimension": l2, "Score": s} for l2, s in l2_dict.items()]
            )
            df.to_excel(writer, sheet_name=dim[:31], index=False)

    return json_path, xlsx_path


def print_bench_scores(bench):
    """Pretty-print bench-level scores to terminal."""
    print("\n" + "=" * 70)
    print("BENCH-LEVEL SCORES")
    print("=" * 70)
    for dim, score in bench["level1"].items():
        s = f"{score:.2f}" if score is not None else "N/A"
        print(f"  L1 {dim:30s}: {s}")
    total = bench["total"]
    total_str = f"{total:.2f}" if total is not None else "N/A"
    print(f"  {'TOTAL':33s}: {total_str}")
    print("-" * 70)
    for dim, l2_dict in bench["level2"].items():
        if not l2_dict:
            continue
        print(f"  [{dim}]")
        for l2, s in l2_dict.items():
            v = f"{s:.2f}" if s is not None else "N/A"
            print(f"    L2 {l2:28s}: {v}")
    print("=" * 70)


def main():
    parser = argparse.ArgumentParser(
        description="Qwen-Image-Bench Judge Model Inference Tool"
    )
    parser.add_argument("--input", required=True, help="Input CSV/JSON/JSONL with ID, prompt, image_path")
    parser.add_argument("--model", required=True, help="HuggingFace model ID or local model path")
    parser.add_argument("--hf-bench-repo", default=None, help="HF dataset repo for bench metadata")
    parser.add_argument("--local-metadata", default=None, help="Local metadata file path (skip HF download)")
    parser.add_argument("--max-batch-size", type=int, default=24,
                        help="ms-swift PtEngine max_batch_size (default: 24)")
    parser.add_argument("--max-new-tokens", type=int, default=4096)

    args = parser.parse_args()
    args.batch_size = args.max_batch_size

    # Load input
    print(f"Loading input: {args.input}")
    input_df = load_input_file(args.input)
    required_cols = {"ID", "prompt", "image_path"}
    missing = required_cols - set(input_df.columns)
    if missing:
        print(f"ERROR: Input file missing required columns: {missing}")
        sys.exit(1)
    print(f"Input: {len(input_df)} rows")

    # Load metadata
    print("Loading bench metadata...")
    metadata_df = load_bench_metadata(
        hf_bench_repo=args.hf_bench_repo,
        local_metadata=args.local_metadata,
    )
    print(f"Metadata: {len(metadata_df)} rows")

    # Run inference
    results, parse_failures, all_dim_raw_scores, image_failures = run_ms_swift_inference(args, input_df, metadata_df)

    # Save per-row JSONL
    saved_path = save_output(results, args.input)
    print(f"\nPer-row results saved to: {saved_path}")
    if image_failures:
        print(f"Skipped (broken images): {image_failures}")
    print(f"Parse failures: {parse_failures}")

    # Compute & save bench-level scores
    bench = compute_bench_scores(all_dim_raw_scores)
    json_path, xlsx_path = save_bench_scores(bench, args.input)
    print(f"Bench scores saved to: {json_path}")
    print(f"Bench scores saved to: {xlsx_path}")
    print_bench_scores(bench)


if __name__ == "__main__":
    main()