Coverage for /home/admin/Documents/AI/applications/lexigram-dev/lexigram/experimental/ai/lexigram-ai-rag/src/lexigram/ai/rag/multimodal/loaders/image.py: 37%

1"""Image loader for multi-modal RAG. 

2 

3This module provides loading capabilities for image documents with support for: 

4- Multiple formats: JPEG, PNG, GIF, WebP, BMP, TIFF, SVG 

5- EXIF metadata extraction 

6- Optional OCR text extraction 

7- Image validation and normalization 

8""" 

9 

10from __future__ import annotations 

11 

12import asyncio 

13import io 

14from pathlib import Path 

15import sys 

16from typing import TYPE_CHECKING, Any, cast 

17 

18from lexigram.logging import ( 

19 get_logger, 

20) 

21 

22if TYPE_CHECKING: 

23 from collections.abc import Mapping 

24 

25logger = get_logger(__name__) 

26 

27# Defer heavy imports (Pillow, pytesseract) until runtime to avoid import-time failures 

28# The loader will attempt to import these when needed and raise helpful errors if missing. 

29Image: Any = None 

30TAGS: Mapping[int, str] = {} 

31PIL_AVAILABLE = False 

32 

33_pytesseract = None 

34TESSERACT_AVAILABLE = False 

35 

36if TYPE_CHECKING: 

37 # Use `from PIL.Image import Image` so we reference the PIL Image *class* for typing 

38 from PIL.Image import Image as PILImage 

39 

40 PILImageType = PILImage 

41 from lexigram.ai.rag.multimodal.types import ( 

42 ImageDocument, 

43 ImageFormat, 

44 ImageMetadata, 

45 ) 

46else: 

47 PILImageType: Any = Any 

48 

49# Make a small runtime alias for the PIL Image type to aid type inference 

50# (the actual Image object is assigned at runtime when Pillow is available) 

51PILImageTypeRuntime: Any = PILImageType 

52 

53from lexigram.ai.rag.exceptions import ImageLoaderError 

54from lexigram.ai.rag.multimodal.types import ( 

55 ImageDocument, 

56 ImageFormat, 

57 ImageMetadata, 

58) 

59 

60 

61def __getattr__(name: str) -> Any: 

62 """Allow module-level attribute assignments for lazy optional imports. 

63 

64 This satisfies mypy while still allowing runtime lazy loading of Pillow 

65 and pytesseract. The actual attributes (Image, TAGS, PIL_AVAILABLE, 

66 _pytesseract, TESSERACT_AVAILABLE) are set at runtime in __init__. 

67 """ 

68 if name in { 

69 "Image", 

70 "TAGS", 

71 "PIL_AVAILABLE", 

72 "_pytesseract", 

73 "TESSERACT_AVAILABLE", 

74 }: 

75 try: 

76 return sys.modules[__name__].__dict__[name] 

77 except KeyError: 

78 pass 

79 msg = f"module {__name__!r} has no attribute {name!r}" 

80 raise AttributeError(msg) 

81 

82 

83class ImageLoader: 

84 """Loader for image documents. 

85 

86 Supports loading images from files or bytes with optional: 

87 - EXIF metadata extraction 

88 - OCR text extraction 

89 - Format validation 

90 - Image normalization 

91 

92 Args: 

93 extract_exif: Whether to extract EXIF metadata 

94 extract_text: Whether to perform OCR text extraction 

95 tesseract_config: Custom tesseract configuration string 

96 max_size: Maximum image dimension (resize if larger) 

97 convert_to_rgb: Convert images to RGB mode 

98 

99 Example: 

100 >>> loader = ImageLoader(extract_text=True) 

101 >>> doc = await loader.load("/path/to/image.jpg") 

102 >>> print(doc.text_content) 

103 "Extracted text from image" 

104 """ 

105 

106 def __init__( 

107 self, 

108 extract_exif: bool = True, 

109 extract_text: bool = False, 

110 tesseract_config: str = "--psm 3", 

111 max_size: int | None = None, 

112 convert_to_rgb: bool = False, 

113 ): 

114 """Initialize image loader.""" 

115 # Try to import Pillow lazily (so importing this module doesn't require Pillow) 

116 mod = sys.modules[__name__] 

117 # If tests have patched PIL_AVAILABLE, respect that and avoid re-importing 

118 # Use getattr on the actual module object to read the live value (which 

119 # may have been patched at runtime) rather than the module-level constant 

120 # captured at import time. 

121 _pil_available = getattr(mod, "PIL_AVAILABLE", False) 

122 if not _pil_available: 

123 try: 

124 from PIL import Image as _Image 

125 from PIL.ExifTags import TAGS as _TAGS 

126 

127 mod.Image = _Image # type: ignore[attr-defined] 

128 mod.TAGS = _TAGS # type: ignore[attr-defined] 

