Update app.py

app.py

@@ -4,26 +4,36 @@ from selenium import webdriver
 from selenium.webdriver.chrome.options import Options
 import ascii_magic
 
+# Function to capture a screenshot of a webpage
 def capture_screenshot(url, output_path='screenshot.png'):
     chrome_options = Options()
-    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')
+    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
 
     driver = webdriver.Chrome(options=chrome_options)
-    driver.get(url)
-    driver.save_screenshot(output_path)
-    driver.quit()
+    try:
+        driver.get(url)  # Open the URL
+        driver.save_screenshot(output_path)  # Save screenshot to file
+    finally:
+        driver.quit()  # Ensure the driver is closed even if an error occurs
 
+# Function to convert an image to ASCII art
 def convert_to_ascii(image_path):
-    output = ascii_magic.from_image_path(image_path, columns=100, mode=ascii_magic.Modes.ASCII)
-    return str(output)
+    try:
+        output = ascii_magic.from_image_path(image_path, columns=100, mode=ascii_magic.Modes.ASCII)
+        return str(output)
+    except Exception as e:
+        raise Exception(f"Failed to convert image to ASCII: {e}")
 
+# Streamlit app layout
 st.title("ASCII Web Screenshot Viewer")
 
+# Input field for URL
 url = st.text_input("Enter a website URL:", "https://example.com")
 
+# Button to trigger the screenshot and conversion process
 if st.button("Generate ASCII Screenshot"):
     try:
         screenshot_file = "screenshot.png"
@@ -33,8 +43,12 @@ if st.button("Generate ASCII Screenshot"):
         st.write("Converting to ASCII art...")
         ascii_output = convert_to_ascii(screenshot_file)
 
+        # Display ASCII art in a text area
         st.text_area("ASCII Art Output", ascii_output, height=600)
-        os.remove(screenshot_file)
 
     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)