app.py

import csv
import itertools
import random
import json
import os
import uuid
from datetime import datetime
from io import BytesIO
from typing import Dict, List, Tuple

import gradio as gr
try:
    from huggingface_hub import HfApi
except Exception:  # optional dependency at runtime
    HfApi = None  # type: ignore


BASE_DIR = os.path.dirname(__file__)
PERSIST_DIR = os.environ.get("PERSIST_DIR", "/data")
# Persistent local storage inside HF Spaces
PERSIST_DIR = os.environ.get("PERSIST_DIR", "/data")
TASK_CONFIG = {
    "Scene Composition & Object Insertion": {
        "folder": "scene_composition_and_object_insertion",
        "score_fields": [
            ("physical_interaction_fidelity_score", "物理交互保真度 (Physical Interaction Fidelity)"),
            ("optical_effect_accuracy_score", "光学效应准确度 (Optical Effect Accuracy)"),
            ("semantic_functional_alignment_score", "语义/功能对齐度 (Semantic/Functional Alignment)"),
            ("overall_photorealism_score", "整体真实感 (Overall Photorealism)"),
        ],
    },
}


def _csv_path_for_task(task_name: str, filename: str) -> str:
    folder = TASK_CONFIG[task_name]["folder"]
    return os.path.join(BASE_DIR, folder, filename)


def _persist_csv_path_for_task(task_name: str) -> str:
    folder = TASK_CONFIG[task_name]["folder"]
    return os.path.join(PERSIST_DIR, folder, "evaluation_results.csv")


def _resolve_image_path(path: str) -> str:
    return path if os.path.isabs(path) else os.path.join(BASE_DIR, path)


def _load_task_rows(task_name: str) -> List[Dict[str, str]]:
    csv_path = _csv_path_for_task(task_name, "results.csv")
    if not os.path.exists(csv_path):
        raise FileNotFoundError(f"未找到任务 {task_name} 的结果文件: {csv_path}")

    with open(csv_path, newline="", encoding="utf-8") as csv_file:
        reader = csv.DictReader(csv_file)
        return [row for row in reader]


def _build_image_pairs(rows: List[Dict[str, str]], task_name: str) -> List[Dict[str, str]]:
    grouped: Dict[Tuple[str, str], List[Dict[str, str]]] = {}
    for row in rows:
        key = (row["test_id"], row["org_img"])
        grouped.setdefault(key, []).append(row)

    pairs: List[Dict[str, str]] = []
    folder = TASK_CONFIG[task_name]["folder"]

    for (test_id, org_img), entries in grouped.items():
        for model_a, model_b in itertools.combinations(entries, 2):
            if model_a["model_name"] == model_b["model_name"]:
                continue

            pair = {
                "test_id": test_id,
                "org_img": os.path.join(folder, org_img),
                "model1_name": model_a["model_name"],
                "model1_res": model_a["res"],
                "model1_path": os.path.join(folder, model_a["path"]),
                "model2_name": model_b["model_name"],
                "model2_res": model_b["res"],
                "model2_path": os.path.join(folder, model_b["path"]),
            }
            pairs.append(pair)

    def sort_key(item: Dict[str, str]):
        test_id = item["test_id"]
        try:
            test_id_key = int(test_id)
        except ValueError:
            test_id_key = test_id
        return (test_id_key, item["model1_name"], item["model2_name"])

    pairs.sort(key=sort_key)
    return pairs


def _read_existing_eval_keys(task_name: str) -> set:
    """Read already-evaluated pair keys from persistent CSV, return a set of keys.
    Key is (test_id, frozenset({model1_name, model2_name}), org_img) to ignore A/B order.
    """
    keys = set()
    csv_path = _persist_csv_path_for_task(task_name)
    if not os.path.exists(csv_path):
        return keys
    try:
        with open(csv_path, newline="", encoding="utf-8") as f:
            reader = csv.DictReader(f)
            for r in reader:
                tid = str(r.get("test_id", "")).strip()
                m1 = str(r.get("model1_name", "")).strip()
                m2 = str(r.get("model2_name", "")).strip()
                org = str(r.get("org_img", "")).strip()
                if tid and m1 and m2 and org:
                    keys.add((tid, frozenset({m1, m2}), org))
    except Exception:
        pass
    return keys


def _schedule_round_robin_by_test_id(pairs: List[Dict[str, str]], seed: int | None = None) -> List[Dict[str, str]]:
    """Interleave pairs across test_ids for balanced coverage; shuffle within each group.
    """
    groups: Dict[str, List[Dict[str, str]]] = {}
    for p in pairs:
        groups.setdefault(p["test_id"], []).append(p)
    rnd = random.Random(seed)
    for lst in groups.values():
        rnd.shuffle(lst)
    # round-robin drain
    ordered: List[Dict[str, str]] = []
    while True:
        progressed = False
        for tid in sorted(groups.keys(), key=lambda x: (int(x) if x.isdigit() else x)):
            if groups[tid]:
                ordered.append(groups[tid].pop())
                progressed = True
        if not progressed:
            break
    return ordered


