Update app.py

app.py

@@ -121,7 +121,24 @@ def run_script(url: str, height: int, width: int, check_b,check_h,wait_time):
             print(html)            
             driver.implicitly_wait(int(wait_time*30))
             driver.set_window_size(int(width), int(height))
-          
+            
+            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, "user-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')
+                    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"
+            
             screenshot = driver.save_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