Update app.py

app.py

@@ -1,2 +1,35 @@
-import subprocess
+import gradio as gr
+from selenium import webdriver
+from selenium.common.exceptions import WebDriverException
+from PIL import Image
+from io import BytesIO
 
+def take_screenshot(url):
+    options = webdriver.ChromeOptions()
+    options.add_argument('--headless')
+    options.add_argument('--no-sandbox')
+    options.add_argument('--disable-dev-shm-usage')
+
+    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.Textbox(label="Website URL", value="https://kargaranamir.github.io"),
+    outputs=gr.Image(type="pil", label="Screenshot", height=360, width=540),  # Adjust the image size here
+    title="Website Screenshot",
+    description="Take a screenshot of a website."
+)
+
+iface.launch()