def load_task(task_name: str):
    if not task_name:
        raise gr.Error("请先选择任务。")

    rows = _load_task_rows(task_name)
    pairs = _build_image_pairs(rows, task_name)
    # Filter out already evaluated pairs from persistent CSV
    done_keys = _read_existing_eval_keys(task_name)
    def key_of(p: Dict[str, str]):
        return (p["test_id"], frozenset({p["model1_name"], p["model2_name"]}), p["org_img"])
    pairs = [p for p in pairs if key_of(p) not in done_keys]

    # Balanced schedule across test_ids with a stable randomization
    seed_env = os.environ.get("SCHEDULE_SEED")
    seed = int(seed_env) if seed_env and seed_env.isdigit() else None
    pairs = _schedule_round_robin_by_test_id(pairs, seed=seed)

    # Assign A/B order to counteract position bias: alternate after scheduling
    for idx, p in enumerate(pairs):
        p["swap"] = bool(idx % 2)  # True -> A=B's image; False -> A=A's image

    if not pairs:
        raise gr.Error("没有找到可评测的图片对，请检查数据文件。")

    return pairs


def _format_pair_header(_pair: Dict[str, str]) -> str:
    # Mask model identity in UI; keep header neutral
    return ""


def _build_eval_row(pair: Dict[str, str], scores: Dict[str, int]) -> Dict[str, object]:
    row = {
        "eval_date": datetime.utcnow().isoformat(),
        "test_id": pair["test_id"],
        "model1_name": pair["model1_name"],
        "model2_name": pair["model2_name"],
        "org_img": pair["org_img"],
        "model1_res": pair["model1_res"],
        "model2_res": pair["model2_res"],
        "model1_path": pair["model1_path"],
        "model2_path": pair["model2_path"],
    }
    row.update(scores)
    return row


def _local_persist_csv_path(task_name: str) -> str:
    folder = TASK_CONFIG[task_name]["folder"]
    return os.path.join(PERSIST_DIR, folder, "evaluation_results.csv")


def _append_local_persist_csv(task_name: str, row: Dict[str, object]) -> bool:
    csv_path = _local_persist_csv_path(task_name)
    os.makedirs(os.path.dirname(csv_path), exist_ok=True)
    csv_exists = os.path.exists(csv_path)
    fieldnames = [
        "eval_date",
        "test_id",
        "model1_name",
        "model2_name",
        "org_img",
        "model1_res",
        "model2_res",
        "model1_path",
        "model2_path",
        "model1_physical_interaction_fidelity_score",
        "model1_optical_effect_accuracy_score",
        "model1_semantic_functional_alignment_score",
        "model1_overall_photorealism_score",
        "model2_physical_interaction_fidelity_score",
        "model2_optical_effect_accuracy_score",
        "model2_semantic_functional_alignment_score",
        "model2_overall_photorealism_score",
    ]
    try:
        with open(csv_path, "a", newline="", encoding="utf-8") as csv_file:
            writer = csv.DictWriter(csv_file, fieldnames=fieldnames)
            if not csv_exists:
                writer.writeheader()
            writer.writerow(row)
        return True
    except Exception:
        return False


def _upload_eval_record_to_dataset(task_name: str, row: Dict[str, object]) -> tuple[bool, str]:
    """Upload a single-eval JSONL record to a dataset repo.
    Repo is taken from EVAL_REPO_ID env or defaults to 'peiranli0930/VisEval'.
    Returns (ok, message) for UI feedback and debugging.
    """
    if HfApi is None:
        return False, "huggingface_hub 未安装"
    token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGINGFACEHUB_API_TOKEN")
    repo_id = os.environ.get("EVAL_REPO_ID", "peiranli0930/VisEval")
    if not token:
        return False, "未找到写入 Token (HF_TOKEN/HUGGINGFACEHUB_API_TOKEN)"
    if not repo_id:
        return False, "未设置 EVAL_REPO_ID"
    try:
        from huggingface_hub import CommitOperationAdd

        api = HfApi(token=token)
        date_prefix = datetime.utcnow().strftime("%Y-%m-%d")
        folder = TASK_CONFIG[task_name]["folder"]
        uid = str(uuid.uuid4())
        path_in_repo = f"submissions/{folder}/{date_prefix}/{uid}.jsonl"
        payload = (json.dumps(row, ensure_ascii=False) + "\n").encode("utf-8")
        operations = [CommitOperationAdd(path_in_repo=path_in_repo, path_or_fileobj=BytesIO(payload))]
        api.create_commit(
            repo_id=repo_id,
            repo_type="dataset",
            operations=operations,
            commit_message=f"Add eval {folder} {row.get('test_id')} {uid}",
        )
        return True, f"上传成功: {repo_id}/{path_in_repo}"
    except Exception as e:
        # Print to logs for debugging in Space
        try:
            print("[VisArena] Upload to dataset failed:", repr(e))
        except Exception:
            pass
        return False, f"异常: {type(e).__name__}: {e}"


