Update app.py

app.py

@@ -5,6 +5,11 @@ import requests
 from typing import List, Dict, Union
 import traceback
 
+from selenium import webdriver
+from selenium.common.exceptions import WebDriverException
+from PIL import Image
+from io import BytesIO
+
 HF_TOKEN = os.getenv("HF_TOKEN")
 hf_client = InferenceClient("CohereForAI/c4ai-command-r-plus-08-2024", token=HF_TOKEN)
 
@@ -86,17 +91,39 @@ def on_select(space):
     try:
         summary = summarize_space(space)
         app_content = get_app_py_content(space['id'])
+        screenshot = take_screenshot(space['url'])
         info = f"선택된 Space: {space['name']} (ID: {space['id']})\n"
         info += f"Author: {space['author']}\n"
         info += f"Likes: {space['likes']}\n"
         info += f"URL: {space['url']}\n\n"
         info += f"요약:\n{summary}"
-        return info, app_content
+        return info, app_content, screenshot
     except Exception as e:
         print(f"Error in on_select: {str(e)}")
         print(traceback.format_exc())
-        return f"오류가 발생했습니다: {str(e)}", ""
+        return f"오류가 발생했습니다: {str(e)}", "", Image.new('RGB', (1, 1))
+        
 
+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)
+        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))
+    
 def create_ui():
     spaces_list = get_most_liked_spaces()
     print(f"Type of spaces_list: {type(spaces_list)}")
@@ -133,8 +160,19 @@ def create_ui():
                 outputs=[info_output, app_py_content]
             )
 
-    return demo
+        with gr.Column(scale=1):
+            info_output = gr.Textbox(label="Space 정보 및 요약", lines=16)
+            app_py_content = gr.Code(language="python", label="app.py 내용")
+            screenshot_output = gr.Image(type="pil", label="Live 화면", height=360, width=540)
 
+        for _, button, space in space_rows:
+            button.click(
+                lambda s=space: on_select(s),
+                inputs=[],
+                outputs=[info_output, app_py_content, screenshot_output]
+            )
+
+    return demo
 if __name__ == "__main__":
     demo = create_ui()
     demo.launch()