Update app.py

app.py

@@ -42,29 +42,27 @@ def download_image(url: str):
         with sync_playwright() as p:
             browser = p.chromium.launch(headless=True)
             page = browser.new_page()
-
+            
             # Set custom headers before navigating to the page
             headers = get_custom_headers(url)
             page.set_extra_http_headers(headers)
-
+            
             # Navigate to the URL and wait for network idle
             page.goto(url, timeout=15000, wait_until="networkidle")
-
-            # Locate the first <img> element on the page
-            image_element = page.query_selector("img")  # You can adjust the selector to target a specific image
-
-            if image_element:
-                # Get the image's source URL
-                img_url = image_element.get_attribute("src")
-                image_data = page.request.get(img_url).body  # Fetch the image data from the source URL
-                
+            
+            # Select the first image on the page
+            img_element = page.query_selector("img")
+            
+            if img_element:
+                # Take a screenshot of the first image element
+                screenshot = img_element.screenshot()
                 browser.close()
                 
-                # Return the image data as a response with image/jpeg or image/png, depending on the image type
-                return Response(content=image_data, media_type="image/jpeg" if img_url.endswith(".jpg") else "image/png")
+                # Return the image as a response
+                return Response(content=screenshot, media_type="image/png")
             else:
                 browser.close()
-                return {"error": "No image found on the page."}
+                return {"error": "No image found on the page"}
     except Exception as e:
         return {"error": str(e)}