npv2k1 commited on
Commit
256f37c
·
1 Parent(s): 1aa5770

Enhance Dockerfile and screenshot functionality: create cache directories, set permissions, and initialize Chrome service for improved screenshot capture.

Browse files
Files changed (2) hide show
  1. Dockerfile +8 -2
  2. src/modules/apps/__init__.py +12 -5
Dockerfile CHANGED
@@ -34,10 +34,16 @@ ENTRYPOINT []
34
  # Run setup.py
35
  RUN python setup.py
36
 
37
- # Create cache directory and set permissions
38
  RUN mkdir -p /.cache/selenium && \
 
39
  chown -R appuser:appuser /.cache/selenium && \
40
- chmod 755 /.cache/selenium
 
 
 
 
 
41
 
42
  # Set user
43
  USER appuser
 
34
  # Run setup.py
35
  RUN python setup.py
36
 
37
+ # Create cache directories and set permissions
38
  RUN mkdir -p /.cache/selenium && \
39
+ mkdir -p /tmp/screenshots && \
40
  chown -R appuser:appuser /.cache/selenium && \
41
+ chown -R appuser:appuser /tmp/screenshots && \
42
+ chmod 755 /.cache/selenium && \
43
+ chmod 755 /tmp/screenshots
44
+
45
+ # Set temporary directory environment variable
46
+ ENV TMPDIR=/tmp/screenshots
47
 
48
  # Set user
49
  USER appuser
src/modules/apps/__init__.py CHANGED
@@ -3,7 +3,9 @@ import tempfile
3
  import os
4
  from selenium import webdriver
5
  from selenium.webdriver.chrome.options import Options
 
6
  print("Gradio app loaded.")
 
7
  def capture_page(url: str, output_file: str = "screenshot.png"):
8
  """
9
  Captures a screenshot of the given webpage.
@@ -14,8 +16,10 @@ def capture_page(url: str, output_file: str = "screenshot.png"):
14
  options = Options()
15
  options.add_argument("--headless") # Run in headless mode
16
  options.add_argument("--window-size=1920,1080") # Set window size
17
-
18
- driver = webdriver.Chrome(options=options)
 
 
19
 
20
  try:
21
  driver.get(url)
@@ -27,9 +31,12 @@ def capture_page(url: str, output_file: str = "screenshot.png"):
27
  def capture_and_show(url: str):
28
  """Capture webpage and return the image"""
29
  try:
30
- # Create temporary file for the screenshot
31
- with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as tmp:
32
- temp_path = tmp.name
 
 
 
33
 
34
  # Capture the webpage
35
  capture_page(url, temp_path)
 
3
  import os
4
  from selenium import webdriver
5
  from selenium.webdriver.chrome.options import Options
6
+ from selenium.webdriver.chrome.service import Service
7
  print("Gradio app loaded.")
8
+
9
  def capture_page(url: str, output_file: str = "screenshot.png"):
10
  """
11
  Captures a screenshot of the given webpage.
 
16
  options = Options()
17
  options.add_argument("--headless") # Run in headless mode
18
  options.add_argument("--window-size=1920,1080") # Set window size
19
+
20
+ # Initialize Chrome service
21
+ service = Service()
22
+ driver = webdriver.Chrome(service=service, options=options)
23
 
24
  try:
25
  driver.get(url)
 
31
  def capture_and_show(url: str):
32
  """Capture webpage and return the image"""
33
  try:
34
+ # Get the temporary directory path (defaulting to /tmp if TMPDIR is not set)
35
+ temp_dir = os.getenv('TMPDIR', '/tmp')
36
+ os.makedirs(temp_dir, exist_ok=True)
37
+
38
+ # Create temporary file in the specified directory
39
+ temp_path = os.path.join(temp_dir, f"screenshot_{os.urandom(8).hex()}.png")
40
 
41
  # Capture the webpage
42
  capture_page(url, temp_path)