129 mod.PIL_AVAILABLE = True # type: ignore[attr-defined] 

130 except (ImportError, ModuleNotFoundError): 

131 mod.PIL_AVAILABLE = False # type: ignore[attr-defined] 

132 

133 _current_pil = getattr(mod, "PIL_AVAILABLE", False) 

134 if not _current_pil: 

135 raise ImportError( 

136 "PIL/Pillow is required for image loading. Install with: pip install Pillow", 

137 ) 

138 

139 # Lazy import of pytesseract only if OCR is requested 

140 if extract_text: 

141 try: 

142 import pytesseract as _pt 

143 

144 mod._pytesseract = _pt # type: ignore[attr-defined] 

145 mod.TESSERACT_AVAILABLE = True # type: ignore[attr-defined] 

146 except (ImportError, ModuleNotFoundError): 

147 mod.TESSERACT_AVAILABLE = False # type: ignore[attr-defined] 

148 

149 _tesseract_available = getattr(mod, "TESSERACT_AVAILABLE", False) 

150 if extract_text and not _tesseract_available: 

151 raise ImportError( 

152 "pytesseract is required for OCR. Install with: pip install pytesseract", 

153 ) 

154 

155 self.extract_exif = extract_exif 

156 self.extract_text = extract_text 

157 

158 # Validate tesseract_config to prevent command injection 

159 # Only allow alphanumeric, spaces, and specific flags 

160 import re 

161 

162 if not re.match(r"^[a-zA-Z0-9\s\-\.\_=]*$", tesseract_config): 

163 msg = f"Insecure tesseract_config detected: {tesseract_config}" 

164 raise ValueError(msg) 

165 

166 self.tesseract_config = tesseract_config 

167 self.max_size = max_size 

168 self.convert_to_rgb = convert_to_rgb 

169 

170 async def load( 

171 self, 

172 source: str | Path | bytes, 

173 caption: str | None = None, 

174 alt_text: str | None = None, 

175 source_url: str | None = None, 

176 **metadata_kwargs: Any, 

177 ) -> ImageDocument: 

178 """Load an image document. 

179 

180 Args: 

181 source: File path, URL, or raw bytes 

182 caption: Optional image caption 

183 alt_text: Optional alt text 

184 source_url: Optional source URL 

185 **metadata_kwargs: Additional metadata fields 

186 

187 Returns: 

188 ImageDocument with loaded image data 

189 

190 Raises: 

191 ImageLoaderError: If loading fails 

192 """ 

193 try: 

194 # Load image 

195 content: bytes | str 

196 file_path: Path | None 

197 

198 if isinstance(source, (str, Path)): 

199 file_path = Path(source) 

200 if not await asyncio.to_thread(file_path.exists): 

201 msg = f"Image file not found: {source}" 

202 raise ImageLoaderError(msg) 

203 

204 # Image.open is lazy, but it still performs initial I/O 

205 img = await asyncio.to_thread(Image.open, file_path) 

206 pil_img = cast("PILImageType", img) 

207 content = str(file_path) 

208 elif isinstance(source, bytes): 

209 # BytesIO is memory-only, but Image.open still does format detection/header parsing 

210 img = await asyncio.to_thread(Image.open, io.BytesIO(source)) 

211 pil_img = cast("PILImageType", img) 

212 content = source 

213 file_path = None 

214 else: 

215 msg = f"Unsupported source type: {type(source)}" 

216 raise ImageLoaderError(msg) 

217 

218 # Detect format (fast) 

219 image_format = self._detect_format(pil_img) 

220 

221 # Convert mode if requested (CPU intensive) 

222 if self.convert_to_rgb and pil_img.mode != "RGB": 

223 pil_img = await asyncio.to_thread(pil_img.convert, "RGB") 

224 

225 # Resize if needed (CPU intensive) 

226 if self.max_size: 

227 pil_img = await asyncio.to_thread( 

228 self._resize_image, 

229 pil_img, 

230 self.max_size, 

231 ) 

232 

233 # Extract EXIF metadata (fast/I/O already done) 

234 exif_data = {} 

235 if self.extract_exif: 

236 exif_data = await asyncio.to_thread(self._extract_exif, pil_img) 

237 

238 # Build metadata 

239 metadata = ImageMetadata( 

240 caption=caption, 

241 alt_text=alt_text, 

242 source=source_url, 

243 exif=exif_data, 

244 **metadata_kwargs, 

245 ) 

246 

247 # Extract text via OCR (Extremely CPU intensive) 

248 text_content = None 

249 if self.extract_text: 

250 text_content = await asyncio.to_thread(self._extract_text, img) 

251 

252 # Create document 

