File size: 6,298 Bytes
d520909 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 |
"""
Image Document Loading
Handles single images and multi-page TIFF documents.
"""
import logging
from pathlib import Path
from typing import Iterator, List, Optional, Tuple, Union
import numpy as np
from PIL import Image
from .base import (
DocumentFormat,
DocumentInfo,
DocumentLoader,
PageInfo,
PageRenderer,
RenderOptions,
)
logger = logging.getLogger(__name__)
class ImageLoader(DocumentLoader):
"""
Image document loader.
Handles common image formats (JPEG, PNG, etc.) and multi-page TIFF.
"""
SUPPORTED_EXTENSIONS = {
".jpg", ".jpeg", ".png", ".bmp", ".gif",
".tif", ".tiff", ".webp"
}
def __init__(self):
self._images: List[Image.Image] = []
self._info: Optional[DocumentInfo] = None
self._path: Optional[Path] = None
def load(self, path: Union[str, Path]) -> DocumentInfo:
"""Load image(s) and extract metadata."""
self._path = Path(path)
if not self._path.exists():
raise FileNotFoundError(f"Image file not found: {self._path}")
suffix = self._path.suffix.lower()
if suffix not in self.SUPPORTED_EXTENSIONS:
raise ValueError(f"Unsupported image format: {suffix}")
# Close any previously loaded images
self.close()
# Load image(s)
img = Image.open(self._path)
# Handle multi-page TIFF
if suffix in {".tif", ".tiff"}:
self._load_multipage_tiff(img)
else:
# Single image
self._images = [img.convert("RGB")]
# Build page info
pages = []
for i, page_img in enumerate(self._images):
dpi = page_img.info.get("dpi", (72, 72))
if isinstance(dpi, tuple):
dpi = int(dpi[0])
else:
dpi = int(dpi)
page_info = PageInfo(
page_number=i + 1,
width_pixels=page_img.width,
height_pixels=page_img.height,
dpi=dpi,
has_images=True
)
pages.append(page_info)
# Determine format
if suffix in {".tif", ".tiff"} and len(self._images) > 1:
doc_format = DocumentFormat.TIFF_MULTIPAGE
else:
doc_format = DocumentFormat.IMAGE
self._info = DocumentInfo(
path=self._path,
format=doc_format,
num_pages=len(self._images),
pages=pages,
file_size_bytes=self._path.stat().st_size,
is_scanned=True, # Images are always "scanned"
has_text_layer=False
)
return self._info
def _load_multipage_tiff(self, img: Image.Image) -> None:
"""Load all pages from a multi-page TIFF."""
self._images = []
try:
page_num = 0
while True:
img.seek(page_num)
# Copy the frame to avoid issues with lazy loading
self._images.append(img.copy().convert("RGB"))
page_num += 1
except EOFError:
# Reached end of TIFF
pass
if not self._images:
raise ValueError("No pages found in TIFF file")
def close(self) -> None:
"""Close all loaded images."""
for img in self._images:
try:
img.close()
except Exception:
pass
self._images = []
def is_loaded(self) -> bool:
"""Check if images are loaded."""
return len(self._images) > 0
@property
def info(self) -> Optional[DocumentInfo]:
"""Get document info."""
return self._info
def get_image(self, page_number: int) -> Image.Image:
"""Get PIL Image for a specific page (1-indexed)."""
if not self._images:
raise RuntimeError("No images loaded")
if page_number < 1 or page_number > len(self._images):
raise ValueError(f"Invalid page number: {page_number}")
return self._images[page_number - 1]
class ImageRenderer(PageRenderer):
"""
Image page renderer.
Renders images with optional resizing and format conversion.
"""
def __init__(self, loader: ImageLoader):
self._loader = loader
def render_page(
self,
page_number: int,
options: Optional[RenderOptions] = None
) -> np.ndarray:
"""Render an image page."""
if not self._loader.is_loaded():
raise RuntimeError("No document loaded")
options = options or RenderOptions()
img = self._loader.get_image(page_number)
# Get original DPI
original_dpi = img.info.get("dpi", (72, 72))
if isinstance(original_dpi, tuple):
original_dpi = original_dpi[0]
# Resize if DPI differs
if options.dpi != original_dpi and original_dpi > 0:
scale = options.dpi / original_dpi
new_size = (int(img.width * scale), int(img.height * scale))
resample = Image.LANCZOS if options.antialias else Image.NEAREST
img = img.resize(new_size, resample=resample)
# Convert color mode
if options.color_mode == "L":
img = img.convert("L")
elif options.color_mode == "RGBA":
img = img.convert("RGBA")
else: # RGB
img = img.convert("RGB")
return np.array(img)
def render_pages(
self,
page_numbers: Optional[List[int]] = None,
options: Optional[RenderOptions] = None
) -> Iterator[Tuple[int, np.ndarray]]:
"""Render multiple pages."""
if not self._loader.is_loaded():
raise RuntimeError("No document loaded")
info = self._loader.info
if page_numbers is None:
page_numbers = list(range(1, info.num_pages + 1))
for page_num in page_numbers:
yield page_num, self.render_page(page_num, options)
def load_image(path: Union[str, Path]) -> Tuple[ImageLoader, ImageRenderer]:
"""
Convenience function to load an image document.
Returns:
Tuple of (loader, renderer)
"""
loader = ImageLoader()
loader.load(path)
renderer = ImageRenderer(loader)
return loader, renderer
|