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 PDF a una lista de PNGs en base64.

  Usa pdf2image + Poppler. Lanza ValueError si una imagen no valida.
  """
  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:
  """Devuelve el número de páginas de un PDF usando pdfinfo_from_path."""
  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 [first_page, last_page] (1-indexed, inclusivo) a PNGs base64.

  Opcionalmente aplica escala de grises y omite validación para acelerar.
  """
  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 TIFF multipágina a PNGs en base64.

  Por defecto, procesa un rango de páginas inclusivo cuando se especifica
  `pages_range` como `(inicio, fin)`, usando índices 1-based. Use None para
  procesar todas las páginas. Puede pasar `None` en uno de los extremos para
  indicar “sin inicio” o “sin fin”.
  """
  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]:
  """Puerta de entrada: decide conversión según la extensión (.pdf, .tif/.tiff).

  `tiff_page_range` permite indicar un rango explícito (inicio, fin) 1-based
  para TIFF. Si es None, se usa `tiff_pages` como valor por defecto.
  `poppler_path` se pasa a las funciones PDF; si es None se usa el default del proyecto.
  """
  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}")