npv2k1 commited on
Commit
f8136bb
·
1 Parent(s): b68547c

Enhance screenshot capture functionality: improve Chrome initialization, add dynamic content loading, and ensure temporary directory permissions.

Browse files
Files changed (1) hide show
  1. 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
- driver.get(url)
41
- # Add a small delay to ensure page loads completely
42
- driver.implicitly_wait(5)
43
- driver.save_screenshot(output_file)
44
- print(f"Screenshot saved: {output_file}")
45
- finally:
46
- driver.quit()
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
- # Capture the webpage
59
- capture_page(url, temp_path)
60
-
61
- # Return the image path
62
- return temp_path
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
63
  except Exception as e:
64
- print(f"Error in capture_and_show: {str(e)}") # Add detailed logging
65
- return None # Return None instead of error string to handle gracefully
 
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"""