trail.py

# sdc.py
import os, re, glob, argparse, json
from statistics import median
from pathlib import Path
from typing import Optional, Sequence

import fitz  # PyMuPDF
from pdf2image import convert_from_path
from PIL import Image
import pytesseract, cv2
import numpy as np
import pandas as pd
from tqdm import tqdm

# ========= DEFAULT CONFIG (you can override via CLI) =========
PDF_DIR_DEFAULT = r"C:\Users\U215438\OneDrive - Trane Technologies\Documents\Python\Yuvaraj_A\EVA_API\PDF"
OUTPUT_DIR_DEFAULT = r"C:\Users\U215438\OneDrive - Trane Technologies\Documents\Python\Yuvaraj_A\EVA_API\output"
TESSERACT_EXE_DEFAULT = r"C:\Users\U215438\AppData\Local\Programs\Tesseract-OCR\tesseract.exe"
POPPLER_BIN_DEFAULT = r"C:\Users\U215438\Downloads\Release-25.07.0-0\poppler-25.07.0\Library\bin"
LANG_DEFAULT = "eng"
# =============================================================

# Enable verbose debug logging to help track failures. Set PDF_DEBUG=0 to silence.
DEBUG = os.getenv("PDF_DEBUG", "1").lower() in {"1", "true", "yes", "on"}


def debug(msg: str):
    """Emit diagnostic messages without disturbing tqdm progress bars."""
    if DEBUG:
        tqdm.write(f"[DEBUG] {msg}")


# Make Pillow safe for huge engineering scans (trusted local files)
Image.MAX_IMAGE_PIXELS = 1_000_000_000

# ------------------- Helpers -------------------
def ensure_dir(p: Path):
    debug(f"Ensuring directory exists: {p}")
    p.mkdir(parents=True, exist_ok=True)

def _validate_pil_image(pil_img: Image.Image, context: str = ""):
    """Raise ValueError if the PIL image is empty (width/height == 0)."""
    if pil_img is None:
        debug(f"PIL validation failed: image is None ({context})")
        raise ValueError(f"Empty image supplied{': ' + context if context else ''}")
    w, h = getattr(pil_img, "width", 0), getattr(pil_img, "height", 0)
    if not w or not h:
        debug(f"PIL validation failed: zero-sized image ({context})")
        raise ValueError(f"Zero-sized image encountered{': ' + context if context else ''}")


def pil2cv(img: Image.Image, context: str = ""):
    debug(f"Converting PIL -> CV2 array ({context}) size={getattr(img, 'size', None)} mode={getattr(img, 'mode', None)}")
    _validate_pil_image(img, context)
    if img.mode != "RGB":
        img = img.convert("RGB")
    arr = np.array(img)
    if arr.size == 0:
        debug(f"PIL -> CV2 conversion produced empty array ({context})")
        raise ValueError(f"Image data empty after conversion{': ' + context if context else ''}")
    return cv2.cvtColor(arr, cv2.COLOR_RGB2BGR)


def cv2pil(img):
    if img is None or img.size == 0:
        debug("CV2 -> PIL conversion requested on empty image")
        raise ValueError("Empty OpenCV image supplied")
    debug(f"Converting CV2 -> PIL image with shape={getattr(img, 'shape', None)}")
    return Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))

def preprocess_for_ocr(pil_img: Image.Image, context: str = "") -> Image.Image:
    """General cleaning that works well on scans/drawings/tables."""
    try:
        debug(f"Preprocessing image for OCR ({context})")
        cv = pil2cv(pil_img, context=context)
        gray = cv2.cvtColor(cv, cv2.COLOR_BGR2GRAY)
        gray = cv2.fastNlMeansDenoising(gray, h=10)
        th = cv2.adaptiveThreshold(gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
                                   cv2.THRESH_BINARY, 35, 15)
        result = cv2pil(cv2.cvtColor(th, cv2.COLOR_GRAY2BGR))
        debug(f"Preprocessing complete ({context}) result_size={result.size}")
        return result
    except (ValueError, cv2.error) as exc:
        msg_ctx = f" ({context})" if context else ""
        print(f"[WARN] Skipping preprocessing{msg_ctx}: {exc}")
        return pil_img.convert("RGB")

