IPSExtractor/core/pdf_utils.py

"""Utilidades para convertir documentos (PDF/TIFF) a imágenes base64.

En Windows se requiere Poppler. Este proyecto incluye una ruta por defecto:
`dependencies/poppler-24.08.0/Library/bin`. Ajuste si su instalación difiere.
"""
import base64
import logging
import os
from io import BytesIO
from typing import List, Tuple, Optional
import time

from pdf2image import convert_from_path, pdfinfo_from_path
from PIL import Image, ImageSequence

Image.MAX_IMAGE_PIXELS = None

logger = logging.getLogger(__name__)

_DEFAULT_POPPLER = 'dependencies/poppler-24.08.0/Library/bin'


def _resolve_poppler(poppler_path: str | None) -> str:
  """Devuelve la ruta a Poppler: usa la provista o el default del proyecto."""
  return poppler_path or _DEFAULT_POPPLER


def pdf_to_images_bs64(
  pdf_path: str,
  validate: bool = True,
  dpi: int = 800,
  thread_count: int = 10,
  poppler_path: str | None = None,
) -> List[str]:
  """
  Convierte un archivo PDF completo a una lista de imágenes PNG en base64.

  Args:
      pdf_path: Ruta al archivo PDF.
      validate: Si es True, intenta abrir cada imagen generada para verificar su integridad.
      dpi: Resolución de la imagen (puntos por pulgada).
      thread_count: Número de hilos para la conversión paralela de páginas.
      poppler_path: Ruta personalizada a los binarios de Poppler (opcional).

  Returns:
      Lista de strings con las imágenes codificadas en base64.

  Raises:
      ValueError: Si la validación de alguna imagen falla.
  """
  poppler_bin = _resolve_poppler(poppler_path)
  t0 = time.time()
  logger.info("Convirtiendo PDF a imágenes", extra={"pdf": pdf_path, "dpi": dpi})
  images = convert_from_path(pdf_path, dpi=dpi, thread_count=thread_count, poppler_path=poppler_bin)
  resultados: List[str] = []
  for indice, imagen in enumerate(images):
    buffer = BytesIO()
    imagen.save(buffer, format="PNG")
    b64 = base64.b64encode(buffer.getvalue()).decode('utf-8')
    if validate:
      try:
        raw = base64.b64decode(b64, validate=True)
        with Image.open(BytesIO(raw)) as im:
          im.verify()
      except Exception as error:
        raise ValueError(f"Imagen base64 inválida en la página {indice + 1}: {error}")
    resultados.append(b64)
  elapsed_ms = int((time.time() - t0) * 1000)
  logger.debug(f"PDF->imgs OK (pages={len(resultados)}, ms={elapsed_ms}, threads={thread_count})", extra={"pdf": pdf_path, "dpi": dpi})
  logger.info("PDF convertido", extra={"pdf": pdf_path, "dpi": dpi})
  return resultados


def get_pdf_num_pages(pdf_path: str, poppler_path: str | None = None) -> int:
  """
  Obtiene el número total de páginas de un PDF.

  Args:
      pdf_path: Ruta al archivo PDF.
      poppler_path: Ruta personalizada a los binarios de Poppler (opcional).

  Returns:
      Número de páginas (0 si ocurre un error).
  """
  info = pdfinfo_from_path(pdf_path, poppler_path=_resolve_poppler(poppler_path))
  try:
    return int(info.get('Pages') or info.get('pages') or 0)
  except Exception:
    return 0


