WeiQing Chen
c392ce21e9
Signed-off-by: David Chen <530634352@qq.com> Signed-off-by: WeiQing Chen <40507679+david6666666@users.noreply.github.com>
928 lines
37 KiB
Python
928 lines
37 KiB
Python
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import statistics
|
|
from collections import defaultdict
|
|
from concurrent.futures import ThreadPoolExecutor, as_completed
|
|
from dataclasses import dataclass
|
|
from datetime import datetime, timezone
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import requests
|
|
from PIL import Image
|
|
|
|
from benchmarks.accuracy.common import (
|
|
VllmOmniImageClient,
|
|
build_openai_url,
|
|
ensure_dir,
|
|
extract_json_object,
|
|
find_first_image,
|
|
load_json,
|
|
pil_to_data_url,
|
|
save_image,
|
|
write_json,
|
|
)
|
|
|
|
TYPE_TO_FOLDER = {
|
|
"type1": "01_single_step",
|
|
"type2": "02_multi_step",
|
|
"type3": "03_trajectory_text_fictionalapp",
|
|
"type4": "04_trajectory_text_realapp",
|
|
"type5": "05_grounding_data",
|
|
}
|
|
SCORE_KEYS = ("goal", "logic", "cons", "ui", "qual")
|
|
DEFAULT_SAMPLES_PER_TYPE = 10
|
|
|
|
|
|
def _utc_timestamp() -> str:
|
|
return datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")
|
|
|
|
|
|
def _write_json_with_timestamp(path: Path, payload: dict[str, Any]) -> Path:
|
|
write_json(path, payload)
|
|
timestamped_path = path.with_name(f"{path.stem}_{_utc_timestamp()}{path.suffix}")
|
|
write_json(timestamped_path, payload)
|
|
return timestamped_path
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class GEBenchSampleSpec:
|
|
sample_path: Path
|
|
metadata: dict[str, Any]
|
|
sample_name: str
|
|
lang_device: str
|
|
|
|
|
|
def summarize_generated_records(records: list[dict[str, Any]]) -> dict[str, Any]:
|
|
by_type: dict[str, list[dict[str, Any]]] = defaultdict(list)
|
|
for record in records:
|
|
by_type[record["data_type"]].append(record)
|
|
|
|
return {
|
|
"count": len(records),
|
|
"by_type": {
|
|
data_type: {
|
|
"count": len(rows),
|
|
}
|
|
for data_type, rows in sorted(by_type.items())
|
|
},
|
|
}
|
|
|
|
|
|
def summarize_gebench_results(results: list[dict[str, Any]]) -> dict[str, Any]:
|
|
by_type: dict[str, list[dict[str, Any]]] = defaultdict(list)
|
|
for result in results:
|
|
by_type[result["data_type"]].append(result)
|
|
|
|
summary: dict[str, Any] = {
|
|
"count": len(results),
|
|
"overall_mean": statistics.fmean(r["overall"] for r in results) if results else 0.0,
|
|
"by_type": {},
|
|
}
|
|
for data_type, rows in by_type.items():
|
|
score_means: dict[str, float] = {}
|
|
all_score_keys = {key for row in rows for key in row.get("scores", {}).keys()}
|
|
for score_key in all_score_keys:
|
|
values = [row["scores"][score_key] for row in rows if score_key in row.get("scores", {})]
|
|
score_means[score_key] = statistics.fmean(values) if values else 0.0
|
|
overall_mean = statistics.fmean(row["overall"] for row in rows)
|
|
summary["by_type"][data_type] = {
|
|
"count": len(rows),
|
|
"overall_mean": overall_mean,
|
|
"overall_mean_100": overall_mean * 100.0,
|
|
"score_means": score_means,
|
|
}
|
|
return summary
|
|
|
|
|
|
def select_balanced_gebench_samples(
|
|
sample_paths_by_type: dict[str, list[Any]],
|
|
*,
|
|
samples_per_type: int | None,
|
|
) -> dict[str, list[Any]]:
|
|
if samples_per_type is None:
|
|
return {data_type: list(paths) for data_type, paths in sample_paths_by_type.items()}
|
|
return {data_type: list(paths)[:samples_per_type] for data_type, paths in sample_paths_by_type.items()}
|
|
|
|
|
|
def collect_gebench_generation_summary(output_root: Path) -> dict[str, Any]:
|
|
records: list[dict[str, Any]] = []
|
|
for data_type, folder_name in TYPE_TO_FOLDER.items():
|
|
type_root = output_root / folder_name
|
|
if not type_root.exists():
|
|
continue
|
|
for lang_dir in sorted(path for path in type_root.iterdir() if path.is_dir()):
|
|
for sample_dir in sorted(path for path in lang_dir.iterdir() if path.is_dir()):
|
|
expected = sample_dir / "frame5.png" if data_type in {"type2", "type3", "type4"} else None
|
|
if expected is None:
|
|
expected = find_first_image(sample_dir)
|
|
elif not expected.exists():
|
|
expected = None
|
|
if expected is None:
|
|
continue
|
|
records.append(
|
|
{
|
|
"data_type": data_type,
|
|
"sample_name": f"{lang_dir.name}/{sample_dir.name}",
|
|
"output_path": str(expected),
|
|
}
|
|
)
|
|
return summarize_generated_records(records)
|
|
|
|
|
|
def _normalize_score_key(key: str) -> str:
|
|
mapping = {
|
|
"goal": "goal",
|
|
"logic": "logic",
|
|
"cons": "cons",
|
|
"consistency": "cons",
|
|
"ui": "ui",
|
|
"qual": "qual",
|
|
"quality": "qual",
|
|
}
|
|
return mapping.get(key.lower(), key.lower())
|
|
|
|
|
|
def _normalize_scores(raw_scores: dict[str, Any]) -> dict[str, int]:
|
|
scores: dict[str, int] = {}
|
|
for key, value in raw_scores.items():
|
|
normalized = _normalize_score_key(key)
|
|
if normalized not in SCORE_KEYS:
|
|
continue
|
|
scalar = value.get("s", 0) if isinstance(value, dict) else value
|
|
try:
|
|
scores[normalized] = int(scalar)
|
|
except (TypeError, ValueError):
|
|
scores[normalized] = 0
|
|
for key in SCORE_KEYS:
|
|
scores.setdefault(key, 0)
|
|
return scores
|
|
|
|
|
|
def _compute_overall(scores: dict[str, int]) -> float:
|
|
return sum(scores.values()) / (len(SCORE_KEYS) * 5.0)
|
|
|
|
|
|
def _iter_sample_paths(dataset_root: Path, data_type: str) -> list[Path]:
|
|
data_dir = dataset_root / TYPE_TO_FOLDER[data_type]
|
|
if not data_dir.exists():
|
|
raise FileNotFoundError(f"Dataset path not found: {data_dir}")
|
|
|
|
samples: list[Path] = []
|
|
for lang_dir in sorted(path for path in data_dir.iterdir() if path.is_dir()):
|
|
for child in sorted(lang_dir.iterdir()):
|
|
if child.is_dir():
|
|
samples.append(child)
|
|
elif child.suffix.lower() == ".json":
|
|
samples.append(child)
|
|
return samples
|
|
|
|
|
|
def _sample_name_from_metadata(metadata: dict[str, Any], sample_path: Path, item_index: int | None = None) -> str:
|
|
sample_id = metadata.get("id") or metadata.get("sample_id") or metadata.get("name")
|
|
if sample_id:
|
|
return str(sample_id)
|
|
if item_index is not None:
|
|
return f"{sample_path.stem}_{item_index:04d}"
|
|
return sample_path.stem if sample_path.is_file() else sample_path.name
|
|
|
|
|
|
def _expand_sample_path(sample_path: Path) -> list[GEBenchSampleSpec]:
|
|
if sample_path.is_dir():
|
|
metadata = _load_metadata(sample_path)
|
|
return [
|
|
GEBenchSampleSpec(
|
|
sample_path=sample_path,
|
|
metadata=metadata,
|
|
sample_name=_sample_name(sample_path),
|
|
lang_device=_lang_device(sample_path, metadata),
|
|
)
|
|
]
|
|
|
|
payload = load_json(sample_path)
|
|
if isinstance(payload, dict):
|
|
return [
|
|
GEBenchSampleSpec(
|
|
sample_path=sample_path,
|
|
metadata=payload,
|
|
sample_name=_sample_name_from_metadata(payload, sample_path),
|
|
lang_device=_lang_device(sample_path, payload),
|
|
)
|
|
]
|
|
|
|
if isinstance(payload, list):
|
|
specs: list[GEBenchSampleSpec] = []
|
|
for item_index, item in enumerate(payload):
|
|
if not isinstance(item, dict):
|
|
continue
|
|
specs.append(
|
|
GEBenchSampleSpec(
|
|
sample_path=sample_path,
|
|
metadata=item,
|
|
sample_name=_sample_name_from_metadata(item, sample_path, item_index),
|
|
lang_device=_lang_device(sample_path, item),
|
|
)
|
|
)
|
|
return specs
|
|
|
|
raise TypeError(f"Unsupported metadata payload type for sample {sample_path}: {type(payload)!r}")
|
|
|
|
|
|
def _iter_sample_specs(dataset_root: Path, data_type: str) -> list[GEBenchSampleSpec]:
|
|
sample_specs: list[GEBenchSampleSpec] = []
|
|
for sample_path in _iter_sample_paths(dataset_root, data_type):
|
|
sample_specs.extend(_expand_sample_path(sample_path))
|
|
return sample_specs
|
|
|
|
|
|
def _load_metadata(sample_path: Path) -> dict[str, Any]:
|
|
if sample_path.is_file():
|
|
return load_json(sample_path)
|
|
for candidate in ("meta_data.json", "metadata.json"):
|
|
meta_path = sample_path / candidate
|
|
if meta_path.exists():
|
|
return load_json(meta_path)
|
|
raise FileNotFoundError(f"Metadata not found for sample: {sample_path}")
|
|
|
|
|
|
def _sample_name(sample_path: Path) -> str:
|
|
return sample_path.stem if sample_path.is_file() else sample_path.name
|
|
|
|
|
|
def _lang_device(sample_path: Path, metadata: dict[str, Any]) -> str:
|
|
return str(metadata.get("lang_device") or sample_path.parent.name)
|
|
|
|
|
|
def _resolve_referenced_image(
|
|
*,
|
|
metadata: dict[str, Any],
|
|
sample_path: Path,
|
|
dataset_root: Path,
|
|
data_type: str,
|
|
) -> Image.Image | None:
|
|
for key in ("image", "input_image", "initial_image", "reference_image"):
|
|
image_ref = metadata.get(key)
|
|
if not image_ref:
|
|
continue
|
|
candidate = dataset_root / TYPE_TO_FOLDER[data_type] / str(image_ref)
|
|
if candidate.exists():
|
|
image = Image.open(candidate)
|
|
image.load()
|
|
return image.convert("RGB")
|
|
if sample_path.is_dir():
|
|
local_image = find_first_image(sample_path)
|
|
if local_image:
|
|
image = Image.open(local_image)
|
|
image.load()
|
|
return image.convert("RGB")
|
|
return None
|
|
|
|
|
|
def _trajectory_steps(metadata: dict[str, Any]) -> list[dict[str, Any]]:
|
|
for key in ("trajectory", "steps", "frames"):
|
|
value = metadata.get(key)
|
|
if isinstance(value, list):
|
|
return [step for step in value if isinstance(step, dict)]
|
|
extracted: list[dict[str, Any]] = []
|
|
for index in range(1, 6):
|
|
value = metadata.get(f"step{index}") or metadata.get(str(index))
|
|
if isinstance(value, dict):
|
|
extracted.append(value)
|
|
return extracted
|
|
|
|
|
|
def _text_or_default(value: Any, default: str = "") -> str:
|
|
return str(value).strip() if value is not None else default
|
|
|
|
|
|
def _type1_prompt(metadata: dict[str, Any]) -> str:
|
|
caption = _text_or_default(metadata.get("caption") or metadata.get("instruction"), "Transform the reference GUI.")
|
|
return (
|
|
"Using the reference GUI screenshot, generate the next GUI state after the requested interaction.\n\n"
|
|
f"Requested change:\n{caption}\n\n"
|
|
"Requirements:\n"
|
|
"- Preserve layout, visual identity, and unrelated regions.\n"
|
|
"- Only apply the requested state change.\n"
|
|
"- Keep all text and controls readable.\n"
|
|
)
|
|
|
|
|
|
def _type2_prompt(goal: str, step_num: int) -> str:
|
|
return (
|
|
"Generate the next GUI state for a multi-step task.\n\n"
|
|
f"Overall goal: {goal}\n"
|
|
f"Current progress step: {step_num}/5\n\n"
|
|
"Requirements:\n"
|
|
"- The change should be incremental and plausible.\n"
|
|
"- Preserve layout and visual identity.\n"
|
|
"- Make text/buttons readable.\n"
|
|
)
|
|
|
|
|
|
def _type34_initial_prompt(metadata: dict[str, Any], first_step: dict[str, Any]) -> str:
|
|
app_name = _text_or_default(metadata.get("app_name"), "App")
|
|
final_goal = _text_or_default(metadata.get("final_goal") or metadata.get("instruction"), "Complete the task.")
|
|
visual_description = _text_or_default(
|
|
metadata.get("visual_description") or first_step.get("visual_description") or first_step.get("description"),
|
|
"A clean product-quality app home screen.",
|
|
)
|
|
return (
|
|
"Generate the first GUI frame for a task trajectory.\n\n"
|
|
f"App name: {app_name}\n"
|
|
f"Final goal: {final_goal}\n"
|
|
f"Visual description:\n{visual_description}\n\n"
|
|
"Requirements:\n"
|
|
"- Generate a production-looking UI screenshot only.\n"
|
|
"- Keep the layout coherent and readable.\n"
|
|
)
|
|
|
|
|
|
def _type34_next_prompt(step_num: int, step_info: dict[str, Any]) -> str:
|
|
action = _text_or_default(step_info.get("action") or step_info.get("instruction"), "Continue the task.")
|
|
visual_description = _text_or_default(
|
|
step_info.get("visual_description") or step_info.get("description"),
|
|
"Reflect the expected next GUI state.",
|
|
)
|
|
return (
|
|
"Using the previous frame as reference, generate the next GUI frame.\n\n"
|
|
f"Step {step_num} action: {action}\n"
|
|
f"Expected visual state:\n{visual_description}\n\n"
|
|
"Requirements:\n"
|
|
"- Only change UI regions affected by this action.\n"
|
|
"- Preserve persistent bars, layout, and style.\n"
|
|
"- Keep text and icons readable.\n"
|
|
)
|
|
|
|
|
|
def _type5_prompt(metadata: dict[str, Any]) -> str:
|
|
grounding = metadata.get("grounding") or {}
|
|
explanation = _text_or_default(
|
|
metadata.get("grounding_explanation") or grounding.get("effect") or grounding.get("description"),
|
|
"Predict the immediate GUI reaction to the indicated target.",
|
|
)
|
|
return (
|
|
"Using the reference GUI screenshot, predict the immediate GUI state after the grounded interaction.\n\n"
|
|
f"Expected effect: {explanation}\n"
|
|
f"Grounding metadata: {json.dumps(grounding, ensure_ascii=False)}\n\n"
|
|
"Requirements:\n"
|
|
"- Apply only the interaction-triggered change.\n"
|
|
"- Preserve unrelated regions.\n"
|
|
"- Keep the UI realistic and readable.\n"
|
|
)
|
|
|
|
|
|
def _make_storyboard_image(
|
|
frames: list[Image.Image],
|
|
*,
|
|
columns: int = 3,
|
|
background_color: tuple[int, int, int] = (255, 255, 255),
|
|
) -> Image.Image:
|
|
if not frames:
|
|
raise ValueError("Expected at least one frame to build a storyboard image.")
|
|
|
|
normalized = [frame.convert("RGB") for frame in frames]
|
|
frame_width = max(frame.width for frame in normalized)
|
|
frame_height = max(frame.height for frame in normalized)
|
|
rows = (len(normalized) + columns - 1) // columns
|
|
storyboard = Image.new("RGB", (frame_width * columns, frame_height * rows), color=background_color)
|
|
|
|
for index, frame in enumerate(normalized):
|
|
x_offset = (index % columns) * frame_width
|
|
y_offset = (index // columns) * frame_height
|
|
if frame.size != (frame_width, frame_height):
|
|
frame = frame.resize((frame_width, frame_height))
|
|
storyboard.paste(frame, (x_offset, y_offset))
|
|
return storyboard
|
|
|
|
|
|
def _trajectory_judge_payload(frames: list[Image.Image]) -> tuple[str, list[Image.Image]]:
|
|
storyboard = _make_storyboard_image(frames, columns=3)
|
|
prompt_suffix = (
|
|
"The attached image is a storyboard containing six frames arranged left-to-right, "
|
|
"top-to-bottom as frame0, frame1, frame2, frame3, frame4, frame5."
|
|
)
|
|
return prompt_suffix, [storyboard]
|
|
|
|
|
|
class LocalJudgeClient:
|
|
def __init__(self, base_url: str, api_key: str, model: str, timeout: int = 600):
|
|
self.base_url = base_url.rstrip("/")
|
|
self.api_key = api_key
|
|
self.model = model
|
|
self.timeout = timeout
|
|
|
|
def _build_scoring_prompt(self, task_prompt: str) -> str:
|
|
return (
|
|
"You are an expert evaluator for GUI image editing and GUI trajectory generation.\n"
|
|
"Evaluate whether the generated image(s) satisfy the task.\n\n"
|
|
"Score these five dimensions from 0 to 5:\n"
|
|
"- goal: whether the user goal is completed correctly\n"
|
|
"- logic: whether the transition/state change is logically correct\n"
|
|
"- cons: whether unrelated regions remain consistent\n"
|
|
"- ui: whether the UI layout/components remain realistic and coherent\n"
|
|
"- qual: whether the images are visually clear and artifact-free\n\n"
|
|
"Return JSON only. Do not add any prose outside JSON.\n"
|
|
"Use exactly this schema:\n"
|
|
"{\n"
|
|
' "goal": 0,\n'
|
|
' "logic": 0,\n'
|
|
' "cons": 0,\n'
|
|
' "ui": 0,\n'
|
|
' "qual": 0,\n'
|
|
' "reasoning": "short explanation"\n'
|
|
"}\n\n"
|
|
"Scoring task:\n"
|
|
f"{task_prompt}"
|
|
)
|
|
|
|
def _request_text(self, prompt: str, images: list[Image.Image]) -> str:
|
|
content: list[dict[str, Any]] = [{"type": "text", "text": prompt}]
|
|
for image in images:
|
|
content.append({"type": "image_url", "image_url": {"url": pil_to_data_url(image)}})
|
|
|
|
response = requests.post(
|
|
build_openai_url(self.base_url, "/chat/completions"),
|
|
json={
|
|
"model": self.model,
|
|
"messages": [{"role": "user", "content": content}],
|
|
"temperature": 0,
|
|
},
|
|
headers={
|
|
"Authorization": f"Bearer {self.api_key}",
|
|
"Content-Type": "application/json",
|
|
},
|
|
timeout=self.timeout,
|
|
)
|
|
response.raise_for_status()
|
|
message_content = response.json()["choices"][0]["message"]["content"]
|
|
if isinstance(message_content, list):
|
|
return "\n".join(part.get("text", "") for part in message_content if part.get("type") == "text")
|
|
return str(message_content)
|
|
|
|
def evaluate(self, *, prompt: str, images: list[Image.Image]) -> dict[str, Any]:
|
|
primary_prompt = self._build_scoring_prompt(prompt)
|
|
text = self._request_text(primary_prompt, images)
|
|
try:
|
|
return extract_json_object(text)
|
|
except ValueError:
|
|
retry_prompt = (
|
|
self._build_scoring_prompt(prompt) + "\n\nYour previous response was not valid JSON. "
|
|
"Return only the JSON object with integer scores."
|
|
)
|
|
retry_text = self._request_text(retry_prompt, images)
|
|
try:
|
|
return extract_json_object(retry_text)
|
|
except ValueError:
|
|
return {
|
|
"goal": 0,
|
|
"logic": 0,
|
|
"cons": 0,
|
|
"ui": 0,
|
|
"qual": 0,
|
|
"reasoning": retry_text.strip() or text.strip() or "Judge response was not valid JSON.",
|
|
}
|
|
|
|
|
|
class GEBenchRunner:
|
|
def __init__(
|
|
self,
|
|
*,
|
|
dataset_root: Path,
|
|
output_root: Path,
|
|
base_url: str,
|
|
model: str,
|
|
api_key: str = "EMPTY",
|
|
width: int = 768,
|
|
height: int = 576,
|
|
num_inference_steps: int = 8,
|
|
output_compression: int | None = 98,
|
|
guidance_scale: float | None = None,
|
|
seed: int | None = 42,
|
|
):
|
|
self.dataset_root = dataset_root
|
|
self.output_root = output_root
|
|
self.model = model
|
|
self.width = width
|
|
self.height = height
|
|
self.num_inference_steps = num_inference_steps
|
|
self.output_compression = output_compression
|
|
self.guidance_scale = guidance_scale
|
|
self.seed = seed
|
|
self.client = VllmOmniImageClient(base_url=base_url, api_key=api_key)
|
|
|
|
def generate(
|
|
self,
|
|
*,
|
|
data_type: str,
|
|
workers: int = 1,
|
|
max_samples: int | None = None,
|
|
samples_per_type: int | None = None,
|
|
) -> list[dict[str, Any]]:
|
|
sample_specs = select_balanced_gebench_samples(
|
|
{data_type: _iter_sample_specs(self.dataset_root, data_type)},
|
|
samples_per_type=samples_per_type,
|
|
)[data_type]
|
|
if max_samples is not None:
|
|
sample_specs = sample_specs[:max_samples]
|
|
|
|
results: list[dict[str, Any]] = []
|
|
if workers <= 1:
|
|
for sample_spec in sample_specs:
|
|
result = self._generate_one(data_type, sample_spec)
|
|
if result:
|
|
results.append(result)
|
|
return results
|
|
|
|
with ThreadPoolExecutor(max_workers=workers) as executor:
|
|
futures = [executor.submit(self._generate_one, data_type, sample_spec) for sample_spec in sample_specs]
|
|
for future in as_completed(futures):
|
|
result = future.result()
|
|
if result:
|
|
results.append(result)
|
|
return results
|
|
|
|
def _generate_one(self, data_type: str, sample_spec: GEBenchSampleSpec) -> dict[str, Any] | None:
|
|
sample_path = sample_spec.sample_path
|
|
metadata = sample_spec.metadata
|
|
lang_device = sample_spec.lang_device
|
|
sample_name = sample_spec.sample_name
|
|
output_dir = ensure_dir(self.output_root / TYPE_TO_FOLDER[data_type] / lang_device / sample_name)
|
|
|
|
if data_type == "type1":
|
|
output_path = output_dir / "generated.png"
|
|
if output_path.exists():
|
|
return {
|
|
"data_type": data_type,
|
|
"sample_name": f"{lang_device}/{sample_name}",
|
|
"output_path": str(output_path),
|
|
}
|
|
source = _resolve_referenced_image(
|
|
metadata=metadata, sample_path=sample_path, dataset_root=self.dataset_root, data_type=data_type
|
|
)
|
|
if source is None:
|
|
return None
|
|
generated = self.client.generate_image_edit(
|
|
model=self.model,
|
|
prompt=_type1_prompt(metadata),
|
|
images=source,
|
|
width=self.width,
|
|
height=self.height,
|
|
num_inference_steps=self.num_inference_steps,
|
|
output_compression=self.output_compression,
|
|
guidance_scale=self.guidance_scale,
|
|
seed=self.seed,
|
|
)
|
|
save_image(output_path, generated)
|
|
return {
|
|
"data_type": data_type,
|
|
"sample_name": f"{lang_device}/{sample_name}",
|
|
"output_path": str(output_path),
|
|
}
|
|
|
|
if data_type == "type2":
|
|
goal = _text_or_default(metadata.get("question") or metadata.get("caption"), "Complete the task.")
|
|
source = _resolve_referenced_image(
|
|
metadata=metadata, sample_path=sample_path, dataset_root=self.dataset_root, data_type=data_type
|
|
)
|
|
if source is None:
|
|
return None
|
|
frame0_path = output_dir / "frame0.png"
|
|
if not frame0_path.exists():
|
|
save_image(frame0_path, source)
|
|
previous = source
|
|
for step_num in range(1, 6):
|
|
frame_path = output_dir / f"frame{step_num}.png"
|
|
if frame_path.exists():
|
|
previous = Image.open(frame_path).convert("RGB")
|
|
continue
|
|
generated = self.client.generate_image_edit(
|
|
model=self.model,
|
|
prompt=_type2_prompt(goal, step_num),
|
|
images=previous,
|
|
width=self.width,
|
|
height=self.height,
|
|
num_inference_steps=self.num_inference_steps,
|
|
output_compression=self.output_compression,
|
|
guidance_scale=self.guidance_scale,
|
|
seed=self.seed,
|
|
)
|
|
save_image(frame_path, generated)
|
|
previous = generated
|
|
output_path = output_dir / "frame5.png"
|
|
return {
|
|
"data_type": data_type,
|
|
"sample_name": f"{lang_device}/{sample_name}",
|
|
"output_path": str(output_path),
|
|
}
|
|
|
|
if data_type in {"type3", "type4"}:
|
|
steps = _trajectory_steps(metadata)
|
|
frame0_path = output_dir / "frame0.png"
|
|
if frame0_path.exists():
|
|
previous = Image.open(frame0_path).convert("RGB")
|
|
else:
|
|
previous = self.client.generate_text_to_image(
|
|
model=self.model,
|
|
prompt=_type34_initial_prompt(metadata, steps[0] if steps else {}),
|
|
width=self.width,
|
|
height=self.height,
|
|
num_inference_steps=self.num_inference_steps,
|
|
output_compression=self.output_compression,
|
|
guidance_scale=self.guidance_scale,
|
|
seed=self.seed,
|
|
)
|
|
save_image(frame0_path, previous)
|
|
|
|
for step_num in range(1, 6):
|
|
frame_path = output_dir / f"frame{step_num}.png"
|
|
if frame_path.exists():
|
|
previous = Image.open(frame_path).convert("RGB")
|
|
continue
|
|
step_info = steps[step_num - 1] if step_num - 1 < len(steps) else {}
|
|
generated = self.client.generate_image_edit(
|
|
model=self.model,
|
|
prompt=_type34_next_prompt(step_num, step_info),
|
|
images=previous,
|
|
width=self.width,
|
|
height=self.height,
|
|
num_inference_steps=self.num_inference_steps,
|
|
output_compression=self.output_compression,
|
|
guidance_scale=self.guidance_scale,
|
|
seed=self.seed,
|
|
)
|
|
save_image(frame_path, generated)
|
|
previous = generated
|
|
output_path = output_dir / "frame5.png"
|
|
return {
|
|
"data_type": data_type,
|
|
"sample_name": f"{lang_device}/{sample_name}",
|
|
"output_path": str(output_path),
|
|
}
|
|
|
|
if data_type == "type5":
|
|
output_path = output_dir / "generated.png"
|
|
if output_path.exists():
|
|
return {
|
|
"data_type": data_type,
|
|
"sample_name": f"{lang_device}/{sample_name}",
|
|
"output_path": str(output_path),
|
|
}
|
|
source = _resolve_referenced_image(
|
|
metadata=metadata, sample_path=sample_path, dataset_root=self.dataset_root, data_type=data_type
|
|
)
|
|
if source is None:
|
|
return None
|
|
generated = self.client.generate_image_edit(
|
|
model=self.model,
|
|
prompt=_type5_prompt(metadata),
|
|
images=source,
|
|
width=self.width,
|
|
height=self.height,
|
|
num_inference_steps=self.num_inference_steps,
|
|
output_compression=self.output_compression,
|
|
guidance_scale=self.guidance_scale,
|
|
seed=self.seed,
|
|
)
|
|
save_image(output_path, generated)
|
|
return {
|
|
"data_type": data_type,
|
|
"sample_name": f"{lang_device}/{sample_name}",
|
|
"output_path": str(output_path),
|
|
}
|
|
|
|
raise ValueError(f"Unsupported data type: {data_type}")
|
|
|
|
|
|
class GEBenchEvaluator:
|
|
def __init__(self, *, dataset_root: Path, output_root: Path, judge: LocalJudgeClient):
|
|
self.dataset_root = dataset_root
|
|
self.output_root = output_root
|
|
self.judge = judge
|
|
|
|
def evaluate(
|
|
self,
|
|
*,
|
|
data_type: str,
|
|
workers: int = 1,
|
|
max_samples: int | None = None,
|
|
samples_per_type: int | None = None,
|
|
) -> dict[str, Any]:
|
|
output_type_dir = self.output_root / TYPE_TO_FOLDER[data_type]
|
|
sample_specs_by_name = {
|
|
(spec.lang_device, spec.sample_name): spec for spec in _iter_sample_specs(self.dataset_root, data_type)
|
|
}
|
|
if not output_type_dir.exists():
|
|
payload = {"data_type": data_type, "results": [], "summary": summarize_gebench_results([])}
|
|
write_json(self.output_root / "evaluations" / f"{data_type}.json", payload)
|
|
return payload
|
|
sample_dirs = [
|
|
sample_dir
|
|
for lang_dir in sorted(path for path in output_type_dir.iterdir() if path.is_dir())
|
|
for sample_dir in sorted(path for path in lang_dir.iterdir() if path.is_dir())
|
|
if (lang_dir.name, sample_dir.name) in sample_specs_by_name
|
|
]
|
|
sample_dirs = select_balanced_gebench_samples(
|
|
{data_type: sample_dirs},
|
|
samples_per_type=samples_per_type,
|
|
)[data_type]
|
|
if max_samples is not None:
|
|
sample_dirs = sample_dirs[:max_samples]
|
|
results: list[dict[str, Any]] = []
|
|
if workers <= 1:
|
|
for sample_dir in sample_dirs:
|
|
result = self._evaluate_one(
|
|
data_type,
|
|
sample_dir,
|
|
sample_specs_by_name[(sample_dir.parent.name, sample_dir.name)],
|
|
)
|
|
if result:
|
|
results.append(result)
|
|
else:
|
|
with ThreadPoolExecutor(max_workers=workers) as executor:
|
|
futures = [
|
|
executor.submit(
|
|
self._evaluate_one,
|
|
data_type,
|
|
sample_dir,
|
|
sample_specs_by_name[(sample_dir.parent.name, sample_dir.name)],
|
|
)
|
|
for sample_dir in sample_dirs
|
|
]
|
|
for future in as_completed(futures):
|
|
result = future.result()
|
|
if result:
|
|
results.append(result)
|
|
|
|
payload = {"data_type": data_type, "results": results, "summary": summarize_gebench_results(results)}
|
|
write_json(self.output_root / "evaluations" / f"{data_type}.json", payload)
|
|
return payload
|
|
|
|
def _evaluate_one(self, data_type: str, sample_dir: Path, sample_spec: GEBenchSampleSpec) -> dict[str, Any] | None:
|
|
lang_device = sample_dir.parent.name
|
|
sample_name = sample_dir.name
|
|
dataset_sample = sample_spec.sample_path
|
|
metadata = sample_spec.metadata
|
|
|
|
if data_type == "type1":
|
|
source = _resolve_referenced_image(
|
|
metadata=metadata, sample_path=dataset_sample, dataset_root=self.dataset_root, data_type=data_type
|
|
)
|
|
generated_path = find_first_image(sample_dir)
|
|
if source is None or generated_path is None:
|
|
return None
|
|
generated = Image.open(generated_path).convert("RGB")
|
|
raw_scores = self.judge.evaluate(prompt=_type1_prompt(metadata), images=[source, generated])
|
|
elif data_type == "type2":
|
|
frames = [Image.open(sample_dir / f"frame{i}.png").convert("RGB") for i in range(6)]
|
|
goal = _text_or_default(metadata.get("question") or metadata.get("caption"), "Complete the task.")
|
|
prompt_suffix, judge_images = _trajectory_judge_payload(frames)
|
|
raw_scores = self.judge.evaluate(
|
|
prompt=f"Evaluate a six-frame GUI trajectory.\nTask: {goal}\n{prompt_suffix}",
|
|
images=judge_images,
|
|
)
|
|
elif data_type in {"type3", "type4"}:
|
|
frames = [Image.open(sample_dir / f"frame{i}.png").convert("RGB") for i in range(6)]
|
|
instruction = _text_or_default(metadata.get("instruction") or metadata.get("caption"), "Complete the task.")
|
|
prompt_suffix, judge_images = _trajectory_judge_payload(frames)
|
|
raw_scores = self.judge.evaluate(
|
|
prompt=f"Evaluate a six-frame GUI trajectory.\nInstruction: {instruction}\n{prompt_suffix}",
|
|
images=judge_images,
|
|
)
|
|
elif data_type == "type5":
|
|
source = _resolve_referenced_image(
|
|
metadata=metadata, sample_path=dataset_sample, dataset_root=self.dataset_root, data_type=data_type
|
|
)
|
|
generated_path = find_first_image(sample_dir)
|
|
if source is None or generated_path is None:
|
|
return None
|
|
generated = Image.open(generated_path).convert("RGB")
|
|
raw_scores = self.judge.evaluate(prompt=_type5_prompt(metadata), images=[source, generated])
|
|
else:
|
|
raise ValueError(f"Unsupported data type: {data_type}")
|
|
|
|
scores = _normalize_scores(raw_scores)
|
|
return {
|
|
"sample_name": f"{lang_device}/{sample_name}",
|
|
"data_type": data_type,
|
|
"scores": scores,
|
|
"overall": _compute_overall(scores),
|
|
"raw_scores": raw_scores,
|
|
}
|
|
|
|
|
|
def _data_types_arg(value: str) -> list[str]:
|
|
return list(TYPE_TO_FOLDER.keys()) if value == "all" else [value]
|
|
|
|
|
|
def build_parser() -> argparse.ArgumentParser:
|
|
parser = argparse.ArgumentParser(description="Run local GEBench generation and scoring against vLLM-Omni.")
|
|
subparsers = parser.add_subparsers(dest="command", required=True)
|
|
|
|
generate = subparsers.add_parser("generate")
|
|
generate.add_argument("--dataset-root", type=Path, required=True)
|
|
generate.add_argument("--output-root", type=Path, required=True)
|
|
generate.add_argument("--base-url", type=str, required=True)
|
|
generate.add_argument("--model", type=str, required=True)
|
|
generate.add_argument("--data-type", choices=["all", *TYPE_TO_FOLDER.keys()], default="all")
|
|
generate.add_argument("--api-key", type=str, default="EMPTY")
|
|
generate.add_argument("--width", type=int, default=768)
|
|
generate.add_argument("--height", type=int, default=576)
|
|
generate.add_argument("--num-inference-steps", type=int, default=8)
|
|
generate.add_argument("--output-compression", type=int, default=98)
|
|
generate.add_argument("--guidance-scale", type=float, default=None)
|
|
generate.add_argument("--seed", type=int, default=42)
|
|
generate.add_argument("--workers", type=int, default=1)
|
|
generate.add_argument("--max-samples", type=int, default=None)
|
|
generate.add_argument("--samples-per-type", type=int, default=None)
|
|
|
|
evaluate = subparsers.add_parser("evaluate")
|
|
evaluate.add_argument("--dataset-root", type=Path, required=True)
|
|
evaluate.add_argument("--output-root", type=Path, required=True)
|
|
evaluate.add_argument("--data-type", choices=["all", *TYPE_TO_FOLDER.keys()], default="all")
|
|
evaluate.add_argument("--judge-base-url", type=str, required=True)
|
|
evaluate.add_argument("--judge-model", type=str, required=True)
|
|
evaluate.add_argument("--judge-api-key", type=str, default="EMPTY")
|
|
evaluate.add_argument("--workers", type=int, default=1)
|
|
evaluate.add_argument("--max-samples", type=int, default=None)
|
|
evaluate.add_argument("--samples-per-type", type=int, default=None)
|
|
|
|
summarize = subparsers.add_parser("summarize")
|
|
summarize.add_argument("--output-root", type=Path, required=True)
|
|
|
|
return parser
|
|
|
|
|
|
def main(argv: list[str] | None = None) -> int:
|
|
parser = build_parser()
|
|
args = parser.parse_args(argv)
|
|
|
|
if args.command == "generate":
|
|
runner = GEBenchRunner(
|
|
dataset_root=args.dataset_root,
|
|
output_root=args.output_root,
|
|
base_url=args.base_url,
|
|
model=args.model,
|
|
api_key=args.api_key,
|
|
width=args.width,
|
|
height=args.height,
|
|
num_inference_steps=args.num_inference_steps,
|
|
output_compression=args.output_compression,
|
|
guidance_scale=args.guidance_scale,
|
|
seed=args.seed,
|
|
)
|
|
records: list[dict[str, Any]] = []
|
|
for data_type in _data_types_arg(args.data_type):
|
|
records.extend(
|
|
runner.generate(
|
|
data_type=data_type,
|
|
workers=args.workers,
|
|
max_samples=args.max_samples,
|
|
samples_per_type=args.samples_per_type,
|
|
)
|
|
)
|
|
payload = {"records": records, "summary": summarize_generated_records(records)}
|
|
write_json(args.output_root / "generation_manifest.json", payload)
|
|
return 0
|
|
|
|
if args.command == "evaluate":
|
|
judge = LocalJudgeClient(
|
|
base_url=args.judge_base_url,
|
|
api_key=args.judge_api_key,
|
|
model=args.judge_model,
|
|
)
|
|
evaluator = GEBenchEvaluator(dataset_root=args.dataset_root, output_root=args.output_root, judge=judge)
|
|
combined_results: list[dict[str, Any]] = []
|
|
for data_type in _data_types_arg(args.data_type):
|
|
payload = evaluator.evaluate(
|
|
data_type=data_type,
|
|
workers=args.workers,
|
|
max_samples=args.max_samples,
|
|
samples_per_type=args.samples_per_type,
|
|
)
|
|
combined_results.extend(payload["results"])
|
|
_write_json_with_timestamp(
|
|
args.output_root / "evaluations" / "summary.json",
|
|
{"summary": summarize_gebench_results(combined_results)},
|
|
)
|
|
return 0
|
|
|
|
if args.command == "summarize":
|
|
generation_summary = collect_gebench_generation_summary(args.output_root)
|
|
evaluation_dir = args.output_root / "evaluations"
|
|
result_records: list[dict[str, Any]] = []
|
|
if evaluation_dir.exists():
|
|
for file_path in sorted(evaluation_dir.glob("type*.json")):
|
|
payload = load_json(file_path)
|
|
result_records.extend(payload.get("results", []))
|
|
payload: dict[str, Any] = {"generation": generation_summary}
|
|
if result_records:
|
|
payload["evaluation"] = summarize_gebench_results(result_records)
|
|
_write_json_with_timestamp(args.output_root / "summary.json", payload)
|
|
print(json.dumps(payload, indent=2, ensure_ascii=False))
|
|
return 0
|
|
|
|
parser.error(f"Unknown command: {args.command}")
|
|
return 1
|