Omnibus's picture
Update app.py
b2b8f0b verified
raw
history blame
9.15 kB
#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()