Update app.py

app.py

@@ -62,26 +62,44 @@ def run_script(url: str, height: int, width: int, check_b,check_h):
             
             obj = driver.find_element(By.CLASS_NAME, "main")            
             messages = driver.find_elements(By.CLASS_NAME, "message")
-            print (len(messages))
-            mes_blocks=len(messages)%2
-            for ea in check_b:
+            #print (len(messages))
+            #mes_blocks=len(messages)%2
+            
+            for i,ea in enumerate(check_b):
                 try:    
                     ea = int(ea)
                     ea = (ea*2)-2
                     print (ea)
-                    messages[ea].screenshot(f'{ea}-{uid}-tmp.png')
-                    messages[ea+1].screenshot(f'{ea+1}-{uid}-tmp.png')
-                    if check_h:
-                        head_paste = get_concat_v_cut(f'head-{uid}-tmp.png',f'{ea}-{uid}-tmp.png')
-                        head_paste.save(f'im1-{uid}-tmp.png')
-                        print(head_paste)
-                        out_box.append(get_concat_v_cut(f'im1-{uid}-tmp.png',f'{ea+1}-{uid}-tmp.png'))
-                    if not check_h:    
-                        out_box.append(get_concat_v_cut(f'{ea}-{uid}-tmp.png',f'{ea+1}-{uid}-tmp.png'))
+                    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')
+                
+
                     print(f'out_box:: {out_box}')
                     #mes_box.clear()
                 except Exception:
                     mes="<center>Some blocks returned an error"
+            if out_box:
+                for i in range(len(out_box)-2):
+                    if i == 0:    
+                        im_roll = get_concat_v_cut(f'{out_box[i]}',f'{out_box[i+1]}')
+                        im_roll.save(f'comb-{uid}-tmp.png')
+                    else:
+                        im_roll = get_concat_v_cut(f'comb-{uid}-tmp.png',f'{out_box[i+1]}')
+                        im_roll.save(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'
+            
+            
+            
+            
+            
             '''
             for i,mes in enumerate(messages):
                 print (mes.text)
@@ -116,15 +134,15 @@ def run_script(url: str, height: int, width: int, check_b,check_h):
             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}',out_box
+            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'), mes,out_box
+        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
+        return None, '<center>Please enter a valid URL of a website/host.',out_box,out
 
 with gr.Blocks() as app:
     
@@ -142,8 +160,9 @@ with gr.Blocks() as app:
     message=gr.HTML('<center>Enter URL')
     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])
+    btn.click(run_script,[inp,height,width,check_b,check_h],[outim,message,outgal,out])
 app.launch()