Update app.py

app.py

@@ -1,11 +1,11 @@
 #from __future__ import annotations
 from selenium import webdriver
-
 import gradio as gr
 import uuid
 import re
 from PIL import Image
 from io import BytesIO
+from pathlib import Path
 from selenium.common.exceptions import WebDriverException
 from selenium.webdriver.common.keys import Keys
 from selenium.webdriver.common.by import By
@@ -13,54 +13,90 @@ from selenium.webdriver.common.by import By
 
 driver_type = 'chromedriver'
 driver=False
-def run_script(url: str, height: int, width: int):
+
+def get_concat_h_cut(in1, in2):
+    im1=Image.open(in1)
+    im2=Image.open(in2)
+    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(in1, in2):
+    im1=Image.open(in1)
+    im2=Image.open(in2)    
+    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, check_b,check_h):
+    mes_box=[]
+    out_box=[]
     uid=uuid.uuid4()
+    out=None
     is_url=True
     if is_url:
         options = webdriver.ChromeOptions()
         options.add_argument('--headless')
         options.add_argument('--no-sandbox')
         options.add_argument('--disable-dev-shm-usage')
-
+        mes='<center>operation success'
         try:
             driver = webdriver.Chrome(options=options)
             driver.get(url)
-            #html=driver.page_source
+            html=driver.page_source
+            #print (driver.html)
             driver.implicitly_wait(30)
             driver.set_window_size(int(width), int(height))
+            #page = driver.find_element(By.TAG_NAME, "html")
+            #driver.execute_script("arguments[0].style.background = 'blue';",page)
+            main_head = driver.find_element(By.CLASS_NAME, "main-content-header")
+            head_shot = main_head.screenshot(f'head-{uid}-tmp.png')
             obj = driver.find_element(By.CLASS_NAME, "main")            
-            #obj = driver.find_element(By.ID, "app-container")            
-            #obj = driver.find_element(By.ID, "conversation-feed")            
-            #obj = driver.find_element_by_id("in_html")
-            ######################
-            #inp = driver.find_element("id","input")
-            '''
-            find_element(By.ID, "id")
-            find_element(By.NAME, "name")
-            find_element(By.XPATH, "xpath")
-            find_element(By.LINK_TEXT, "link text")
-            find_element(By.PARTIAL_LINK_TEXT, "partial link text")
-            find_element(By.TAG_NAME, "tag name")
-            find_element(By.CLASS_NAME, "class name")
-            find_element(By.CSS_SELECTOR, "css selector")       
-            '''
-            #inp.send_keys("test")
-            ######################
-            cookie_jar = []
-            #cookie_jar.append(driver.get_cookies())
-            #print(cookie_jar)
+            messages = driver.find_elements(By.CLASS_NAME, "message")
+            for i,ea in enumerate(check_b):
+                try:    
+                    ea = int(ea)
+                    ea = (ea*2)-2
+                    print (ea)
+                    messages[ea].screenshot(f'{i}-{uid}-tmp.png')
+                    messages[ea+1].screenshot(f'{i+1}-{uid}-tmp.png')
+                    im_roll = get_concat_v_cut(f'{i}-{uid}-tmp.png',f'{i+1}-{uid}-tmp.png')
+                    im_roll.save(f'comb{i}-{uid}-tmp.png')
+                    out_box.append(f'comb{i}-{uid}-tmp.png')
+                except Exception:
+                    mes="<center>Some blocks returned an error"
+            if out_box:
+                if len(out_box)>1:
+                    im_roll = get_concat_v_cut(f'{out_box[0]}',f'{out_box[1]}')
+                    im_roll.save(f'comb-{uid}-tmp.png')                
+                    for i in range(2,len(out_box)):
+                        im_roll = get_concat_v_cut(f'comb-{uid}-tmp.png',f'{out_box[i]}')
+                        im_roll.save(f'comb-{uid}-tmp.png')
+                    out = f'comb-{uid}-tmp.png'
+                else:
+                    tmp_im = Image.open(out_box[0])
+                    tmp_im.save(f'comb-{uid}-tmp.png')
+                    out = f'comb-{uid}-tmp.png'
+                if check_h:
+                    head_paste = get_concat_v_cut(f'head-{uid}-tmp.png',f'comb-{uid}-tmp.png')
+                    head_paste.save(f'comb-{uid}-tmp.png')
+                    out = f'comb-{uid}-tmp.png'
             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}',out_box,out
         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'), mes,out_box,out
     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.',out_box,out
 
 with gr.Blocks() as app:
     
@@ -70,10 +106,17 @@ with gr.Blocks() as app:
             btn= gr.Button()
             
         with gr.Column():
+            check_h=gr.Checkbox(label="Show Header", value=True)
+            check_b=gr.CheckboxGroup(label="Chatblocks", choices=[1,2,3,4,5,6,7,8,9,10], value=[1])
             with gr.Row():
                 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():
+        with gr.Column():
+            out=gr.Image()
+            outgal=gr.Gallery()
+        outim = gr.Image()
+
+    btn.click(run_script,[inp,height,width,check_b,check_h],[outim,message,outgal,out])
 app.launch()