Spaces:
Running on Zero
Running on Zero
| """Wave-UI (agentsea/wave-ui): UI elements with instructions and boxes, used as clicks in the agent | |
| layout and as element boxes. The ScreenSpot source and every image that also appears in ScreenSpot | |
| are removed, so ScreenSpot stays a clean test set.""" | |
| import hashlib | |
| from . import load_cache, save_cache | |
| # sources whose "instruction" is element-type metadata; their element name reads like an instruction | |
| NAME_SOURCES = ("webui", "roboflow", "mind2web_test_domain", "mind2web_test_task", "mind2web_test_website") | |
| def screenspot_hashes(): | |
| from datasets import load_dataset | |
| key = lambda im: hashlib.md5(im.convert("L").resize((64, 64)).tobytes()).hexdigest() | |
| return key, {key(im) for im in load_dataset("bevaya/ScreenSpot", split="test")["image"]} | |
| def clean(split, n, seed, key, banned): | |
| from datasets import load_dataset | |
| ds = load_dataset("agentsea/wave-ui", split=split).shuffle(seed=seed) | |
| ds = ds.filter(lambda src: src != "screenspot", input_columns=["source"]) | |
| out = [] | |
| for r in ds: | |
| if len(out) >= n: | |
| break | |
| x1, y1, x2, y2 = r["bbox"] | |
| W, H = r["image"].size | |
| if not r["instruction"] or x2 <= x1 or y2 <= y1 or x2 > W * 1.02 or y2 > H * 1.02: | |
| continue | |
| if key(r["image"]) in banned: | |
| continue | |
| if r["source"] in NAME_SOURCES: | |
| r = dict(r, instruction=r["name"]) | |
| if not r["instruction"]: | |
| continue | |
| out.append(r) | |
| return out | |
| def clicks(n, skip=0, cache=True): | |
| """Wave-UI elements as click steps in the agent layout (keeps the pointing sharp).""" | |
| name = f"waveui_clicks_{n}_{skip}" | |
| if cache and (got := load_cache(name)) is not None: | |
| return got | |
| key, banned = screenspot_hashes() | |
| out = [] | |
| for r in clean("train", skip + n, 0, key, banned)[skip:]: | |
| W, H = r["image"].size | |
| x1, y1, x2, y2 = r["bbox"] | |
| out.append(dict(image=r["image"], user=f"Instruction: {r['instruction']}\n\nPrevious actions:\nNone", | |
| action="click", source="wave-ui", | |
| args=dict(x=(x1 + x2) / 2 / W, y=(y1 + y2) / 2 / H, box=(x2 - x1) / W, boy=(y2 - y1) / H))) | |
| if cache: | |
| save_cache(name, out) | |
| return out | |
| def boxes(n, skip=2000, cache=True): | |
| """Wave-UI element boxes (pixels). Rows [skip, skip + n) are disjoint from the clicks [0, n_click).""" | |
| name = f"waveui_boxes_{n}_{skip}" | |
| if cache and (got := load_cache(name)) is not None: | |
| return got | |
| key, banned = screenspot_hashes() | |
| out = [dict(image=r["image"], instruction=r["instruction"], box=list(r["bbox"]), source=r["source"]) | |
| for r in clean("train", skip + n, 0, key, banned)[skip:]] | |
| if cache: | |
| save_cache(name, out) | |
| return out | |
| def val_boxes(n): | |
| key, banned = screenspot_hashes() | |
| return [dict(image=r["image"], instruction=r["instruction"], box=list(r["bbox"]), source=r["source"]) | |
| for r in clean("validation", n, 1, key, banned)] | |