from pathlib import Path
from datetime import datetime
from landingai_ade import LandingAIADE
from PIL import Image
import pymupdf
def save_chunks_as_images(parse_response, document_path, output_base_dir="groundings"):
"""Save each parsed chunk as a separate image file."""
# Create timestamped output directory
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
document_name = Path(document_path).stem
output_dir = Path(output_base_dir) / f"{document_name}_{timestamp}"
def save_page_chunks(image, chunks, page_num):
"""Save all chunks for a specific page."""
img_width, img_height = image.size
# Create page-specific directory
page_dir = output_dir / f"page_{page_num}"
page_dir.mkdir(parents=True, exist_ok=True)
for chunk in chunks:
# Check if chunk belongs to this page
if chunk.grounding.page != page_num:
continue
box = chunk.grounding.box
# Convert normalized coordinates to pixel coordinates
x1 = int(box.left * img_width)
y1 = int(box.top * img_height)
x2 = int(box.right * img_width)
y2 = int(box.bottom * img_height)
# Crop the chunk region
chunk_img = image.crop((x1, y1, x2, y2))
# Save with descriptive filename
filename = f"{chunk.type}.{chunk.id}.png"
output_path = page_dir / filename
chunk_img.save(output_path)
print(f"Saved chunk: {output_path}")
if document_path.suffix.lower() == '.pdf':
pdf = pymupdf.open(document_path)
total_pages = len(pdf)
for page_num in range(total_pages):
page = pdf[page_num]
pix = page.get_pixmap(matrix=pymupdf.Matrix(2, 2)) # 2x scaling for clarity
img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
# Save chunks for this page
save_page_chunks(img, parse_response.chunks, page_num)
pdf.close()
else:
# Load image file directly
img = Image.open(document_path)
if img.mode != "RGB":
img = img.convert("RGB")
# Save chunks for single page
save_page_chunks(img, parse_response.chunks, 0)
print(f"\nAll chunks saved to: {output_dir}")
return output_dir
# Initialize client (uses the API key from the VISION_AGENT_API_KEY environment variable)
client = LandingAIADE()
# Replace with your file path
document_path = Path("/path/to/file/document")
# Parse the document
print("Parsing document...")
parse_response = client.parse(
document=document_path,
model="dpt-2-latest"
)
print("Parsing complete!")
# Save chunks as images
save_chunks_as_images(parse_response, document_path)