Update app.py

app.py

@@ -1,6 +1,6 @@
 #from __future__ import annotations
 from selenium import webdriver
-
+from PIL import Image
 import gradio as gr
 import uuid
 import re
@@ -13,7 +13,26 @@ from selenium.webdriver.common.by import By
 
 driver_type = 'chromedriver'
 driver=False
+
+def get_concat_h_cut(im1, im2):
+    dst = Image.new('RGB', (im1.width + im2.width,
+                            min(im1.height, im2.height)))
+    dst.paste(im1, (0, 0))
+    dst.paste(im2, (im1.width, 0))
+    return dst
+
+
+def get_concat_v_cut(im1, im2):
+    dst = Image.new(
+        'RGB', (min(im1.width, im2.width), im1.height + im2.height))
+    dst.paste(im1, (0, 0))
+    dst.paste(im2, (0, im1.height))
+    return dst
+
+
+
 def run_script(url: str, height: int, width: int):
+    mes_box=[]
     uid=uuid.uuid4()
     is_url=True
     if is_url:
@@ -38,8 +57,10 @@ def run_script(url: str, height: int, width: int):
             obj = driver.find_element(By.CLASS_NAME, "main")            
             messages = driver.find_elements(By.CLASS_NAME, "message")
             print (len(messages))
-            for mes in messages:
-                print (messages.text)
+            for i,mes in enumerate(messages):
+                print (mes.text)
+                mes_shot = obj.screenshot(f'{i}{uid}-tmp.png')
+                mes_box.append(mes_shot)
             #obj = driver.find_element(By.ID, "app-container")            
             #obj = driver.find_element(By.ID, "conversation-feed")            
             #obj = driver.find_element_by_id("in_html")
@@ -63,15 +84,15 @@ def run_script(url: str, height: int, width: int):
             screenshot = obj.screenshot(f'{uid}-tmp.png')
             #screenshot = obj.get_screenshot_as_png()
         except WebDriverException as e:
-            return [Image.new('RGB', (1, 1)), f'<center>{e}']
+            return [Image.new('RGB', (1, 1)), f'<center>{e}',mes_box]
         finally:
             if driver:
                 driver.quit()
 
         #return [Image.open(BytesIO(screenshot)), 'operation success.',cookie_jar,html]
-        return [Image.open(f'{uid}-tmp.png'), '<center>operation success.']
+        return [Image.open(f'{uid}-tmp.png'), '<center>operation success.',mes_box]
     else:
-        return [None, '<center>Please enter a valid URL of a website/host.']
+        return [None, '<center>Please enter a valid URL of a website/host.',mes_box]
 
 with gr.Blocks() as app:
     
@@ -85,6 +106,8 @@ with gr.Blocks() as app:
                 height=gr.Number(label="Height", value=4096)
                 width=gr.Number(label="Width",value=800)
     message=gr.HTML('<center>Enter URL')
-    outim = gr.Image()
-    btn.click(run_script,[inp,height,width],[outim,message])
+    with gr.Row():
+        outim = gr.Image()
+        outgal=gr.Gallery()
+    btn.click(run_script,[inp,height,width],[outim,message,outgal])
 app.launch()