Update app.py

try to remove background by specifying viewport width and height

app.py

@@ -79,37 +79,11 @@ async def capture_screenshot(image_type: str):
        # await page.wait_for_selector("img", timeout=120000)  # Wait for the <img> to be visible
         await page.wait_for_timeout(5000)  # Allow sketch to fully render
         print("Capturing screenshot...")
-        await page.screenshot(path="img.png")
+        await page.set_viewport_size({"width": 800, "height": 800})
+        await page.locator("iframe").screenshot(path="img.png")
         await browser.close()
         print("Screenshot saved!")
-
-# def crop_based_on_bg(image_path: str, bg_color=(59, 59, 59), tolerance=10):
-#     img = Image.open(image_path).convert("RGB")
-#     img_array = np.array(img)
-
-#     # Define fixed crop for top header
-#     top_crop = 50  
-#     #width, height = img.size
-
-#     # Get height & width
-#     height, width, _ = img_array.shape
-
-#     # Extract the portion of the image below the fixed header
-#     img_no_header = img_array[top_crop:, :, :]
-
-#     # Compute a mask for pixels that are NOT within the background tolerance
-#     lower_bound = np.array(bg_color) - tolerance
-#     upper_bound = np.array(bg_color) + tolerance
-#     mask = np.any((img_no_header < lower_bound) | (img_no_header > upper_bound), axis=2)
-
-#     # Find leftmost and rightmost non-background pixels
-#     col_sums = mask.sum(axis=0)
-#     left_crop = np.argmax(col_sums > 0)  # First column with content
-#     right_crop = width - np.argmax(col_sums[::-1] > 0)  # Last column with content
-
-#     # Crop and save
-#     cropped_img = img.crop((left_crop, top_crop, right_crop, height))
-#     return cropped_img      
+   
      
 @tool
 def grab_image(image_type: str) -> Image:
@@ -131,14 +105,10 @@ def grab_image(image_type: str) -> Image:
 
     print("Loading image for Gradio...")
 
-    #cropped_img = crop_based_on_bg("img.png", (59, 59, 59), 10)
     img = Image.open("img.png")
-    header_height = 50
-    left_crop = 0
-    right_crop = 0
-
-    cropped_img = img.crop((left_crop, header_height, right_crop, img.height))
-    return cropped_img
+    
+    
+    return img
 
 @tool
 def get_joke() -> str: