Update app.py

app.py

@@ -37,25 +37,34 @@ def fetch_html(url: str):
         return {"error": str(e)}
 
 @app.get("/download")
-def download_screenshot(url: str):
+def download_image(url: str):
     try:
         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")
-            
-            # Take a screenshot of the page once content is fully loaded
-            screenshot = page.screenshot(full_page=True)  # Set full_page=True to capture the entire page
-            browser.close()
-        
-        # Return the screenshot as a response with image/png content type
-        return Response(content=screenshot, media_type="image/png")
+
+            # 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
+                
+                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")
+            else:
+                browser.close()
+                return {"error": "No image found on the page."}
     except Exception as e:
         return {"error": str(e)}