Spaces:
Running on Zero
Running on Zero
| """Interactive Seg2Click demo for Hugging Face Spaces.""" | |
| # ruff: noqa: I001 -- spaces must be imported before torch and its transitive importers. | |
| import spaces # Import before torch so ZeroGPU can configure the runtime. | |
| import time | |
| import gradio as gr | |
| import numpy as np | |
| import torch | |
| from huggingface_hub import snapshot_download | |
| from PIL import Image, ImageDraw | |
| from seg2click.predictor import Seg2ClickPredictor | |
| MODEL_ID = "jcwang0602/Seg2Click-Qwen3.5-0.8B" | |
| MASK_THRESHOLD = 0.5 | |
| EXAMPLES = [ | |
| ["examples/vscode_refresh.png", "Refresh the file explorer."], | |
| ["examples/excel_sort.png", "open the sort dialog"], | |
| ["examples/word_insert_picture.png", "the pictures inserting function in word"], | |
| ["examples/photoshop_shape_mode.png", "change to shape mode"], | |
| ["examples/macos_cancel_extraction.png", "cancel extraction"], | |
| ] | |
| print(f"Downloading {MODEL_ID} ...", flush=True) | |
| checkpoint_dir = snapshot_download(MODEL_ID) | |
| print(f"Loading model from {checkpoint_dir} ...", flush=True) | |
| predictor = Seg2ClickPredictor.from_pretrained( | |
| checkpoint_dir, | |
| device="cuda", | |
| attn_implementation="sdpa", | |
| dtype=torch.bfloat16, | |
| ) | |
| print("Seg2Click is ready.", flush=True) | |
| def render_overlay(image: Image.Image, probability: np.ndarray, x: float, y: float) -> Image.Image: | |
| """Overlay the action-mask heatmap and predicted click on a screenshot.""" | |
| canvas = image.convert("RGB") | |
| width, height = canvas.size | |
| probability_image = Image.fromarray( | |
| (np.clip(probability, 0, 1) * 255).astype(np.uint8), mode="L" | |
| ).resize((width, height), resample=Image.Resampling.BILINEAR) | |
| values = np.asarray(probability_image, dtype=np.float32) / 255.0 | |
| heatmap = np.stack( | |
| [ | |
| np.clip(1.5 - np.abs(4.0 * values - 3.0), 0, 1), | |
| np.clip(1.5 - np.abs(4.0 * values - 2.0), 0, 1), | |
| np.clip(1.5 - np.abs(4.0 * values - 1.0), 0, 1), | |
| ], | |
| axis=-1, | |
| ) | |
| base = np.asarray(canvas, dtype=np.float32) | |
| blended = base * (1 - 0.55 * values[..., None]) + heatmap * 255 * (0.55 * values[..., None]) | |
| foreground = (values >= MASK_THRESHOLD).astype(np.float32) | |
| blended *= 0.75 + 0.25 * foreground[..., None] | |
| canvas = Image.fromarray(blended.astype(np.uint8), mode="RGB") | |
| px = round(x * width) | |
| py = round(y * height) | |
| marker = Image.new("RGBA", canvas.size, (0, 0, 0, 0)) | |
| draw = ImageDraw.Draw(marker) | |
| draw.ellipse((px - 14, py - 14, px + 14, py + 14), fill=(255, 0, 0, 190)) | |
| draw.ellipse((px - 26, py - 26, px + 26, py + 26), outline=(255, 215, 0, 255), width=5) | |
| return Image.alpha_composite(canvas.convert("RGBA"), marker).convert("RGB") | |
| def predict_click(screenshot: Image.Image, instruction: str): | |
| """Predict an action mask and executable click point.""" | |
| if screenshot is None: | |
| raise gr.Error("Please upload a screenshot first.") | |
| instruction = (instruction or "").strip() | |
| if not instruction: | |
| raise gr.Error("Please enter an instruction, such as 'click the Submit button'.") | |
| image = screenshot.convert("RGB") | |
| started = time.perf_counter() | |
| with torch.inference_mode(): | |
| prediction = predictor.predict(image, instruction) | |
| elapsed = time.perf_counter() - started | |
| x, y = prediction.point | |
| width, height = image.size | |
| overlay = render_overlay(image, prediction.probability.numpy(), x, y) | |
| report = ( | |
| f"**Click position:** `({x:.4f}, {y:.4f})` (normalized) \n" | |
| f"**Pixel position:** `({round(x * width)}, {round(y * height)})` \n" | |
| f"**Decode:** `{prediction.decode_reason}` \n" | |
| f"**Inference time:** `{elapsed:.1f}s`" | |
| ) | |
| return overlay, report | |
| CSS = "#app-container { max-width: 1200px; margin: 0 auto; }" | |
| with gr.Blocks(title="Seg2Click — GUI Visual Grounding") as demo: | |
| gr.Markdown( | |
| """ | |
| # 🖱️ Seg2Click: Pixel-Level Action Masks for GUI Visual Grounding | |
| Upload a GUI screenshot and describe what you want to click. Seg2Click predicts the complete | |
| executable region as a pixel-level action mask, then returns its centroid as the click point. | |
| This demo runs [Seg2Click-Qwen3.5-0.8B](https://huggingface.co/jcwang0602/Seg2Click-Qwen3.5-0.8B). | |
| [Paper](https://arxiv.org/abs/2610.16234) · [Code](https://github.com/jcwang0602/Seg2Click) | |
| """ | |
| ) | |
| with gr.Row(elem_id="app-container"): | |
| with gr.Column(): | |
| screenshot_input = gr.Image( | |
| type="pil", | |
| label="GUI screenshot", | |
| height=440, | |
| sources=["upload", "clipboard"], | |
| ) | |
| instruction_input = gr.Textbox( | |
| label="Instruction", | |
| placeholder="e.g. click the refresh button in the sidebar", | |
| ) | |
| predict_button = gr.Button("Predict click", variant="primary") | |
| with gr.Column(): | |
| output_image = gr.Image(type="pil", label="Action mask and click point", height=440) | |
| output_report = gr.Markdown() | |
| gr.Markdown("## Examples\nClick an example to run fresh model inference on the selected screenshot.") | |
| gr.Examples( | |
| examples=EXAMPLES, | |
| inputs=[screenshot_input, instruction_input], | |
| outputs=[output_image, output_report], | |
| fn=predict_click, | |
| cache_examples=False, | |
| run_on_click=True, | |
| ) | |
| predict_button.click( | |
| predict_click, | |
| inputs=[screenshot_input, instruction_input], | |
| outputs=[output_image, output_report], | |
| api_name="predict_click", | |
| ) | |
| instruction_input.submit( | |
| predict_click, | |
| inputs=[screenshot_input, instruction_input], | |
| outputs=[output_image, output_report], | |
| ) | |
| demo.launch(theme=gr.themes.Citrus(), css=CSS) | |