| | from __future__ import annotations
|
| |
|
| | import base64
|
| | import io
|
| | import os
|
| | import shutil
|
| | import uuid
|
| | from pathlib import Path
|
| |
|
| | from PIL import Image
|
| |
|
| | from state import ImageRef
|
| |
|
| |
|
| | class ImageStore:
|
| | """Disk-based image manager. LangGraph state only carries lightweight
|
| | ``ImageRef`` dicts; all heavy image bytes live on disk."""
|
| |
|
| | def __init__(self, base_dir: str):
|
| | self.base_dir = Path(base_dir)
|
| | self.base_dir.mkdir(parents=True, exist_ok=True)
|
| | self._pages_dir = self.base_dir / "pages"
|
| | self._crops_dir = self.base_dir / "crops"
|
| | self._annotated_dir = self.base_dir / "annotated"
|
| | for d in (self._pages_dir, self._crops_dir, self._annotated_dir):
|
| | d.mkdir(exist_ok=True)
|
| |
|
| |
|
| |
|
| |
|
| |
|
| | def save_page_image(self, page_num: int, image_bytes: bytes) -> ImageRef:
|
| | img = Image.open(io.BytesIO(image_bytes))
|
| | fname = f"page_{page_num}.png"
|
| | path = self._pages_dir / fname
|
| | img.save(str(path), format="PNG")
|
| | return ImageRef(
|
| | id=f"page_{page_num}",
|
| | path=str(path),
|
| | label=f"Page {page_num} (full page)",
|
| | page_num=page_num,
|
| | crop_type="full_page",
|
| | width=img.width,
|
| | height=img.height,
|
| | )
|
| |
|
| | def save_crop(
|
| | self,
|
| | page_num: int,
|
| | crop_id: str,
|
| | image: Image.Image,
|
| | label: str,
|
| | ) -> ImageRef:
|
| | fname = f"page_{page_num}_{crop_id}.png"
|
| | path = self._crops_dir / fname
|
| | image.save(str(path), format="PNG")
|
| | return ImageRef(
|
| | id=f"page_{page_num}_{crop_id}",
|
| | path=str(path),
|
| | label=label,
|
| | page_num=page_num,
|
| | crop_type="crop",
|
| | width=image.width,
|
| | height=image.height,
|
| | )
|
| |
|
| | def save_annotated(
|
| | self,
|
| | source_ref: ImageRef,
|
| | annotated_image: Image.Image,
|
| | ) -> ImageRef:
|
| | ann_id = f"{source_ref['id']}_ann"
|
| | fname = f"{ann_id}.png"
|
| | path = self._annotated_dir / fname
|
| | annotated_image.save(str(path), format="PNG")
|
| | return ImageRef(
|
| | id=ann_id,
|
| | path=str(path),
|
| | label=f"{source_ref['label']} [annotated]",
|
| | page_num=source_ref["page_num"],
|
| | crop_type="annotated",
|
| | width=annotated_image.width,
|
| | height=annotated_image.height,
|
| | )
|
| |
|
| |
|
| |
|
| |
|
| |
|
| | def load_image(self, ref: ImageRef) -> Image.Image:
|
| | return Image.open(ref["path"])
|
| |
|
| | def load_bytes(self, ref: ImageRef) -> bytes:
|
| | with open(ref["path"], "rb") as f:
|
| | return f.read()
|
| |
|
| | def get_page_image_path(self, page_num: int) -> str:
|
| | return str(self._pages_dir / f"page_{page_num}.png")
|
| |
|
| | def load_page_bytes(self, page_num: int) -> bytes:
|
| | path = self.get_page_image_path(page_num)
|
| | with open(path, "rb") as f:
|
| | return f.read()
|
| |
|
| |
|
| |
|
| |
|
| |
|
| | def to_gemini_part(self, ref: ImageRef):
|
| | """Return a ``google.genai.types.Part`` for Gemini multimodal prompts."""
|
| | from google.genai import types
|
| | return types.Part.from_bytes(
|
| | data=self.load_bytes(ref),
|
| | mime_type="image/png",
|
| | )
|
| |
|
| | def to_openai_base64(self, ref: ImageRef) -> dict:
|
| | """Return an OpenAI-compatible image content block (base64 data URI)."""
|
| | b64 = base64.b64encode(self.load_bytes(ref)).decode("utf-8")
|
| | return {
|
| | "type": "image_url",
|
| | "image_url": {"url": f"data:image/png;base64,{b64}"},
|
| | }
|
| |
|
| | def create_thumbnail(self, ref: ImageRef, max_size: int = 400) -> bytes:
|
| | img = self.load_image(ref)
|
| | img.thumbnail((max_size, max_size))
|
| | buf = io.BytesIO()
|
| | img.save(buf, format="PNG")
|
| | return buf.getvalue()
|
| |
|
| |
|
| |
|
| |
|
| |
|
| | def cleanup(self):
|
| | if self.base_dir.exists():
|
| | shutil.rmtree(self.base_dir, ignore_errors=True)
|
| |
|