253 return ImageDocument( 

254 content=content, 

255 format=image_format, 

256 width=pil_img.width, 

257 height=pil_img.height, 

258 metadata=metadata, 

259 text_content=text_content, 

260 file_path=file_path, 

261 ) 

262 

263 except Exception as e: 

264 if isinstance(e, ImageLoaderError): 

265 raise 

266 msg = f"Failed to load image {source!r}: {e}" 

267 raise ImageLoaderError(msg) from e 

268 

269 async def load_batch( 

270 self, 

271 sources: list[str | Path | bytes], 

272 **kwargs: Any, 

273 ) -> list[ImageDocument]: 

274 """Load multiple images. 

275 

276 Args: 

277 sources: List of image sources 

278 **kwargs: Common metadata for all images 

279 

280 Returns: 

281 List of loaded ImageDocuments 

282 """ 

283 if not sources: 

284 return [] 

285 

286 async def _safe_load(source) -> Any: 

287 try: 

288 return await self.load(source, **kwargs) 

289 except ImageLoaderError as e: 

290 logger.warning("Failed to load image %r: %s", source, e) 

291 return None 

292 

293 tasks = [_safe_load(source) for source in sources] 

294 results = await asyncio.gather(*tasks) 

295 return [doc for doc in results if doc is not None] 

296 

297 def _detect_format(self, img: PILImage) -> ImageFormat: 

298 """Detect image format. 

299 

300 Args: 

301 img: PIL Image object 

302 

303 Returns: 

304 ImageFormat enum value 

305 """ 

306 format_str = img.format 

307 if not format_str: 

308 msg = "Unable to detect image format" 

309 raise ImageLoaderError(msg) 

310 

311 format_str = format_str.lower() 

312 

313 # Map PIL format to our enum 

314 format_mapping = { 

315 "jpeg": ImageFormat.JPEG, 

316 "jpg": ImageFormat.JPG, 

317 "png": ImageFormat.PNG, 

318 "gif": ImageFormat.GIF, 

319 "webp": ImageFormat.WEBP, 

320 "bmp": ImageFormat.BMP, 

321 "tiff": ImageFormat.TIFF, 

322 "svg": ImageFormat.SVG, 

323 } 

324 

325 if format_str not in format_mapping: 

326 msg = f"Unsupported image format: {format_str}" 

327 raise ImageLoaderError(msg) 

328 

329 return format_mapping[format_str] 

330 

331 def _extract_exif(self, img: PILImage) -> dict[str, Any]: 

332 """Extract EXIF metadata from image. 

333 

334 Args: 

335 img: PIL Image object 

336 

337 Returns: 

338 Dictionary of EXIF data 

339 """ 

340 exif_data = {} 

341 

342 try: 

343 exif = img.getexif() 

344 if exif: 

345 for tag_id, raw_value in exif.items(): 

346 tag = str(TAGS.get(tag_id, tag_id)) 

347 # Convert bytes to string for JSON serialization 

348 value: str | bytes = raw_value 

349 if isinstance(raw_value, bytes): 

350 try: 

351 value = raw_value.decode("utf-8") 

352 except UnicodeDecodeError: 

353 value = str(raw_value) 

354 exif_data[tag] = value 

355 except (ValueError, TypeError, RuntimeError) as e: 

356 # EXIF extraction is optional, don't fail 

357 logger.warning("Failed to extract EXIF data: %s", e) 

358 

359 return exif_data 

360 

361 def _extract_text(self, img: PILImage) -> str: 

362 """Extract text from image using OCR. 

363 

364 Args: 

365 img: PIL Image object 

366 

367 Returns: 

368 Extracted text 

369 """ 

370 if not self.extract_text: 

371 return "" 

372 

373 if not TESSERACT_AVAILABLE or _pytesseract is None: 

374 msg = "Tesseract OCR not available or not installed" 

375 raise ImportError(msg) 

376 

377 try: 

378 text = _pytesseract.image_to_string(img, config=self.tesseract_config) 

379 return text.strip() 

380 except (ValueError, TypeError, RuntimeError, OSError) as e: 

381 logger.warning("OCR extraction failed: %s", e) 

382 return "" 

383 

384 def _resize_image(self, img: PILImage, max_size: int) -> PILImage: 

385 """Resize image if larger than max_size. 

386 

387 Args: 

388 img: PIL Image object 

389 max_size: Maximum dimension 

390 

391 Returns: 

392 Resized image (or original if already small enough) 

393 """ 

394 if max(img.width, img.height) <= max_size: 

395 return img 

396 

397 # Calculate new dimensions maintaining aspect ratio 

398 if img.width > img.height: 

399 new_width = max_size 

400 new_height = int(img.height * (max_size / img.width)) 

401 else: 

402 new_height = max_size 

403 new_width = int(img.width * (max_size / img.height)) 

404 

405 return img.resize((new_width, new_height), Image.Resampling.LANCZOS)