Spaces:
Sleeping
Sleeping
Enhance screenshot capture functionality: improve Chrome initialization, add dynamic content loading, and ensure temporary directory permissions.
Browse files- src/modules/apps/__init__.py +98 -26
src/modules/apps/__init__.py
CHANGED
|
@@ -4,7 +4,11 @@ import os
|
|
| 4 |
from selenium import webdriver
|
| 5 |
from selenium.webdriver.chrome.options import Options
|
| 6 |
from selenium.webdriver.chrome.service import Service
|
|
|
|
|
|
|
| 7 |
from subprocess import PIPE, STDOUT
|
|
|
|
|
|
|
| 8 |
print("Gradio app loaded.")
|
| 9 |
|
| 10 |
def capture_page(url: str, output_file: str = "screenshot.png"):
|
|
@@ -15,54 +19,122 @@ def capture_page(url: str, output_file: str = "screenshot.png"):
|
|
| 15 |
:param output_file: The filename to save the screenshot.
|
| 16 |
"""
|
| 17 |
options = Options()
|
|
|
|
| 18 |
options.add_argument('--headless')
|
| 19 |
options.add_argument('--no-sandbox') # Required in Docker
|
| 20 |
options.add_argument('--disable-dev-shm-usage') # Required in Docker
|
|
|
|
|
|
|
| 21 |
options.add_argument('--disable-gpu') # Required in Docker
|
| 22 |
options.add_argument('--disable-software-rasterizer')
|
| 23 |
-
options.add_argument('--window-size=1920,1080')
|
| 24 |
options.add_argument('--disable-extensions')
|
| 25 |
options.add_argument('--disable-infobars')
|
| 26 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 27 |
# Set up Chrome service with explicit path to chromedriver and logging
|
| 28 |
service = Service(
|
| 29 |
executable_path='/usr/local/bin/chromedriver',
|
| 30 |
-
log_output=PIPE # Redirect logs to pipe
|
| 31 |
-
|
| 32 |
-
|
| 33 |
-
# Initialize Chrome with the service and options
|
| 34 |
-
driver = webdriver.Chrome(
|
| 35 |
-
service=service,
|
| 36 |
-
options=options
|
| 37 |
)
|
| 38 |
|
| 39 |
try:
|
| 40 |
-
|
| 41 |
-
|
| 42 |
-
|
| 43 |
-
|
| 44 |
-
|
| 45 |
-
|
| 46 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 47 |
|
| 48 |
def capture_and_show(url: str):
|
| 49 |
"""Capture webpage and return the image"""
|
| 50 |
try:
|
| 51 |
# Get the temporary directory path (defaulting to /tmp if TMPDIR is not set)
|
| 52 |
temp_dir = os.getenv('TMPDIR', '/tmp')
|
| 53 |
-
os.makedirs(temp_dir, exist_ok=True)
|
| 54 |
-
|
| 55 |
-
# Create temporary file in the specified directory
|
| 56 |
-
temp_path = os.path.join(temp_dir, f"screenshot_{os.urandom(8).hex()}.png")
|
| 57 |
|
| 58 |
-
|
| 59 |
-
|
| 60 |
-
|
| 61 |
-
|
| 62 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 63 |
except Exception as e:
|
| 64 |
-
print(f"Error in capture_and_show: {str(e)}")
|
| 65 |
-
|
|
|
|
| 66 |
|
| 67 |
def create_gradio_app():
|
| 68 |
"""Create the main Gradio application with all components"""
|
|
|
|
| 4 |
from selenium import webdriver
|
| 5 |
from selenium.webdriver.chrome.options import Options
|
| 6 |
from selenium.webdriver.chrome.service import Service
|
| 7 |
+
from selenium.webdriver.support.ui import WebDriverWait
|
| 8 |
+
from selenium.webdriver.support import expected_conditions as EC
|
| 9 |
from subprocess import PIPE, STDOUT
|
| 10 |
+
import traceback
|
| 11 |
+
|
| 12 |
print("Gradio app loaded.")
|
| 13 |
|
| 14 |
def capture_page(url: str, output_file: str = "screenshot.png"):
|
|
|
|
| 19 |
:param output_file: The filename to save the screenshot.
|
| 20 |
"""
|
| 21 |
options = Options()
|
| 22 |
+
# Basic options
|
| 23 |
options.add_argument('--headless')
|
| 24 |
options.add_argument('--no-sandbox') # Required in Docker
|
| 25 |
options.add_argument('--disable-dev-shm-usage') # Required in Docker
|
| 26 |
+
|
| 27 |
+
# Performance and stability options
|
| 28 |
options.add_argument('--disable-gpu') # Required in Docker
|
| 29 |
options.add_argument('--disable-software-rasterizer')
|
|
|
|
| 30 |
options.add_argument('--disable-extensions')
|
| 31 |
options.add_argument('--disable-infobars')
|
| 32 |
|
| 33 |
+
# Resource configuration
|
| 34 |
+
options.add_argument('--window-size=1920,1080')
|
| 35 |
+
options.add_argument('--remote-debugging-port=9222') # Fix DevTools port issue
|
| 36 |
+
options.add_argument('--disable-features=site-per-process') # Reduce memory usage
|
| 37 |
+
options.add_argument('--memory-pressure-off') # Prevent memory-related crashes
|
| 38 |
+
|
| 39 |
+
# Additional stability options
|
| 40 |
+
options.add_argument('--ignore-certificate-errors')
|
| 41 |
+
options.add_argument('--allow-insecure-localhost')
|
| 42 |
+
options.add_argument('--disable-setuid-sandbox')
|
| 43 |
+
options.add_argument('--disable-web-security')
|
| 44 |
+
|
| 45 |
# Set up Chrome service with explicit path to chromedriver and logging
|
| 46 |
service = Service(
|
| 47 |
executable_path='/usr/local/bin/chromedriver',
|
| 48 |
+
log_output=PIPE, # Redirect logs to pipe
|
| 49 |
+
service_args=['--verbose'] # Enable verbose logging
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 50 |
)
|
| 51 |
|
| 52 |
try:
|
| 53 |
+
print("Initializing Chrome...")
|
| 54 |
+
driver = webdriver.Chrome(
|
| 55 |
+
service=service,
|
| 56 |
+
options=options
|
| 57 |
+
)
|
| 58 |
+
|
| 59 |
+
print("Chrome initialized successfully")
|
| 60 |
+
|
| 61 |
+
try:
|
| 62 |
+
print(f"Navigating to URL: {url}")
|
| 63 |
+
driver.get(url)
|
| 64 |
+
|
| 65 |
+
# Wait for page load
|
| 66 |
+
print("Waiting for page to load...")
|
| 67 |
+
driver.implicitly_wait(10) # Increased wait time
|
| 68 |
+
|
| 69 |
+
# Additional wait for dynamic content
|
| 70 |
+
from selenium.webdriver.support.ui import WebDriverWait
|
| 71 |
+
from selenium.webdriver.support import expected_conditions as EC
|
| 72 |
+
WebDriverWait(driver, 10).until(
|
| 73 |
+
lambda d: d.execute_script('return document.readyState') == 'complete'
|
| 74 |
+
)
|
| 75 |
+
|
| 76 |
+
print("Taking screenshot...")
|
| 77 |
+
driver.save_screenshot(output_file)
|
| 78 |
+
print(f"Screenshot saved: {output_file}")
|
| 79 |
+
return True
|
| 80 |
+
|
| 81 |
+
except Exception as e:
|
| 82 |
+
print(f"Error during page capture: {str(e)}")
|
| 83 |
+
raise
|
| 84 |
+
finally:
|
| 85 |
+
print("Closing Chrome...")
|
| 86 |
+
driver.quit()
|
| 87 |
+
|
| 88 |
+
except Exception as e:
|
| 89 |
+
print(f"Error initializing Chrome: {str(e)}")
|
| 90 |
+
raise Exception(f"Failed to initialize Chrome: {str(e)}")
|
| 91 |
|
| 92 |
def capture_and_show(url: str):
|
| 93 |
"""Capture webpage and return the image"""
|
| 94 |
try:
|
| 95 |
# Get the temporary directory path (defaulting to /tmp if TMPDIR is not set)
|
| 96 |
temp_dir = os.getenv('TMPDIR', '/tmp')
|
|
|
|
|
|
|
|
|
|
|
|
|
| 97 |
|
| 98 |
+
try:
|
| 99 |
+
# Ensure temp directory exists and has correct permissions
|
| 100 |
+
os.makedirs(temp_dir, mode=0o777, exist_ok=True)
|
| 101 |
+
print(f"Using temp directory: {temp_dir}")
|
| 102 |
+
|
| 103 |
+
# Verify directory is writable
|
| 104 |
+
if not os.access(temp_dir, os.W_OK):
|
| 105 |
+
print(f"Warning: Temp directory {temp_dir} is not writable")
|
| 106 |
+
# Try to create a user-specific temp directory instead
|
| 107 |
+
temp_dir = os.path.join('/tmp', f'chrome_screenshots_{os.getuid()}')
|
| 108 |
+
os.makedirs(temp_dir, mode=0o777, exist_ok=True)
|
| 109 |
+
print(f"Created user-specific temp directory: {temp_dir}")
|
| 110 |
+
|
| 111 |
+
# Create temporary file in the specified directory
|
| 112 |
+
temp_path = os.path.join(temp_dir, f"screenshot_{os.urandom(8).hex()}.png")
|
| 113 |
+
print(f"Temp file path: {temp_path}")
|
| 114 |
+
|
| 115 |
+
# Capture the webpage
|
| 116 |
+
success = capture_page(url, temp_path)
|
| 117 |
+
if not success:
|
| 118 |
+
print("Screenshot capture returned False")
|
| 119 |
+
return None
|
| 120 |
+
|
| 121 |
+
# Verify file was created
|
| 122 |
+
if not os.path.exists(temp_path):
|
| 123 |
+
print("Screenshot file was not created")
|
| 124 |
+
return None
|
| 125 |
+
|
| 126 |
+
print("Screenshot captured successfully")
|
| 127 |
+
return temp_path
|
| 128 |
+
|
| 129 |
+
except OSError as e:
|
| 130 |
+
print(f"OS Error: {str(e)}")
|
| 131 |
+
print(f"Stack trace: {traceback.format_exc()}")
|
| 132 |
+
return None
|
| 133 |
+
|
| 134 |
except Exception as e:
|
| 135 |
+
print(f"Error in capture_and_show: {str(e)}")
|
| 136 |
+
print(f"Stack trace: {traceback.format_exc()}")
|
| 137 |
+
return None
|
| 138 |
|
| 139 |
def create_gradio_app():
|
| 140 |
"""Create the main Gradio application with all components"""
|