def pdf_pages_to_images_bs64(
  pdf_path: str,
  first_page: int,
  last_page: int,
  *,
  dpi: int,
  grayscale: bool = True,
  validate: bool = False,
  thread_count: int = 10,
  poppler_path: str | None = None,
) -> List[str]:
  """
  Convierte un rango específico de páginas de un PDF a imágenes base64.

  Args:
      pdf_path: Ruta al archivo PDF.
      first_page: Índice de la primera página (1-indexed).
      last_page: Índice de la última página (inclusivo).
      dpi: Resolución de la imagen.
      grayscale: Si es True, convierte las imágenes a escala de grises para ahorrar ancho de banda/tokens.
      validate: Si es True, verifica la integridad de las imágenes generadas.
      thread_count: Número de hilos para procesamiento paralelo.
      poppler_path: Ruta personalizada a Poppler.

  Returns:
      Lista de imágenes en base64.
  """
  poppler_bin = _resolve_poppler(poppler_path)
  t0 = time.time()
  logger.debug(
    f"Convirtiendo rango PDF→imgs (first={first_page}, last={last_page}, dpi={dpi}, gray={grayscale})",
    extra={"pdf": pdf_path, "dpi": dpi},
  )
  images = convert_from_path(
    pdf_path,
    dpi=dpi,
    thread_count=thread_count,
    poppler_path=poppler_bin,
    first_page=first_page,
    last_page=last_page,
  )
  resultados: List[str] = []
  for idx, img in enumerate(images, start=first_page):
    if grayscale:
      img = img.convert('L')
    buffer = BytesIO()
    img.save(buffer, format="PNG")
    b64 = base64.b64encode(buffer.getvalue()).decode('utf-8')
    if validate:
      try:
        raw = base64.b64decode(b64, validate=True)
        with Image.open(BytesIO(raw)) as im:
          im.verify()
      except Exception as error:
        raise ValueError(f"Imagen base64 inválida en la página {idx}: {error}")
    resultados.append(b64)
  logger.debug(
    f"Rango convertido (count={len(resultados)}, ms={int((time.time()-t0)*1000)})",
    extra={"pdf": pdf_path, "dpi": dpi},
  )
  return resultados

def tiff_to_images_bs64(tiff_path: str, validate: bool = True, pages_range: Optional[Tuple[int, int]] = (2, 3)) -> List[str]:
  """
  Convierte un archivo TIFF multipágina a una lista de imágenes PNG en base64.

  Args:
      tiff_path: Ruta al archivo TIFF.
      validate: Si es True, verifica la integridad de las imágenes generadas.
      pages_range: Tupla (inicio, fin) de páginas a procesar (1-indexed, inclusivo).
                  Use None para procesar todas las páginas.

  Returns:
      Lista de imágenes en base64.
  """
  resultados: List[str] = []
  t0 = time.time()
  logger.info("Convirtiendo TIFF a imágenes", extra={"pdf": tiff_path})
  with Image.open(tiff_path) as im:
    frames = []
    for idx, frame in enumerate(ImageSequence.Iterator(im), start=1):
      if pages_range is not None:
        start, end = pages_range
        # Permitir extremos None para significar “sin límite” por ese lado
        if start is not None and idx < start:
          continue
        if end is not None and idx > end:
          continue
      frames.append(frame.copy())
    for indice, frame in enumerate(frames):
      buffer = BytesIO()
      frame.convert("RGB").save(buffer, format="PNG")
      b64 = base64.b64encode(buffer.getvalue()).decode('utf-8')
      if validate:
        try:
          raw = base64.b64decode(b64, validate=True)
          with Image.open(BytesIO(raw)) as imv:
            imv.verify()
        except Exception as error:
          raise ValueError(f"Imagen base64 inválida en la página TIFF {indice + 1}: {error}")
      resultados.append(b64)
  elapsed_ms = int((time.time() - t0) * 1000)
  logger.debug(f"TIFF->imgs OK (pages={len(resultados)}, ms={elapsed_ms})", extra={"pdf": tiff_path})
  logger.info("TIFF convertido", extra={"pdf": tiff_path})
  return resultados

def document_to_images_bs64(
  path: str,
  validate: bool = True,
  dpi: int = 600,
  thread_count: int = 10,
  tiff_pages: Optional[Tuple[int, int]] = (2, 3),
  tiff_page_range: Optional[Tuple[int, int]] = None,
  poppler_path: str | None = None,
) -> List[str]:
  """
  Función unificada para convertir documentos (PDF o TIFF) a imágenes base64.

  Detecta automáticamente el formato por la extensión del archivo.

  Args:
      path: Ruta al archivo.
      validate: Si es True, valida las imágenes generadas.
      dpi: Resolución para archivos PDF.
      thread_count: Hilos paralelos para archivos PDF.
      tiff_pages: Páginas TIFF por defecto si tiff_page_range es None.
      tiff_page_range: Rango explícito de páginas para TIFF (inicio, fin).
      poppler_path: Ruta personalizada a Poppler para PDFs.

  Returns:
      Lista de imágenes en base64.
  """
  ext = os.path.splitext(path)[1].lower()
  if ext == ".pdf":
    return pdf_to_images_bs64(path, validate=validate, dpi=dpi, thread_count=thread_count, poppler_path=poppler_path)
  if ext in (".tif", ".tiff"):
    effective_range = tiff_pages
    if tiff_page_range is not None:
      effective_range = None if tiff_page_range == (None, None) else tiff_page_range
    return tiff_to_images_bs64(path, validate=validate, pages_range=effective_range)
  raise ValueError(f"Formato no soportado: {ext}")