Spaces:
Runtime error
Runtime error
File size: 7,144 Bytes
d73c58e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 |
import subprocess
import base64
from pathlib import Path
from PIL import ImageGrab
from uuid import uuid4
from screeninfo import get_monitors
import platform
if platform.system() == "Darwin":
import Quartz # uncomment this line if you are on macOS
from PIL import ImageGrab
from functools import partial
from .base import BaseAnthropicTool, ToolError, ToolResult
OUTPUT_DIR = "./tmp/outputs"
def get_screenshot(selected_screen: int = 0, resize: bool = True, target_width: int = 1920, target_height: int = 1080):
# print(f"get_screenshot selected_screen: {selected_screen}")
# Get screen width and height using Windows command
display_num = None
offset_x = 0
offset_y = 0
selected_screen = selected_screen
width, height = _get_screen_size()
"""Take a screenshot of the current screen and return a ToolResult with the base64 encoded image."""
output_dir = Path(OUTPUT_DIR)
output_dir.mkdir(parents=True, exist_ok=True)
path = output_dir / f"screenshot_{uuid4().hex}.png"
ImageGrab.grab = partial(ImageGrab.grab, all_screens=True)
# Detect platform
system = platform.system()
if system == "Windows":
# Windows: Use screeninfo to get monitor details
screens = get_monitors()
# Sort screens by x position to arrange from left to right
sorted_screens = sorted(screens, key=lambda s: s.x)
if selected_screen < 0 or selected_screen >= len(screens):
raise IndexError("Invalid screen index.")
screen = sorted_screens[selected_screen]
bbox = (screen.x, screen.y, screen.x + screen.width, screen.y + screen.height)
elif system == "Darwin": # macOS
# macOS: Use Quartz to get monitor details
max_displays = 32 # Maximum number of displays to handle
active_displays = Quartz.CGGetActiveDisplayList(max_displays, None, None)[1]
# Get the display bounds (resolution) for each active display
screens = []
for display_id in active_displays:
bounds = Quartz.CGDisplayBounds(display_id)
screens.append({
'id': display_id,
'x': int(bounds.origin.x),
'y': int(bounds.origin.y),
'width': int(bounds.size.width),
'height': int(bounds.size.height),
'is_primary': Quartz.CGDisplayIsMain(display_id) # Check if this is the primary display
})
# Sort screens by x position to arrange from left to right
sorted_screens = sorted(screens, key=lambda s: s['x'])
# print(f"Darwin sorted_screens: {sorted_screens}")
if selected_screen < 0 or selected_screen >= len(screens):
raise IndexError("Invalid screen index.")
screen = sorted_screens[selected_screen]
bbox = (screen['x'], screen['y'], screen['x'] + screen['width'], screen['y'] + screen['height'])
else: # Linux or other OS
cmd = "xrandr | grep ' primary' | awk '{print $4}'"
try:
output = subprocess.check_output(cmd, shell=True).decode()
resolution = output.strip().split()[0]
width, height = map(int, resolution.split('x'))
bbox = (0, 0, width, height) # Assuming single primary screen for simplicity
except subprocess.CalledProcessError:
raise RuntimeError("Failed to get screen resolution on Linux.")
# Take screenshot using the bounding box
screenshot = ImageGrab.grab(bbox=bbox)
# Set offsets (for potential future use)
offset_x = screen['x'] if system == "Darwin" else screen.x
offset_y = screen['y'] if system == "Darwin" else screen.y
# # Resize if
if resize:
screenshot = screenshot.resize((target_width, target_height))
# Save the screenshot
screenshot.save(str(path))
if path.exists():
# Return a ToolResult instance instead of a dictionary
return screenshot, path
raise ToolError(f"Failed to take screenshot: {path} does not exist.")
def _get_screen_size(selected_screen: int = 0):
if platform.system() == "Windows":
# Use screeninfo to get primary monitor on Windows
screens = get_monitors()
# Sort screens by x position to arrange from left to right
sorted_screens = sorted(screens, key=lambda s: s.x)
if selected_screen is None:
primary_monitor = next((m for m in get_monitors() if m.is_primary), None)
return primary_monitor.width, primary_monitor.height
elif selected_screen < 0 or selected_screen >= len(screens):
raise IndexError("Invalid screen index.")
else:
screen = sorted_screens[selected_screen]
return screen.width, screen.height
elif platform.system() == "Darwin":
# macOS part using Quartz to get screen information
max_displays = 32 # Maximum number of displays to handle
active_displays = Quartz.CGGetActiveDisplayList(max_displays, None, None)[1]
# Get the display bounds (resolution) for each active display
screens = []
for display_id in active_displays:
bounds = Quartz.CGDisplayBounds(display_id)
screens.append({
'id': display_id,
'x': int(bounds.origin.x),
'y': int(bounds.origin.y),
'width': int(bounds.size.width),
'height': int(bounds.size.height),
'is_primary': Quartz.CGDisplayIsMain(display_id) # Check if this is the primary display
})
# Sort screens by x position to arrange from left to right
sorted_screens = sorted(screens, key=lambda s: s['x'])
if selected_screen is None:
# Find the primary monitor
primary_monitor = next((screen for screen in screens if screen['is_primary']), None)
if primary_monitor:
return primary_monitor['width'], primary_monitor['height']
else:
raise RuntimeError("No primary monitor found.")
elif selected_screen < 0 or selected_screen >= len(screens):
raise IndexError("Invalid screen index.")
else:
# Return the resolution of the selected screen
screen = sorted_screens[selected_screen]
return screen['width'], screen['height']
else: # Linux or other OS
cmd = "xrandr | grep ' primary' | awk '{print $4}'"
try:
output = subprocess.check_output(cmd, shell=True).decode()
resolution = output.strip().split()[0]
width, height = map(int, resolution.split('x'))
return width, height
except subprocess.CalledProcessError:
raise RuntimeError("Failed to get screen resolution on Linux.")
|