Update app.py

app.py

@@ -1,21 +1,35 @@
 import gradio as gr
-from extract import take_webdata
+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')
 
-def visualize(url: str):
     try:
-        img, title = take_webdata(url)
-        return img, f"✅ Title: {title}"
-    except Exception as e:
-        return None, f"❌ Error: {e}"
+        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()
 
-with gr.Blocks() as demo:
-    gr.Markdown("## 🌐 Website Screenshot + Title Extractor (Playwright-powered)")
-    url_box = gr.Textbox(label="Enter URL (e.g., https://www.amazon.in)")
-    btn = gr.Button("Proceed")
+    return Image.open(BytesIO(screenshot))
 
-    with gr.Row():
-        img_out = gr.Image(type="pil", label="Screenshot")
-        text_out = gr.Textbox(label="Page Title / Error")
+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."
+)
 
-    btn.click(fn=visualize, inputs=[url_box], outputs=[img_out, text_out])
-demo.launch()
+iface.launch()