Update app.py

app.py

@@ -1,43 +1,50 @@
 import os
+import time
 import streamlit as st
 from selenium import webdriver
 from selenium.webdriver.chrome.options import Options
-from PIL import Image
+import tempfile
 
 # Function to capture a screenshot of a webpage
-def capture_screenshot(url, output_path='screenshot.png'):
+def capture_screenshot(url, output_path):
     chrome_options = Options()
-    chrome_options.add_argument('--headless')  # Run Chrome in headless mode
-    chrome_options.add_argument('--disable-gpu')  # Disable GPU acceleration
-    chrome_options.add_argument('--no-sandbox')  # Disable sandbox for better compatibility
-    chrome_options.add_argument('--window-size=1920,1080')  # Set window size for consistent screenshots
+    chrome_options.add_argument('--headless')  
+    chrome_options.add_argument('--disable-gpu')  
+    chrome_options.add_argument('--no-sandbox')  
+    chrome_options.add_argument('--window-size=1920,1080')  
 
     driver = webdriver.Chrome(options=chrome_options)
     try:
-        driver.get(url)  # Open the URL
-        driver.save_screenshot(output_path)  # Save screenshot to file
+        driver.get(url)
+        driver.save_screenshot(output_path)
     finally:
-        driver.quit()  # Ensure the driver is closed even if an error occurs
+        driver.quit()
 
-# Streamlit app layout
-st.title("Web Screenshot Viewer")
+# Streamlit UI
+st.title("Timed Web Screenshot Viewer")
 
-# Input field for URL
 url = st.text_input("Enter a website URL:", "https://example.com")
+wait_time = st.number_input("Time to wait before taking screenshot (in seconds):", min_value=1, max_value=30, value=5, step=1)
 
-# Button to trigger the screenshot process
 if st.button("Generate Screenshot"):
+    screenshot_path = None  # Initialize to ensure it's defined
     try:
-        screenshot_file = "screenshot.png"
-        st.write("Capturing website screenshot...")
-        capture_screenshot(url.strip('"'))  # Strip any extra quotes from the URL
+        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp_file:
+            screenshot_path = tmp_file.name
 
-        # Display the screenshot in the Streamlit app
-        st.image(screenshot_file, caption='Screenshot of the webpage', use_column_width=True)
+        # Show a countdown
+        with st.empty():
+            for i in range(wait_time, 0, -1):
+                st.info(f"Waiting... {i} seconds left")
+                time.sleep(1)
+
+        st.write("Capturing screenshot...")
+        capture_screenshot(url.strip('"'), screenshot_path)
+
+        st.image(screenshot_path, caption='Screenshot of the webpage', use_column_width=True)
 
     except Exception as e:
         st.error(f"An error occurred: {e}")
     finally:
-        # Clean up the screenshot file if it exists
-        if os.path.exists(screenshot_file):
-            os.remove(screenshot_file)
+        if screenshot_path and os.path.exists(screenshot_path):
+            os.remove(screenshot_path)