app.py

#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
#from selenium_webdriver import WebElement

driver_type = 'chromedriver'
driver=False

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, theme):
    im1=Image.open(in1)
    im2=Image.open(in2)
    if theme=='dark':
        color=(31,41,55)
    if theme=='light':
        color=(255,255,255)
    dst = Image.new('RGB', (min(im1.width, im2.width), im1.height + im2.height+30),color=color)
    dst.paste(im1, (0, 10))
    dst.paste(im2, (0, im1.height+20))
    return dst



def run_script1(url: str, height: int, width: int, check_b,check_h):
    mes_box=[]
    out_box=[]
    uid=uuid.uuid4()
    out=None
    is_url=True
    url='https://omnibus-snap.hf.space/?chat=[(123,555),(2342,32344)]&__theme=dark'
    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.current_window_handle
            driver.get(url)
            html=driver.page_source
            #print(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")            
            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')
        except WebDriverException as e:
            return Image.new('RGB', (1, 1)), f'<center>Please enter a valid URL of a website/host.',out_box,out
        finally:
            if driver:
                driver.quit()
        return Image.open(f'{uid}-tmp.png'), mes,out_box,out
    else:
        return None, '<center>Please enter a valid URL of a website/host.',out_box,out

def run_script(chat: list, height: int, width: int, check_b,check_h,check_t,wait_time):
    mes_box=[]
    out_box=[]
    uid=uuid.uuid4()
    out=None
    is_url=True
    url=f'https://omnibus-snap.hf.space/?chat={chat}&__theme={check_t}'
    if wait_time>10000:
        wait_time=10000
    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
            print(html)            
            driver.implicitly_wait(int(wait_time*30))
            driver.set_window_size(int(width), int(height))
            obj = driver.find_element(By.ID, "component-1")            

            screenshot = obj.screenshot(f'{uid}-tmp.png')

            yield Image.open(f'{uid}-tmp.png'), '<center>Finished loading. If loading is incomplete, restart with longer wait time',out_box,out

            
            #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")            
            #messages = driver.find_elements(By.CLASS_NAME, "panel")
            messages = driver.find_elements(By.CLASS_NAME, "message-row")
            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',check_t)
                    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]}',check_t)
                    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]}',check_t)
                        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'            
        except WebDriverException as e:
            return Image.new('RGB', (1, 1)), f'<center>Please enter a valid URL of a website/host.',out_box,out
        finally:
            if driver:
                driver.quit()
        yield Image.open(f'{uid}-tmp.png'), mes,out_box,out
    else:
        yield None, '<center>Please enter a valid URL of a website/host.',out_box,out
def load_frame(inp):
        html_mod=f"""
        <div id="my_game" style="height:600px;width:100%">
            <iframe
            id="myIframe"
            src="https://omnibus-snap.hf.space/?chat=[(123,555),(2342,32344)]&__theme=dark"        
        	frameborder="0"
        	width="100%"
            height="100%"
            ></iframe>
        </div>"""
        return html_mod

with gr.Blocks() as app:
    gr.HTML("""<head><meta http-equiv="refresh" content="180"></head>""")

    with gr.Row():
        with gr.Column():
            inp = gr.Textbox(label="Chat: [('user','bot'),('user','bot')]",lines=6,value=[('123','555'),('2342','32344')])
            wait_t=gr.Slider(1, 10000, value=2000, label="Wait time", info="Wait for page to load. More data = Longer wait")
            with gr.Row():
                btn= gr.Button("Snapshot")
                stop_btn=gr.Button("Stop")
            
        with gr.Column():
            with gr.Row():
                check_t=gr.Radio(label="Theme", choices=["light","dark"], value="dark")
                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,11,12,13,14,15,16,17,18,19,20], 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')
    with gr.Row():
        with gr.Column():
            out=gr.Image()
            outgal=gr.Gallery()
        outim = gr.Image()
    prev=gr.HTML()
    #load_btn.click(load_frame,inp,prev)
    go = btn.click(run_script,[inp,height,width,check_b,check_h,check_t,wait_t],[outim,message,outgal,out])
    stop_btn.click(None,None,None,cancels=go)
app.queue(default_concurrency_limit=10).launch()