def on_task_change(task_name: str, _state_pairs: List[Dict[str, str]]):
    pairs = load_task(task_name)
    pair = pairs[0]
    header = _format_pair_header(pair)
    # Defaults for A and B (8 sliders total)
    default_scores = [3, 3, 3, 3, 3, 3, 3, 3]
    # Pick display order according to swap flag
    a_path = pair["model2_path"] if pair.get("swap") else pair["model1_path"]
    b_path = pair["model1_path"] if pair.get("swap") else pair["model2_path"]
    return (
        pairs,
        gr.update(value=0, minimum=0, maximum=len(pairs) - 1, visible=(len(pairs) > 1)),
        gr.update(value=header),
        _resolve_image_path(pair["org_img"]),
        _resolve_image_path(a_path),
        _resolve_image_path(b_path),
        *default_scores,
        gr.update(value=f"共 {len(pairs)} 个待评测的图片对。"),
    )


def on_pair_navigate(index: int, pairs: List[Dict[str, str]]):
    if not pairs:
        raise gr.Error("请先选择任务。")
    index = int(index)
    index = max(0, min(index, len(pairs) - 1))
    pair = pairs[index]
    header = _format_pair_header(pair)
    a_path = pair["model2_path"] if pair.get("swap") else pair["model1_path"]
    b_path = pair["model1_path"] if pair.get("swap") else pair["model2_path"]
    return (
        gr.update(value=index),
        gr.update(value=header),
        _resolve_image_path(pair["org_img"]),
        _resolve_image_path(a_path),
        _resolve_image_path(b_path),
        3, 3, 3, 3,  # A
        3, 3, 3, 3,  # B
    )


def on_submit(
    task_name: str,
    index: int,
    pairs: List[Dict[str, str]],
    a_physical_score: int,
    a_optical_score: int,
    a_semantic_score: int,
    a_overall_score: int,
    b_physical_score: int,
    b_optical_score: int,
    b_semantic_score: int,
    b_overall_score: int,
):
    if not task_name:
        raise gr.Error("请先选择任务。")

    if not pairs:
        raise gr.Error("当前任务没有加载任何图片对。")

    pair = pairs[index]
    score_map = {
        # Model A
        "model1_physical_interaction_fidelity_score": int(a_physical_score),
        "model1_optical_effect_accuracy_score": int(a_optical_score),
        "model1_semantic_functional_alignment_score": int(a_semantic_score),
        "model1_overall_photorealism_score": int(a_overall_score),
        # Model B
        "model2_physical_interaction_fidelity_score": int(b_physical_score),
        "model2_optical_effect_accuracy_score": int(b_optical_score),
        "model2_semantic_functional_alignment_score": int(b_semantic_score),
        "model2_overall_photorealism_score": int(b_overall_score),
    }
    # Map A/B scores to the correct model columns depending on swap
    if pair.get("swap"):
        # UI A == model2, UI B == model1
        score_map = {
            "model1_physical_interaction_fidelity_score": int(b_physical_score),
            "model1_optical_effect_accuracy_score": int(b_optical_score),
            "model1_semantic_functional_alignment_score": int(b_semantic_score),
            "model1_overall_photorealism_score": int(b_overall_score),
            "model2_physical_interaction_fidelity_score": int(a_physical_score),
            "model2_optical_effect_accuracy_score": int(a_optical_score),
            "model2_semantic_functional_alignment_score": int(a_semantic_score),
            "model2_overall_photorealism_score": int(a_overall_score),
        }
    row = _build_eval_row(pair, score_map)
    ok_local = _append_local_persist_csv(task_name, row)
    ok_hub, hub_msg = _upload_eval_record_to_dataset(task_name, row)

    next_index = min(index + 1, len(pairs) - 1)
    info = f"已保存 Test ID {pair['test_id']} 的评价结果。"
    info += " 本地持久化" + ("成功" if ok_local else "失败") + "。"
    info += " 上传Hub" + ("成功" if ok_hub else "失败") + (f"（{hub_msg}）" if hub_msg else "") + "。"

    if next_index != index:
        pair = pairs[next_index]
        header = _format_pair_header(pair)
        a_path = pair["model2_path"] if pair.get("swap") else pair["model1_path"]
        b_path = pair["model1_path"] if pair.get("swap") else pair["model2_path"]
        return (
            gr.update(value=next_index),
            gr.update(value=header),
            _resolve_image_path(pair["org_img"]),
            _resolve_image_path(a_path),
            _resolve_image_path(b_path),
            3, 3, 3, 3,
            3, 3, 3, 3,
            gr.update(value=info + f" 自动跳转到下一组（{next_index + 1}/{len(pairs)}）。"),
        )

    return (
        gr.update(),
        gr.update(),
        gr.update(),
        gr.update(),
        gr.update(),
        3, 3, 3, 3,
        3, 3, 3, 3,
        gr.update(value=info + " 已经是最后一组。"),
    )


