Layout Detector MCP

"""MCP server for layout detection.""" import asyncio import json from pathlib import Path from mcp.server import Server from mcp.server.stdio import stdio_server from mcp.types import Tool, TextContent from .template_match import find_all_assets, get_screenshot_dimensions from .layout_analysis import analyze_layout # Create the MCP server server = Server("layout-detector") @server.list_tools() async def list_tools() -> list[Tool]: """List available tools.""" return [ Tool( name="find_assets_in_screenshot", description=( "Find known image assets within a screenshot. Uses template matching " "to locate each asset and return its position (x, y, width, height). " "Useful for determining where specific images appear in a webpage screenshot." ), inputSchema={ "type": "object", "properties": { "screenshot_path": { "type": "string", "description": "Absolute path to the screenshot image file", }, "asset_paths": { "type": "array", "items": {"type": "string"}, "description": "List of absolute paths to asset images to find", }, "threshold": { "type": "number", "description": "Match confidence threshold (0-1). Default 0.8", "default": 0.8, }, }, "required": ["screenshot_path", "asset_paths"], }, ), Tool( name="analyze_layout", description=( "Analyze the layout of assets in a screenshot. Finds all assets, " "identifies the center/hero element, calculates relative positions " "(angle and distance from center), and detects the layout pattern " "(radial, grid, or freeform). Returns structured data for rebuilding " "the layout with semantic CSS." ), inputSchema={ "type": "object", "properties": { "screenshot_path": { "type": "string", "description": "Absolute path to the screenshot image file", }, "asset_paths": { "type": "array", "items": {"type": "string"}, "description": "List of absolute paths to asset images to find", }, "threshold": { "type": "number", "description": "Match confidence threshold (0-1). Default 0.8", "default": 0.8, }, }, "required": ["screenshot_path", "asset_paths"], }, ), Tool( name="get_screenshot_info", description=( "Get basic information about a screenshot image, including its " "dimensions (width and height in pixels)." ), inputSchema={ "type": "object", "properties": { "screenshot_path": { "type": "string", "description": "Absolute path to the screenshot image file", }, }, "required": ["screenshot_path"], }, ), ] @server.call_tool() async def call_tool(name: str, arguments: dict) -> list[TextContent]: """Handle tool calls.""" if name == "find_assets_in_screenshot": screenshot_path = arguments["screenshot_path"] asset_paths = arguments["asset_paths"] threshold = arguments.get("threshold", 0.8) # Validate paths if not Path(screenshot_path).exists(): return [TextContent( type="text", text=json.dumps({"error": f"Screenshot not found: {screenshot_path}"}) )] missing_assets = [p for p in asset_paths if not Path(p).exists()] if missing_assets: return [TextContent( type="text", text=json.dumps({"error": f"Assets not found: {missing_assets}"}) )] # Find assets matches = find_all_assets(screenshot_path, asset_paths, threshold) result = { "found": len(matches), "total_assets": len(asset_paths), "matches": [m.to_dict() for m in matches], } return [TextContent(type="text", text=json.dumps(result, indent=2))] elif name == "analyze_layout": screenshot_path = arguments["screenshot_path"] asset_paths = arguments["asset_paths"] threshold = arguments.get("threshold", 0.8) # Validate paths if not Path(screenshot_path).exists(): return [TextContent( type="text", text=json.dumps({"error": f"Screenshot not found: {screenshot_path}"}) )] missing_assets = [p for p in asset_paths if not Path(p).exists()] if missing_assets: return [TextContent( type="text", text=json.dumps({"error": f"Assets not found: {missing_assets}"}) )] # Analyze layout analysis = analyze_layout(screenshot_path, asset_paths, threshold) return [TextContent(type="text", text=json.dumps(analysis.to_dict(), indent=2))] elif name == "get_screenshot_info": screenshot_path = arguments["screenshot_path"] if not Path(screenshot_path).exists(): return [TextContent( type="text", text=json.dumps({"error": f"Screenshot not found: {screenshot_path}"}) )] width, height = get_screenshot_dimensions(screenshot_path) result = { "path": screenshot_path, "width": width, "height": height, } return [TextContent(type="text", text=json.dumps(result, indent=2))] else: return [TextContent( type="text", text=json.dumps({"error": f"Unknown tool: {name}"}) )] def main(): """Run the MCP server.""" async def run(): async with stdio_server() as (read_stream, write_stream): await server.run( read_stream, write_stream, server.create_initialization_options(), ) asyncio.run(run()) if __name__ == "__main__": main()