Update app.py

try to crop background

app.py

@@ -44,9 +44,6 @@ def load_image_sources():
 
 image_sources = load_image_sources()
 
-# def get_image_url(image_type: str):
-#     return image_sources.get(image_type, None)  
-
 def get_image_url(image_type: str):
     """Finds the best match for the given image type using fuzzy matching."""
     choices = list(image_sources.keys())  # Get all available keys
@@ -59,20 +56,6 @@ def get_image_url(image_type: str):
        #return None  # No good match found
         return "https://editor.p5js.org/kfahn/full/2XD5Y8MiV"
 
-# def get_image_url(image_type: str):
-#     """Finds an approximate match using simple substring matching."""
-#     image_type = image_type.lower()
-    
-#     # Check for exact match first
-#     if image_type in image_sources:
-#         return image_sources[image_type]
-
-#     # If no exact match, look for a partial match
-#     for key in image_sources:
-#         if image_type in key.lower():  # Check if input is a substring of a key
-#             return image_sources[key]  
-#    return None  # No match found
-
 async def capture_screenshot(image_type: str):
     """Launches Playwright and uses user input,  if any, to captures a screenshot of an image from p5.js."""
     print("Launching Playwright...")
@@ -98,6 +81,32 @@ async def capture_screenshot(image_type: str):
         await page.screenshot(path="img.png")
         await browser.close()
         print("Screenshot saved!")
+
+def crop_based_on_bg(img: Image, bg_color=(59, 59, 59)):
+    # img = Image.open(image_path).convert("RGB")
+    # pixels = img.load()
+
+    # Define fixed crop for top header
+    top_crop = 50  
+    width, height = img.size
+
+    # Find leftmost and rightmost non-background pixels
+    left_crop, right_crop = 0, width
+    for x in range(width):
+        column = [pixels[x, y] for y in range(top_crop, height)]
+        if any(pixel != bg_color for pixel in column):
+            left_crop = x
+            break
+
+    for x in range(width - 1, -1, -1):
+        column = [pixels[x, y] for y in range(top_crop, height)]
+        if any(pixel != bg_color for pixel in column):
+            right_crop = x
+            break
+
+    # 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:
@@ -118,16 +127,15 @@ def grab_image(image_type: str) -> Image:
     loop.run_until_complete(capture_screenshot(image_type))
 
     print("Loading image for Gradio...")
-    img = Image.open("img.png")
+    img = Image.open("img.png").RGB
     header_height = 50  
     left_crop = 0  # Adjust to remove background
     right_crop = img.width - 0 
 
-    cropped_img = img.crop((left_crop, header_height, right_crop, img.height))
+    #cropped_img = img.crop((left_crop, header_height, right_crop, img.height))
+    cropped_img = crop_based_on_bg(img, left_crop, header_height, right_crop, img.height)
     return cropped_img
 
-
-
 @tool
 def get_joke() -> str:
     """