with gr.Blocks(title="VisArena Human Evaluation") as demo:
    gr.Markdown(
        """
        # VisArena Human Evaluation
        请选择任务并对模型生成的图像进行评分。每项评分范围为 **1（效果极差）** 到 **5（效果极佳）**。
        """
    )

    with gr.Row():
        task_selector = gr.Dropdown(
            label="Task",
            choices=list(TASK_CONFIG.keys()),
            interactive=True,
            value="Scene Composition & Object Insertion",
        )
        index_slider = gr.Slider(
            label="Pair Index",
            value=0,
            minimum=0,
            maximum=0,
            step=1,
            interactive=True,
            visible=False,
        )

    pair_state = gr.State([])

    pair_header = gr.Markdown("")

    # Layout: Original on top, two outputs below with their own sliders
    with gr.Row():
        with gr.Column(scale=12):
            orig_image = gr.Image(type="filepath", label="原图 Original", interactive=False)

    with gr.Row():
        with gr.Column(scale=6):
            model1_image = gr.Image(type="filepath", label="模型 A 输出", interactive=False)
            a_physical_input = gr.Slider(1, 5, value=3, step=1, label="A: 物理交互保真度")
            a_optical_input = gr.Slider(1, 5, value=3, step=1, label="A: 光学效应准确度")
            a_semantic_input = gr.Slider(1, 5, value=3, step=1, label="A: 语义/功能对齐度")
            a_overall_input = gr.Slider(1, 5, value=3, step=1, label="A: 整体真实感")
        with gr.Column(scale=6):
            model2_image = gr.Image(type="filepath", label="模型 B 输出", interactive=False)
            b_physical_input = gr.Slider(1, 5, value=3, step=1, label="B: 物理交互保真度")
            b_optical_input = gr.Slider(1, 5, value=3, step=1, label="B: 光学效应准确度")
            b_semantic_input = gr.Slider(1, 5, value=3, step=1, label="B: 语义/功能对齐度")
            b_overall_input = gr.Slider(1, 5, value=3, step=1, label="B: 整体真实感")

    submit_button = gr.Button("Submit Evaluation", variant="primary")
    feedback_box = gr.Markdown("")

    # Event bindings
    task_selector.change(
        fn=on_task_change,
        inputs=[task_selector, pair_state],
        outputs=[
            pair_state,
            index_slider,
            pair_header,
            orig_image,
            model1_image,
            model2_image,
            a_physical_input,
            a_optical_input,
            a_semantic_input,
            a_overall_input,
            b_physical_input,
            b_optical_input,
            b_semantic_input,
            b_overall_input,
            feedback_box,
        ],
    )

    index_slider.release(
        fn=on_pair_navigate,
        inputs=[index_slider, pair_state],
        outputs=[
            index_slider,
            pair_header,
            orig_image,
            model1_image,
            model2_image,
            a_physical_input,
            a_optical_input,
            a_semantic_input,
            a_overall_input,
            b_physical_input,
            b_optical_input,
            b_semantic_input,
            b_overall_input,
        ],
    )

    submit_button.click(
        fn=on_submit,
        inputs=[
            task_selector,
            index_slider,
            pair_state,
            a_physical_input,
            a_optical_input,
            a_semantic_input,
            a_overall_input,
            b_physical_input,
            b_optical_input,
            b_semantic_input,
            b_overall_input,
        ],
        outputs=[
            index_slider,
            pair_header,
            orig_image,
            model1_image,
            model2_image,
            a_physical_input,
            a_optical_input,
            a_semantic_input,
            a_overall_input,
            b_physical_input,
            b_optical_input,
            b_semantic_input,
            b_overall_input,
            feedback_box,
        ],
    )

    # Auto-load default task on startup
    demo.load(
        fn=on_task_change,
        inputs=[task_selector, pair_state],
        outputs=[
            pair_state,
            index_slider,
            pair_header,
            orig_image,
            model1_image,
            model2_image,
            a_physical_input,
            a_optical_input,
            a_semantic_input,
            a_overall_input,
            b_physical_input,
            b_optical_input,
            b_semantic_input,
            b_overall_input,
            feedback_box,
        ],
    )


if __name__ == "__main__":
    demo.queue().launch()