wehub-resource-sync
413 lines
14 KiB
Python
413 lines
14 KiB
Python
"""Screenshot utilities using Selenium WebDriver.
|
|
|
|
Selenium and webdriver_manager are optional — only needed when capturing
|
|
screenshots (--url-screenshot mode). Import is deferred to function scope
|
|
so the eval script works without them for --local-api users.
|
|
"""
|
|
|
|
import base64
|
|
import logging
|
|
import os
|
|
import time
|
|
|
|
from PIL import Image
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def setup_driver(window_width=1024, window_height=2000, device_scale_factor=1):
|
|
"""Set up Chrome WebDriver.
|
|
|
|
Args:
|
|
window_width: Viewport width (1024 = tile width, ensures screenshots align with tile grid).
|
|
window_height: Initial viewport height.
|
|
device_scale_factor: Pixel density (1 = standard, 2 = retina quality).
|
|
"""
|
|
from selenium import webdriver
|
|
from selenium.webdriver.chrome.service import Service
|
|
from webdriver_manager.chrome import ChromeDriverManager
|
|
|
|
import shutil
|
|
|
|
snap_chromedriver = "/snap/bin/chromium.chromedriver"
|
|
if os.path.exists(snap_chromedriver):
|
|
driver_path = snap_chromedriver
|
|
elif shutil.which("chromedriver"):
|
|
driver_path = shutil.which("chromedriver")
|
|
else:
|
|
driver_path = ChromeDriverManager().install()
|
|
service = Service(driver_path)
|
|
options = webdriver.ChromeOptions()
|
|
|
|
# Find Chrome binary path
|
|
chrome_binary = None
|
|
for chrome_path in [
|
|
"/usr/bin/google-chrome-stable",
|
|
"/usr/bin/google-chrome",
|
|
"/usr/bin/chromium-browser",
|
|
"/usr/bin/chromium",
|
|
]:
|
|
if os.path.exists(chrome_path):
|
|
chrome_binary = chrome_path
|
|
break
|
|
|
|
if chrome_binary:
|
|
options.binary_location = chrome_binary
|
|
|
|
options.add_argument("--headless=new")
|
|
options.add_argument("--disable-gpu")
|
|
options.add_argument(f"--window-size={window_width},{window_height}")
|
|
options.add_argument("--no-sandbox")
|
|
options.add_argument("--disable-dev-shm-usage")
|
|
options.add_argument("--disable-extensions")
|
|
options.add_argument("--disable-setuid-sandbox")
|
|
options.add_argument("--no-zygote")
|
|
options.add_argument("--remote-debugging-port=0")
|
|
options.add_argument(
|
|
"--user-agent=Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 "
|
|
"(KHTML, like Gecko) Chrome/145.0.0.0 Safari/537.36"
|
|
)
|
|
# Retina-quality rendering (2x pixel density)
|
|
if device_scale_factor and device_scale_factor > 1:
|
|
options.add_argument(f"--force-device-scale-factor={device_scale_factor}")
|
|
driver = webdriver.Chrome(service=service, options=options)
|
|
return driver
|
|
|
|
|
|
def _capture_with_scroll(driver, output_path, scroll_pause=0.8, max_scrolls=100):
|
|
"""Capture full page by scrolling and stitching screenshots.
|
|
|
|
Works for PDF viewers, infinite scroll pages, and other dynamic content.
|
|
Uses image comparison to detect when scrolling has stopped.
|
|
"""
|
|
from selenium.webdriver.common.action_chains import ActionChains
|
|
from selenium.webdriver.common.keys import Keys
|
|
import tempfile
|
|
import hashlib
|
|
|
|
def get_screenshot_hash(path):
|
|
"""Get hash of screenshot to detect changes."""
|
|
with Image.open(path) as img:
|
|
return hashlib.md5(img.tobytes()).hexdigest()
|
|
|
|
viewport_height = driver.execute_script("return window.innerHeight")
|
|
viewport_width = driver.execute_script("return window.innerWidth")
|
|
|
|
# Try to scroll to top using multiple methods
|
|
driver.execute_script("window.scrollTo(0, 0)")
|
|
actions = ActionChains(driver)
|
|
actions.send_keys(Keys.HOME)
|
|
actions.perform()
|
|
time.sleep(scroll_pause)
|
|
|
|
screenshots = []
|
|
last_hash = None
|
|
scroll_count = 0
|
|
consecutive_same = 0
|
|
|
|
while scroll_count < max_scrolls:
|
|
# Take screenshot
|
|
temp_file = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
|
|
driver.save_screenshot(temp_file.name)
|
|
|
|
# Check if image changed (detect end of scrolling)
|
|
current_hash = get_screenshot_hash(temp_file.name)
|
|
|
|
if current_hash == last_hash:
|
|
consecutive_same += 1
|
|
os.unlink(temp_file.name)
|
|
if consecutive_same >= 2:
|
|
# Scrolling stopped, we've reached the end
|
|
break
|
|
else:
|
|
consecutive_same = 0
|
|
screenshots.append(temp_file.name)
|
|
last_hash = current_hash
|
|
|
|
# Scroll down using CDP mouse wheel event (works for PDF viewers)
|
|
center_x = viewport_width // 2
|
|
center_y = viewport_height // 2
|
|
try:
|
|
driver.execute_cdp_cmd(
|
|
"Input.dispatchMouseEvent",
|
|
{
|
|
"type": "mouseWheel",
|
|
"x": center_x,
|
|
"y": center_y,
|
|
"deltaX": 0,
|
|
"deltaY": int(viewport_height * 0.8), # Scroll 80% of viewport
|
|
},
|
|
)
|
|
except Exception:
|
|
# Fallback to ActionChains
|
|
actions = ActionChains(driver)
|
|
actions.send_keys(Keys.PAGE_DOWN)
|
|
actions.perform()
|
|
|
|
time.sleep(scroll_pause)
|
|
scroll_count += 1
|
|
|
|
if not screenshots:
|
|
driver.save_screenshot(output_path)
|
|
return
|
|
|
|
if len(screenshots) == 1:
|
|
# Only one screenshot, just use it
|
|
os.rename(screenshots[0], output_path)
|
|
return
|
|
|
|
# Stitch screenshots vertically
|
|
# Each screenshot is viewport_height, but they overlap
|
|
# We'll stack them with some overlap detection
|
|
images = [Image.open(p) for p in screenshots]
|
|
|
|
# Simple stacking: assume each Page Down scrolls ~80% of viewport
|
|
overlap = int(viewport_height * 0.2)
|
|
total_height = viewport_height + (len(images) - 1) * (viewport_height - overlap)
|
|
|
|
stitched = Image.new("RGB", (viewport_width, total_height), (255, 255, 255))
|
|
|
|
y_offset = 0
|
|
for i, img in enumerate(images):
|
|
if i == 0:
|
|
stitched.paste(img, (0, 0))
|
|
y_offset = viewport_height - overlap
|
|
else:
|
|
# Crop top overlap region and paste
|
|
cropped = img.crop((0, overlap, viewport_width, viewport_height))
|
|
stitched.paste(cropped, (0, y_offset))
|
|
y_offset += viewport_height - overlap
|
|
|
|
# Close images and clean up
|
|
for img in images:
|
|
img.close()
|
|
for p in screenshots:
|
|
if os.path.exists(p):
|
|
os.unlink(p)
|
|
|
|
# Trim any white space at bottom
|
|
stitched = stitched.crop((0, 0, viewport_width, y_offset + overlap))
|
|
stitched.save(output_path)
|
|
|
|
|
|
def _eager_load_images(driver):
|
|
"""Force lazy images to load by promoting data-src and setting loading='eager'."""
|
|
driver.execute_script("""
|
|
(function() {
|
|
var imgs = document.querySelectorAll('img');
|
|
for (var i = 0; i < imgs.length; i++) {
|
|
var img = imgs[i];
|
|
try {
|
|
if (img.loading === 'lazy') img.loading = 'eager';
|
|
var dataSrc = img.getAttribute('data-src') || (img.dataset && img.dataset.src);
|
|
var dataSrcset = img.getAttribute('data-srcset') || (img.dataset && img.dataset.srcset);
|
|
if (dataSrc) img.setAttribute('src', dataSrc);
|
|
if (dataSrcset) img.setAttribute('srcset', dataSrcset);
|
|
} catch(e) {}
|
|
}
|
|
})();
|
|
""")
|
|
|
|
|
|
def _wait_for_images(driver, timeout=10):
|
|
"""Wait for all document images to finish loading (load or error)."""
|
|
deadline = time.time() + timeout
|
|
while time.time() < deadline:
|
|
pending = driver.execute_script("""
|
|
return Array.from(document.images || [])
|
|
.filter(function(img) { return !(img.complete && img.naturalWidth > 0); }).length;
|
|
""")
|
|
if pending == 0:
|
|
return
|
|
time.sleep(0.3)
|
|
|
|
|
|
def _scroll_to_trigger_lazy_load(driver, page_height):
|
|
"""Scroll through the page to trigger lazy-loaded content, then scroll back to top."""
|
|
viewport_height = driver.execute_script("return window.innerHeight") or 1080
|
|
y = 0
|
|
while y < page_height:
|
|
driver.execute_script(f"window.scrollTo(0, {y})")
|
|
time.sleep(0.15)
|
|
_eager_load_images(driver)
|
|
y += viewport_height
|
|
# Wait for all images to load after full scroll
|
|
_wait_for_images(driver, timeout=10)
|
|
# Scroll back to top
|
|
driver.execute_script("window.scrollTo(0, 0)")
|
|
time.sleep(0.3)
|
|
|
|
|
|
def capture_screenshot(url, output_path, full_page=False, scroll_capture=False):
|
|
"""Capture screenshot of a URL.
|
|
|
|
Args:
|
|
url: URL to capture
|
|
output_path: Path to save screenshot
|
|
full_page: If True, resize window to capture full page (works for normal pages)
|
|
scroll_capture: If True, scroll and stitch screenshots (works for PDF viewers, etc.)
|
|
"""
|
|
driver = None
|
|
try:
|
|
driver = setup_driver()
|
|
driver.get(url)
|
|
time.sleep(3) # Wait for initial load
|
|
|
|
# Force lazy images to load eagerly
|
|
_eager_load_images(driver)
|
|
_wait_for_images(driver, timeout=5)
|
|
|
|
if scroll_capture:
|
|
# Scroll-based capture for PDF viewers and similar
|
|
# just for PDF
|
|
_capture_with_scroll(driver, output_path)
|
|
elif full_page:
|
|
# Get page height, keep original window width to avoid horizontal tiling
|
|
total_height = driver.execute_script("return document.body.scrollHeight")
|
|
current_window = driver.get_window_size()
|
|
|
|
# Scroll through page to trigger lazy-loaded images
|
|
_scroll_to_trigger_lazy_load(driver, total_height)
|
|
|
|
# Re-measure height (may change after lazy content loads)
|
|
total_height = driver.execute_script("return document.body.scrollHeight")
|
|
driver.set_window_size(current_window["width"], total_height)
|
|
time.sleep(0.5)
|
|
|
|
# Final wait for any images triggered by resize
|
|
_eager_load_images(driver)
|
|
_wait_for_images(driver, timeout=5)
|
|
|
|
driver.save_screenshot(output_path)
|
|
else:
|
|
driver.save_screenshot(output_path)
|
|
|
|
# Convert to RGB (remove alpha channel if present)
|
|
with Image.open(output_path) as img:
|
|
if img.mode in ("RGBA", "LA") or (
|
|
img.mode == "P" and "transparency" in img.info
|
|
):
|
|
bg = Image.new("RGB", img.size, (255, 255, 255))
|
|
if img.mode != "RGBA":
|
|
img = img.convert("RGBA")
|
|
bg.paste(img, mask=img.split()[3])
|
|
img = bg
|
|
img.save(output_path)
|
|
elif img.mode != "RGB":
|
|
img = img.convert("RGB")
|
|
img.save(output_path)
|
|
|
|
return True
|
|
except Exception as e:
|
|
print(f"Screenshot failed for {url}: {e}")
|
|
return False
|
|
finally:
|
|
if driver:
|
|
driver.quit()
|
|
|
|
|
|
def encode_image(image_path, max_pixels: int = 150_000_000, max_height: int = 8000):
|
|
"""Encode image to base64, compressing if too large.
|
|
|
|
Used for Vector DB retrieval where consistent image sizes help with embedding.
|
|
|
|
Args:
|
|
image_path: Path to image file.
|
|
max_pixels: Maximum pixels allowed (default 150M).
|
|
max_height: Maximum height in pixels (default 8000).
|
|
"""
|
|
if not os.path.exists(image_path):
|
|
return None
|
|
|
|
try:
|
|
# Increase PIL limit temporarily
|
|
Image.MAX_IMAGE_PIXELS = 300_000_000
|
|
|
|
with Image.open(image_path) as img:
|
|
# Check if compression needed
|
|
total_pixels = img.width * img.height
|
|
needs_compression = total_pixels > max_pixels or img.height > max_height
|
|
|
|
if not needs_compression:
|
|
# Just read and encode directly
|
|
with open(image_path, "rb") as f:
|
|
return base64.b64encode(f.read()).decode("utf-8")
|
|
|
|
# Compress: resize to fit within limits
|
|
if img.mode != "RGB":
|
|
img = img.convert("RGB")
|
|
|
|
# Calculate new size
|
|
if img.height > max_height:
|
|
ratio = max_height / img.height
|
|
new_width = int(img.width * ratio)
|
|
new_height = max_height
|
|
else:
|
|
# Scale down to fit max_pixels
|
|
ratio = (max_pixels / total_pixels) ** 0.5
|
|
new_width = int(img.width * ratio)
|
|
new_height = int(img.height * ratio)
|
|
|
|
img = img.resize((new_width, new_height), Image.Resampling.LANCZOS)
|
|
|
|
# Encode to JPEG
|
|
from io import BytesIO
|
|
|
|
buffered = BytesIO()
|
|
img.save(buffered, format="JPEG", quality=85)
|
|
return base64.b64encode(buffered.getvalue()).decode("utf-8")
|
|
|
|
except Exception as e:
|
|
print(f"Failed to encode image {image_path}: {e}")
|
|
return None
|
|
|
|
|
|
def encode_image_for_vlm(image_path, max_pixels: int = 89_000_000):
|
|
"""Encode image to base64 for VLM ground truth, minimal processing.
|
|
|
|
For Ground Truth evaluation, we want to preserve original image quality
|
|
and let the VLM handle resizing according to its own requirements.
|
|
Only applies PIL safety limit (89M pixels).
|
|
|
|
Args:
|
|
image_path: Path to image file.
|
|
max_pixels: Maximum pixels (default 89M, PIL's safety limit).
|
|
"""
|
|
if not os.path.exists(image_path):
|
|
return None
|
|
|
|
try:
|
|
# Increase PIL limit temporarily
|
|
Image.MAX_IMAGE_PIXELS = 300_000_000
|
|
|
|
with Image.open(image_path) as img:
|
|
total_pixels = img.width * img.height
|
|
|
|
# Only compress if exceeds PIL safety limit
|
|
if total_pixels <= max_pixels:
|
|
# Just read and encode directly - no resize
|
|
with open(image_path, "rb") as f:
|
|
return base64.b64encode(f.read()).decode("utf-8")
|
|
|
|
# Compress only if exceeds max_pixels
|
|
if img.mode != "RGB":
|
|
img = img.convert("RGB")
|
|
|
|
# Scale down to fit max_pixels
|
|
ratio = (max_pixels / total_pixels) ** 0.5
|
|
new_width = int(img.width * ratio)
|
|
new_height = int(img.height * ratio)
|
|
|
|
img = img.resize((new_width, new_height), Image.Resampling.LANCZOS)
|
|
|
|
# Encode to JPEG
|
|
from io import BytesIO
|
|
|
|
buffered = BytesIO()
|
|
img.save(buffered, format="JPEG", quality=85)
|
|
return base64.b64encode(buffered.getvalue()).decode("utf-8")
|
|
|
|
except Exception as e:
|
|
print(f"Failed to encode image {image_path}: {e}")
|
|
return None
|