def autorotate(pil_img: Image.Image) -> Image.Image:
    """Use Tesseract OSD to correct rotation when OCR’ing."""
    try:
        debug(f"Attempting autorotate size={pil_img.size}")
        osd = pytesseract.image_to_osd(pil_img)
        m = re.search(r"Rotate:\s+(\d+)", osd)
        angle = int(m.group(1)) if m else 0
        debug(f"Autorotate detected angle={angle}")
        return pil_img.rotate(-angle, expand=True, fillcolor="white") if angle else pil_img
    except Exception:
        debug("Autorotate failed; returning original image")
        return pil_img

def page_profile(page: fitz.Page):
    """
    Inspect a page with PyMuPDF:
      has_text: bool (nontrivial text layer?)
      img_coverage: fraction of page covered by image blocks (0..1)
      median_font_pt: median font size (pt) if any text spans, else None
    """
    w, h = page.rect.width, page.rect.height
    area = max(w * h, 1.0)
    rd = page.get_text("rawdict")

    text_chars = 0
    font_sizes = []
    img_area = 0.0

    for block in rd["blocks"]:
        btype = block.get("type", 0)
        bbox = block.get("bbox", (0,0,0,0))
        bw = max(0.0, bbox[2] - bbox[0])
        bh = max(0.0, bbox[3] - bbox[1])

        if btype == 0:  # text
            for line in block.get("lines", []):
                for span in line.get("spans", []):
                    text_chars += len(span.get("text", ""))
                    fs = span.get("size")
                    if fs:
                        font_sizes.append(fs)
        elif btype == 1:  # image
            img_area += (bw * bh)

    has_text = text_chars >= 20
    med_font = median(font_sizes) if font_sizes else None
    coverage = img_area / area
    debug(
        "Page profile computed: "
        f"has_text={has_text} text_chars={text_chars} img_coverage={coverage:.3f} "
        f"median_font={med_font}"
    )
    return has_text, coverage, med_font

def choose_dpi(has_text: bool, img_cov: float, median_pt: float | None) -> int:
    """Adaptive DPI policy per page."""
    if not has_text or img_cov >= 0.90:
        dpi = 400  # image-only / dominant scan
    elif median_pt is not None and median_pt < 8.5:
        dpi = 350  # tiny fonts
    elif median_pt is not None and median_pt > 13:
        dpi = 250  # big fonts
    else:
        dpi = 300  # normal text
    debug(
        "choose_dpi result: "
        f"has_text={has_text} img_cov={img_cov:.3f} median_pt={median_pt} -> dpi={dpi}"
    )
    return dpi

def ocr_image(pil_img: Image.Image, lang: str, table_mode=False, context: str = "") -> str:
    debug(
        f"Starting OCR ({context}) table_mode={table_mode} lang={lang} size={pil_img.size}"
    )
    pil_img = autorotate(pil_img)
    pil_img = preprocess_for_ocr(pil_img, context=context or "ocr_image")
    if table_mode:
        cfg = r'--oem 3 --psm 6 -c tessedit_char_whitelist=ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789./-()+°%,±Øø'
    else:
        cfg = r'--oem 3 --psm 6'
    txt = pytesseract.image_to_string(pil_img, lang=lang, config=cfg)
    debug(f"OCR complete ({context}) chars={len(txt)}")
    return txt

