Update app.py

app.py

@@ -1,36 +1,35 @@
 import gradio as gr
 from selenium import webdriver
-from selenium.webdriver.support.ui import WebDriverWait
-from selenium.webdriver.support import expected_conditions as EC
-from selenium.webdriver.common.by import By
+from selenium.common.exceptions import WebDriverException
+from PIL import Image
+from io import BytesIO
 
-def scrape_website(url):
-    # Create a new instance of the Chrome driver
-    driver = webdriver.Chrome()
-    
-    # Navigate to the provided URL
-    driver.get(url)
-    
-    # Wait until all processes are complete
-    wait = WebDriverWait(driver, 10)  # Maximum wait time in seconds
-    wait.until(EC.presence_of_element_located((By.TAG_NAME, 'body')))
-    
-    # Extract the page title
-    title = driver.title
-    
-    # Close the browser
-    driver.quit()
-    
-    return title
+def take_screenshot(url):
+    options = webdriver.ChromeOptions()
+    options.add_argument('--headless')
+    options.add_argument('--no-sandbox')
+    options.add_argument('--disable-dev-shm-usage')
 
-# Define the Gradio interface
-interface = gr.Interface(
-    fn=scrape_website,
-    inputs="text",
-    outputs="text",
-    title="Web Scraping Application",
-    description="Enter a URL to extract the page title."
+    try:
+        wd = webdriver.Chrome(options=options)
+        wd.set_window_size(1080, 720)  # Adjust the window size here
+        wd.get(url)
+        wd.implicitly_wait(10)
+        screenshot = wd.get_screenshot_as_png()
+    except WebDriverException as e:
+        return Image.new('RGB', (1, 1))
+    finally:
+        if wd:
+            wd.quit()
+
+    return Image.open(BytesIO(screenshot))
+
+iface = gr.Interface(
+    fn=take_screenshot,
+    inputs=gr.inputs.Textbox(label="Website URL", default="https://kargaranamir.github.io"),
+    outputs=gr.Image(type="pil", height=360, width=540),  # Adjust the image size here
+    title="Website Screenshot",
+    description="Take a screenshot of a website.",
 )
 
-# Run the Gradio app
-interface.launch()
+iface.launch()