Create app.py

app.py

@@ -0,0 +1,281 @@
+import spaces
+
+import re
+from typing import Tuple, Optional
+
+import gradio as gr
+import numpy as np
+from PIL import Image, ImageDraw, ImageFont
+from smolvlm_inference import TransformersModel
+
+from prompt import OS_SYSTEM_PROMPT
+
+# --- Configuration ---
+MODEL_ID = "smolagents/SmolVLM2-2.2B-Instruct-Agentic-GUI"
+
+# --- Model and Processor Loading (Load once) ---
+print(f"Loading model and processor for {MODEL_ID}...")
+model = None
+processor = None
+model_loaded = False
+load_error_message = ""
+
+
+
+model = TransformersModel(
+    model_id=MODEL_ID,
+    to_device="cuda:0",
+)
+
+
+title = "Smol2Operator Demo"
+
+description = """
+This is a demo of the Smol2Operator model designed to interact with graphical user interfaces (GUIs) and perform actions within them.
+This proof-of-concept (POC) version, described in [blogpost], showcases the model’s core capabilities.
+This compact release is intentionally scoped to fundamental tasks, with complex workflows planned for future iterations. :hugging_face: 
+"""
+
+
+
+SYSTEM_PROMPT: str = OS_SYSTEM_PROMPT
+
+
+def get_navigation_prompt(task, image, step=1):
+    """
+    Get the prompt for the navigation task.
+    - task: The task to complete
+    - image: The current screenshot of the web page
+    - step: The current step of the task
+    """
+    system_prompt = SYSTEM_PROMPT
+    return [
+        {
+            "role": "system",
+            "content": [
+                {"type": "text", "text": system_prompt},
+            ],
+        },
+        {
+            "role": "user",
+            "content": [
+                {
+                    "type": "image",
+                    "image": image,
+                },
+                {"type": "text", "text": f"Please generate the next move according to the UI screenshot, instruction and previous actions.\n\nInstruction: {task}\n\nPrevious actions:\nNone"},
+            ],
+        },
+    ]
+
+
+def array_to_image(image_array: np.ndarray) -> Image.Image:
+    if image_array is None:
+        raise ValueError("No image provided. Please upload an image before submitting.")
+    # Convert numpy array to PIL Image
+    img = Image.fromarray(np.uint8(image_array))
+    return img
+
+
+def parse_actions_from_response(response: str) -> list[str]:
+    """Parse actions from model response using regex pattern."""
+    pattern = r"<code>\n(.*?)\n</code>"
+    matches = re.findall(pattern, response, re.DOTALL)
+    return matches
+
+
+def extract_coordinates_from_action(action_code: str) -> list[dict]:
+    """Extract coordinates from action code for localization actions."""
+    localization_actions = []
+    
+    # Patterns for different action types
+    patterns = {
+        'click': r'click\((?:x=)?([0-9.]+)(?:,\s*(?:y=)?([0-9.]+))?\)',
+        'double_click': r'double_click\((?:x=)?([0-9.]+)(?:,\s*(?:y=)?([0-9.]+))?\)',
+        'move_mouse': r'move_mouse\((?:self,\s*)?(?:x=)?([0-9.]+)(?:,\s*(?:y=)?([0-9.]+))\)',
+        'drag': r'drag\(\[([0-9.]+),\s*([0-9.]+)\],\s*\[([0-9.]+),\s*([0-9.]+)\]\)'
+    }
+    
+    for action_type, pattern in patterns.items():
+        matches = re.finditer(pattern, action_code)
+        for match in matches:
+            if action_type == 'drag':
+                # Drag has from and to coordinates
+                from_x, from_y, to_x, to_y = match.groups()
+                localization_actions.append({
+                    'type': 'drag_from',
+                    'x': float(from_x),
+                    'y': float(from_y),
+                    'action': action_type
+                })
+                localization_actions.append({
+                    'type': 'drag_to',
+                    'x': float(to_x),
+                    'y': float(to_y),
+                    'action': action_type
+                })
+            else:
+                # Single coordinate actions
+                x_val = match.group(1)
+                y_val = match.group(2) if match.group(2) else x_val  # Handle single coordinate case
+                if x_val and y_val:
+                    localization_actions.append({
+                        'type': action_type,
+                        'x': float(x_val),
+                        'y': float(y_val),
+                        'action': action_type
+                    })
+    
+    return localization_actions
+
+
+def create_localized_image(original_image: Image.Image, coordinates: list[dict]) -> Optional[Image.Image]:
+    """Create an image with localization markers drawn on it."""
+    if not coordinates:
+        return None
+    
+    # Create a copy of the original image
+    img_copy = original_image.copy()
+    draw = ImageDraw.Draw(img_copy)
+    
+    # Get image dimensions
+    width, height = img_copy.size
+    
+    # Try to load a font, fallback to default if not available
+    font = ImageFont.load_default()
+
+    
+    # Color scheme for different actions
+    colors = {
+        'click': 'red',
+        'double_click': 'blue',
+        'move_mouse': 'green',
+        'drag_from': 'orange',
+        'drag_to': 'purple'
+    }
+    
+    for i, coord in enumerate(coordinates):
+        # Convert normalized coordinates to pixel coordinates
+        pixel_x = int(coord['x'] * width)
+        pixel_y = int(coord['y'] * height)
+        
+        # Get color for this action type
+        color = colors.get(coord['type'], 'red')
+        
+        # Draw a circle at the coordinate
+        circle_radius = 8
+        draw.ellipse([
+            pixel_x - circle_radius, pixel_y - circle_radius,
+            pixel_x + circle_radius, pixel_y + circle_radius
+        ], fill=color, outline='white', width=2)
+        
+        # Add text label
+        label = f"{coord['type']}({coord['x']:.2f},{coord['y']:.2f})"
+        if font:
+            draw.text((pixel_x + 10, pixel_y - 10), label, fill=color, font=font)
+        else:
+            draw.text((pixel_x + 10, pixel_y - 10), label, fill=color)
+        
+        # For drag actions, draw an arrow
+        if coord['type'] == 'drag_from' and i + 1 < len(coordinates) and coordinates[i + 1]['type'] == 'drag_to':
+            next_coord = coordinates[i + 1]
+            end_x = int(next_coord['x'] * width)
+            end_y = int(next_coord['y'] * height)
+            
+            # Draw arrow line
+            draw.line([pixel_x, pixel_y, end_x, end_y], fill='orange', width=3)
+            
+            # Draw arrowhead
+            arrow_size = 10
+            dx = end_x - pixel_x
+            dy = end_y - pixel_y
+            length = (dx**2 + dy**2)**0.5
+            if length > 0:
+                dx_norm = dx / length
+                dy_norm = dy / length
+                
+                # Arrowhead points
+                arrow_x1 = end_x - arrow_size * dx_norm + arrow_size * dy_norm * 0.5
+                arrow_y1 = end_y - arrow_size * dy_norm - arrow_size * dx_norm * 0.5
+                arrow_x2 = end_x - arrow_size * dx_norm - arrow_size * dy_norm * 0.5
+                arrow_y2 = end_y - arrow_size * dy_norm + arrow_size * dx_norm * 0.5
+                
+                draw.polygon([end_x, end_y, arrow_x1, arrow_y1, arrow_x2, arrow_y2], fill='orange')
+    
+    return img_copy
+
+
+# --- Gradio processing function ---
+@spaces.GPU
+def navigate(input_numpy_image: np.ndarray, task: str) -> Tuple[str, Optional[Image.Image]]:
+    input_pil_image = array_to_image(input_numpy_image)
+    assert isinstance(input_pil_image, Image.Image)
+
+    prompt = get_navigation_prompt(task, input_pil_image)
+
+
+    if model is None:
+        raise ValueError("Model not loaded")
+    
+    navigation_str = model.generate(prompt, max_new_tokens=500)
+    print(f"Navigation string: {navigation_str}")
+    navigation_str = navigation_str.strip()
+
+    # Parse actions from the response
+    actions = parse_actions_from_response(navigation_str)
+    
+    # Extract coordinates from all actions
+    all_coordinates = []
+    for action_code in actions:
+        coordinates = extract_coordinates_from_action(action_code)
+        all_coordinates.extend(coordinates)
+    
+    # Create localized image if there are coordinates
+    localized_image = None
+    if all_coordinates:
+        localized_image = create_localized_image(input_pil_image, all_coordinates)
+        print(f"Found {len(all_coordinates)} localization actions")
+    
+    return navigation_str, localized_image
+
+
+# --- Load Example Data ---
+example_1_image: str = "./assets/google.png"
+example_1_image = Image.open(example_1_image)
+example_1_task = "Search for the name of the current UK Prime Minister."
+
+example_2_image: str = "./assets/huggingface.png"
+example_2_image = Image.open(example_2_image)
+example_2_task = "Find the most trending model."
+
+
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown(f"<h1 style='text-align: center;'>{title}</h1>")
+    # gr.Markdown(description)
+
+    with gr.Row():
+        input_image_component = gr.Image(label="UI Image", height=500)
+    with gr.Row():
+        with gr.Column():
+            task_component = gr.Textbox(
+                label="task",
+                placeholder="e.g., Search for the name of the current UK Prime Minister.",
+                info="Type the task you want the model to complete.",
+            )
+            submit_button = gr.Button("Call Agent", variant="primary")
+
+        with gr.Column():
+            output_coords_component = gr.Textbox(label="Agent Output", lines=10)
+
+    submit_button.click(navigate, [input_image_component, task_component], [output_coords_component, input_image_component])
+
+    gr.Examples(
+        examples=[[example_1_image, example_1_task], [example_2_image, example_2_task]],
+        inputs=[input_image_component, task_component],
+        outputs=[output_coords_component, input_image_component],
+        fn=navigate,
+        cache_examples=True,
+    )
+
+demo.queue(api_open=False)
+demo.launch(debug=True, share=True)