# ------ Basic grid/table detector (best effort) ------
def extract_table_to_csv(pil_img: Image.Image, lang: str, out_csv: Path, context: str = "") -> tuple[bool, list[list[str]]]:
    """Detect grid lines, OCR cells, and write CSV.

    Returns
    -------
    tuple[bool, list[list[str]]]
        A flag indicating whether a table was detected along with the table
        contents (each inner list represents one row).
    """
    ensure_dir(out_csv.parent)
    ctx = context or f"table:{out_csv.name}"
    debug(f"Attempting table extraction ({ctx}) -> {out_csv}")
    processed = preprocess_for_ocr(pil_img, context=ctx)
    cv = pil2cv(processed, context=ctx)
    gray = cv2.cvtColor(cv, cv2.COLOR_BGR2GRAY)
    inv = 255 - gray

    H, W = inv.shape
    debug(f"Table candidate image size=({W},{H})")
    horiz_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (max(10, W // 60), 1))
    vert_kernel  = cv2.getStructuringElement(cv2.MORPH_RECT, (1, max(10, H // 40)))

    horiz = cv2.dilate(cv2.erode(inv, horiz_kernel, 1), horiz_kernel, 1)
    vert  = cv2.dilate(cv2.erode(inv, vert_kernel, 1),  vert_kernel, 1)
    grid  = cv2.add(horiz, vert)

    contours, _ = cv2.findContours(grid, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
    debug(f"Found {len(contours)} contours for table detection ({ctx})")
    boxes = []
    for c in contours:
        x,y,w,h = cv2.boundingRect(c)
        if w <= 1 or h <= 1:
            # Reject degenerate boxes that would rasterize to empty crops later
            continue
        if w*h < 200:       # skip tiny
            continue
        if w > W*0.98 and h > H*0.15:  # skip large outer borders
            continue
        boxes.append((x,y,w,h))
    debug(f"Filtered to {len(boxes)} potential table cells ({ctx})")
    if not boxes:
        debug(f"No valid boxes detected for table ({ctx})")
        return False, []

    boxes = sorted(boxes, key=lambda b: (b[1], b[0]))
    # group into rows
    rows, current, row_y = [], [boxes[0]], boxes[0][1]
    row_tol = max(10, H // 100)
    for b in boxes[1:]:
        if abs(b[1] - row_y) <= row_tol:
            current.append(b)
        else:
            rows.append(sorted(current, key=lambda x: x[0]))
            current, row_y = [b], b[1]
    rows.append(sorted(current, key=lambda x: x[0]))

    col_count = max(len(r) for r in rows)
    cfg_whitelist = r'--oem 3 --psm 7 -c tessedit_char_whitelist=ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789./-()+°%±Øø,'
    table = []
    for row_index, r in enumerate(rows):
        if len(r) < col_count:
            r = r + [r[-1]]*(col_count - len(r))
        row_text = []
        for col_index, (x,y,w,h) in enumerate(r[:col_count]):
            pad = 2
            x0,y0,x1,y1 = max(0,x+pad), max(0,y+pad), min(W,x+w-pad), min(H,y+h-pad)
            if x1 <= x0 or y1 <= y0:
                # Degenerate crop after padding – treat as empty cell
                row_text.append("")
                debug(f"Empty crop for cell ({row_index},{col_index}) in {ctx}")
                continue
            cell_img = cv[y0:y1, x0:x1]
            if cell_img.size == 0:
                row_text.append("")
                debug(f"Zero-sized cell ({row_index},{col_index}) in {ctx}")
                continue
            cell = cv2pil(cell_img)
            txt = pytesseract.image_to_string(cell, lang=lang, config=cfg_whitelist)
            row_text.append(txt.strip().replace("\n", " "))
        table.append(row_text)

    pd.DataFrame(table).to_csv(out_csv, index=False, header=False)
    debug(f"Table extraction complete ({ctx}) rows={len(table)} cols={col_count}")
    return True, table

# ------------------- Per-page runner -------------------
def render_pdf_page(pdf_path: Path, page_index: int, dpi: int, poppler_bin: str | None) -> Image.Image | None:
    """Render a single PDF page into a PIL image, handling empty results safely."""
    try:
        debug(
            f"Rendering page {page_index+1} of {pdf_path.name} at {dpi} dpi with poppler={poppler_bin}"
        )
        images = convert_from_path(
            str(pdf_path),
            first_page=page_index + 1,
            last_page=page_index + 1,
            dpi=dpi,
            poppler_path=poppler_bin,
        )
    except Exception as exc:
        print(f"[ERROR] {pdf_path.name} page {page_index+1}: convert_from_path failed ({exc})")
        return None

    if not images:
        print(f"[WARN] {pdf_path.name} page {page_index+1}: no rasterized images returned")
        return None

    img = images[0]
    debug(
        f"Rendered page {page_index+1} of {pdf_path.name}: image_size={getattr(img, 'size', None)}"
    )
    try:
        _validate_pil_image(img, context=f"{pdf_path.name}#p{page_index+1}")
    except ValueError as exc:
        print(f"[WARN] {pdf_path.name} page {page_index+1}: {exc}")
        return None
    return img


def process_page(doc, pdf_path: Path, page_index: int, out_dir: Path, lang: str,
                 poppler_bin: str | None, tiles=(3,3), tile_overlap=0.04):
    debug(f"Processing {pdf_path.name} page {page_index+1}")
    page = doc[page_index]
    has_text, img_cov, med_pt = page_profile(page)
    dpi = choose_dpi(has_text, img_cov, med_pt)
    debug(
        f"Page {page_index+1}: has_text={has_text} img_cov={img_cov:.3f} "
        f"median_font={med_pt} dpi={dpi}"
    )

    page_txt = out_dir / f"page_{page_index+1:03d}.txt"
    page_csv = out_dir / f"page_{page_index+1:03d}_table.csv"

    page_info: dict[str, object] = {
        "page_number": page_index + 1,
        "dpi": dpi,
        "image_coverage": img_cov,
        "mode": None,
        "raw_text": None,
        "ocr_text": None,
        "table": None,
        "notes": [],
    }

    if has_text and img_cov < 0.90:
        # Prefer native text; OCR image regions if worthwhile
        txt_native = page.get_text("raw").strip()
        debug(
            f"Page {page_index+1}: using native text (len={len(txt_native)}) with img_cov={img_cov:.3f}"
        )
        if img_cov > 0.15:
            pil = render_pdf_page(pdf_path, page_index, dpi, poppler_bin)
            if pil is None:
                print(f"[WARN] Skipping OCR image regions for {pdf_path.name} page {page_index+1}: could not render image")
                page_txt.write_text(txt_native, encoding="utf-8")
                page_info["mode"] = "native"
                page_info["raw_text"] = txt_native
                page_info["notes"].append("image-region-render-failed")
                return dpi, "native", img_cov, page_info
            W,H = pil.size
            debug(f"Page {page_index+1}: rendered native+image region size=({W},{H})")
            img_region = pil.crop((0, int(H*0.4), W, H))
            made_csv, table_rows = extract_table_to_csv(img_region, lang, page_csv,
                                                        context=f"{pdf_path.name}#p{page_index+1}-img")
            if made_csv:
                debug(f"Page {page_index+1}: table detected in image region rows={len(table_rows)}")
                page_info["table"] = {
                    "csv": page_csv.name,
                    "rows": table_rows,
                }
            else:
                ocr_txt = ocr_image(img_region, lang, table_mode=True,
                                    context=f"{pdf_path.name}#p{page_index+1}-img").strip()
                debug(f"Page {page_index+1}: OCR image region chars={len(ocr_txt)}")
                page_info["ocr_text"] = ocr_txt
        page_txt.write_text(txt_native, encoding="utf-8")
        page_info["mode"] = "native+ocr" if img_cov > 0.15 else "native"
        page_info["raw_text"] = txt_native
        debug(f"Page {page_index+1}: finished native branch mode={page_info['mode']}")
        return dpi, page_info["mode"], img_cov, page_info

    # Image-only / dominant → OCR with tiling + try table CSV
    pil = render_pdf_page(pdf_path, page_index, dpi, poppler_bin)
    if pil is None:
        print(f"[ERROR] {pdf_path.name} page {page_index+1}: failed to rasterize for OCR")
        page_info["mode"] = "render-error"
        page_info["notes"].append("render-failed")
        return dpi, "render-error", img_cov, page_info
    pil = autorotate(pil)
    pil_clean = preprocess_for_ocr(pil, context=f"{pdf_path.name}#p{page_index+1}")
    debug(f"Page {page_index+1}: cleaned image size={pil_clean.size}")

    made_csv, table_rows = extract_table_to_csv(pil_clean, lang, page_csv,
                                                context=f"{pdf_path.name}#p{page_index+1}")
    if made_csv:
        debug(f"Page {page_index+1}: table detected rows={len(table_rows)}")
        page_info["table"] = {
            "csv": page_csv.name,
            "rows": table_rows,
        }

    W, H = pil_clean.size
    nx, ny = tiles
    tw, th = int(W/nx), int(H/ny)
    dx, dy = int(tw*tile_overlap), int(th*tile_overlap)

    tile_texts = []
    debug(
        f"Page {page_index+1}: tiling image W={W} H={H} nx={nx} ny={ny} overlap={tile_overlap}"
    )
    with tqdm(total=nx*ny, desc=f"  tiles@p{page_index+1}", leave=False) as pbar_tiles:
        for iy in range(ny):
            for ix in range(nx):
                x0 = max(0, ix*tw - dx); y0 = max(0, iy*th - dy)
                x1 = min(W, (ix+1)*tw + dx); y1 = min(H, (iy+1)*th + dy)
                tile = pil_clean.crop((x0,y0,x1,y1))
                cfg = r'--oem 3 --psm 6 -c tessedit_char_whitelist=ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789./-()+°%,±Øø'
                txt = pytesseract.image_to_string(tile, lang=lang, config=cfg)
                tile_texts.append(f"[tile {ix},{iy}]\n{txt.strip()}\n")
                debug(
                    f"Page {page_index+1}: tile ({ix},{iy}) crop=({x0},{y0},{x1},{y1}) chars={len(txt)}"
                )
                pbar_tiles.update(1)

    ocr_text = "\n".join(tile_texts)
    page_txt.write_text(ocr_text, encoding="utf-8")
    page_info["mode"] = "ocr-tiled"
    page_info["ocr_text"] = ocr_text.strip()
    page_info["raw_text"] = None
    debug(f"Page {page_index+1}: finished OCR tiled branch chars={len(page_info['ocr_text'])}")
    return dpi, "ocr-tiled", img_cov, page_info

# ------------------- PDF runner -------------------
def process_pdf(pdf_path: Path, out_root: Path, lang: str, poppler_bin: str | None):
    out_dir = out_root / pdf_path.stem
    ensure_dir(out_dir)
    debug(f"Processing PDF: {pdf_path} -> {out_dir}")
    page_summaries = []
    page_infos = []

    try:
        with fitz.open(str(pdf_path)) as doc:
            debug(f"Opened PDF {pdf_path.name} pages={len(doc)}")
            if len(doc) == 0:
                print(f"[WARN] Empty PDF: {pdf_path.name}")
                return
            pages_iter = tqdm(range(len(doc)), desc=f"pages: {pdf_path.name}", leave=False)
            for p in pages_iter:
                try:
                    dpi, mode, imgcov, page_info = process_page(doc, pdf_path, p, out_dir, lang, poppler_bin)
                    page_summaries.append((p+1, dpi, mode, imgcov))
                    page_infos.append(page_info)
                    imgcov_display = f"{imgcov:.3f}" if imgcov is not None else "n/a"
                    debug(
                        f"Page {p+1} summary: dpi={dpi} mode={mode} img_cov={imgcov_display}"
                    )
                except Exception as e:
                    (out_dir / "errors.log").open("a", encoding="utf-8").write(f"page {p+1}: {e}\n")
                    print(f"[ERROR] {pdf_path.name} page {p+1}: {e}")
                    page_infos.append({
                        "page_number": p + 1,
                        "dpi": None,
                        "image_coverage": None,
                        "mode": "error",
                        "raw_text": None,
                        "ocr_text": None,
                        "table": None,
                        "notes": [f"exception: {e}"],
                    })
    except Exception as e:
        print(f"[FATAL] Could not open {pdf_path}: {e}")
        return

    # Combine text & save manifest
    combined = []
    debug(f"Combining per-page text files for {pdf_path.name}")
    for p, _, _, _ in page_summaries:
        pt = out_dir / f"page_{p:03d}.txt"
        if pt.exists():
            combined.append(f"\n===== PAGE {p} =====\n")
            combined.append(pt.read_text(encoding="utf-8", errors="ignore"))
    if combined:
        (out_dir / "combined.txt").write_text("".join(combined), encoding="utf-8")

    pd.DataFrame(page_summaries, columns=["page", "dpi", "mode", "image_coverage"])\
      .to_csv(out_dir / "manifest.csv", index=False)
    debug(f"Wrote manifest for {pdf_path.name} with {len(page_summaries)} entries")

    storage_doc = {
        "document": pdf_path.name,
        "pages": page_infos,
    }
    (out_dir / "storage.json").write_text(
        json.dumps(storage_doc, ensure_ascii=False, indent=2),
        encoding="utf-8",
    )
    debug(f"Wrote storage.json for {pdf_path.name}")

# ------------------- CLI / Main -------------------
def parse_args(argv: Optional[Sequence[str]] = None):
    """Parse CLI args but gracefully ignore stray ones from notebook launchers."""

    ap = argparse.ArgumentParser(description="Smart PDF text extractor (adaptive OCR).")
    ap.add_argument("--pdf-dir", default=PDF_DIR_DEFAULT, help="Folder containing PDFs.")
    ap.add_argument("--output-dir", default=OUTPUT_DIR_DEFAULT, help="Output folder.")
    ap.add_argument("--tesseract", default=TESSERACT_EXE_DEFAULT, help="Path to tesseract.exe on Windows.")
    ap.add_argument("--poppler", default=POPPLER_BIN_DEFAULT, help="Path to Poppler bin (pdfinfo.exe, pdftoppm.exe).")
    ap.add_argument("--lang", default=LANG_DEFAULT, help="Tesseract languages, e.g. 'eng' or 'eng+deu'.")

    args, unknown = ap.parse_known_args(argv)
    if unknown:
        print(f"[WARN] Ignoring unrecognized arguments: {unknown}")
    return args

def main():
    args = parse_args()

    pytesseract.pytesseract.tesseract_cmd = args.tesseract
    poppler_bin = args.poppler if args.poppler and len(args.poppler.strip()) else None

    pdf_dir = Path(args.pdf_dir).resolve()
    out_dir = Path(args.output_dir).resolve()
    ensure_dir(out_dir)

    print(f"[INFO] cwd: {Path.cwd()}")
    print(f"[INFO] PDF_DIR: {pdf_dir}")
    print(f"[INFO] OUTPUT_DIR: {out_dir}")
    print(f"[INFO] POPPLER_BIN: {poppler_bin}")
    print(f"[INFO] Tesseract: {pytesseract.pytesseract.tesseract_cmd}")
    print(f"[INFO] Debug logging enabled: {DEBUG}")

    # List contents (helps with OneDrive placeholders)
    try:
        names = [p.name for p in pdf_dir.iterdir()]
        print(f"[INFO] PDF_DIR exists: {pdf_dir.exists()}  items: {len(names)}  sample: {names[:10]}")
    except Exception as e:
        print(f"[WARN] Could not list PDF_DIR: {e}")

    pdf_candidates = sorted(glob.glob(str(pdf_dir / "*.pdf"))) + sorted(glob.glob(str(pdf_dir / "*.PDF")))
    # Use dict.fromkeys to preserve order while removing duplicates (Windows is case-insensitive)
    pdf_files = list(dict.fromkeys(pdf_candidates))
    print(f"[INFO] PDFs found: {len(pdf_files)} -> {[Path(p).name for p in pdf_files]}")

    if not pdf_files:
        print("[HINT] Put PDFs in the folder above OR pass --pdf-dir FULL_PATH.")
        print("[HINT] In OneDrive, right-click PDFs → 'Always keep on this device'.")
        return

    for pdf in tqdm(pdf_files, desc="PDFs"):
        process_pdf(Path(pdf), out_dir, args.lang, poppler_bin)

    print(f"\nDone. See results in: {out_dir}")

if __name__ == "__main__":
    main()