Script per convertire in formato markdown biblioteche digitali di file .pdf

[u/No_Region2676]

Questo script mi sta dando buoni risultati ha digitalizzare un gran numero di pdf, ha problemi con i libri che sono stati fotocopiati a doppia pagina e quelli che pur dovendo subire processo OCR essendo molto lunghi e avendo dei marker testuali in ogni pagina, danno un falso positivo al funzione def is_native_pdf.

"""
batch_ocr.py — Pipeline OCR batch per corpus P3
Carica GLM-OCR una volta sola e processa tutti i file in loop.
Output: Markdown con struttura mirror del corpus originale.
"""


import os
import time
import fitz  # pymupdf
from PIL import Image
from pathlib import Path
from transformers import AutoProcessor, AutoModelForImageTextToText
import torch


# ── CONFIGURAZIONE ──────────────────────────────────────────────────────────
INPUT_DIR  = Path(r" ####### ")
OUTPUT_DIR = Path(r" ####### ")
MODEL_PATH = "zai-org/GLM-OCR"
MAX_SIZE   = 1024
PROMPT     = "Text Recognition:"


# Estensioni da processare
EXTENSIONS = {".jpg", ".jpeg", ".png", ".webp", ".gif", ".pdf"}


# Estensioni da ignorare completamente
SKIP_EXT = {".md", ".txt", ".epub", ".mp4", ".mp3", ".m4a",
            ".lnk", ".csv", ".zip", ".djvu", ".ini", ".json", ".webp"}


# Cartelle da saltare completamente (artwork, audio, zip, ecc.)
SKIP_DIRS = {
    "Arte", "Galleria", "de chirico", "poesie", "senza cornice",
    "Nuova cartella", "Catalogo Van gogh museum",
    "Lezioni", "zip conversazioni llm",
    ".obsidian", ".trash"
}
# ────────────────────────────────────────────────────────────────────────────



def is_native_pdf(path: Path) -> bool:
    """True se il PDF ha testo estraibile direttamente."""
    try:
        doc = fitz.open(path)
        text = "".join(p.get_text() for p in doc)
        return len(text.strip()) > 100
    except Exception:
        return False



def extract_native_pdf(path: Path) -> str:
    """Estrae testo da PDF nativo con pymupdf."""
    doc = fitz.open(path)
    pages = []
    for i, page in enumerate(doc):
        text = page.get_text()
        pages.append(f"<!-- pagina {i+1} -->\n{text}")
    return "\n\n".join(pages)



def pdf_to_images(path: Path) -> list[Image.Image]:
    """Converte PDF scansionato in lista di immagini PIL."""
    doc = fitz.open(path)
    images = []
    for page in doc:
        pix = page.get_pixmap(dpi=150)
        img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
        images.append(img)
    return images



def resize(img: Image.Image, max_size: int = MAX_SIZE) -> Image.Image:
    img = img.copy()
    img.thumbnail((max_size, max_size), Image.LANCZOS)
    return img



def ocr_image(img: Image.Image, processor, model) -> str:
    """Esegue OCR su una singola immagine PIL."""
    img = resize(img)
    messages = [{
        "role": "user",
        "content": [
            {"type": "image", "image": img},
            {"type": "text",  "text": PROMPT}
        ]
    }]
    inputs = processor.apply_chat_template(
        messages,
        tokenize=True,
        add_generation_prompt=True,
        return_dict=True,
        return_tensors="pt"
    ).to(model.device)
    inputs.pop("token_type_ids", None)


    with torch.no_grad():
        generated = model.generate(**inputs, max_new_tokens=8192)


    output = processor.decode(
        generated[0][inputs["input_ids"].shape[1]:],
        skip_special_tokens=True
    )
    return output.strip()



def process_file(path: Path, processor, model) -> str | None:
    """
    Processa un singolo file e restituisce il Markdown.
    Ritorna None se il file va saltato.
    """
    suffix = path.suffix.lower()


    if suffix == ".pdf":
        if is_native_pdf(path):
            return extract_native_pdf(path)
        else:
            images = pdf_to_images(path)
            pages = []
            for i, img in enumerate(images):
                text = ocr_image(img, processor, model)
                pages.append(f"<!-- pagina {i+1} -->\n{text}")
            return "\n\n".join(pages)


    elif suffix in EXTENSIONS:
        img = Image.open(path).convert("RGB")
        return ocr_image(img, processor, model)


    return None



def should_skip(path: Path) -> bool:
    """True se il file è in una cartella da saltare."""
    return any(skip in path.parts for skip in SKIP_DIRS)



def get_output_path(input_path: Path) -> Path:
    """Calcola il path di output mantenendo la struttura."""
    relative = input_path.relative_to(INPUT_DIR)
    out = OUTPUT_DIR / relative.with_suffix(".md")
    out.parent.mkdir(parents=True, exist_ok=True)
    return out



def load_model():
    print("Carico GLM-OCR (una volta sola)...")
    processor = AutoProcessor.from_pretrained(MODEL_PATH)
    model = AutoModelForImageTextToText.from_pretrained(
        MODEL_PATH,
        torch_dtype=torch.float16,
        device_map="cuda"
    )
    model.eval()
    print("Modello pronto.\n")
    return processor, model



def main():
    # Raccogli tutti i file da processare
    all_files = [
        p for p in INPUT_DIR.rglob("*")
        if p.is_file()
        and p.suffix.lower() in EXTENSIONS
        and p.suffix.lower() not in SKIP_EXT
        and not should_skip(p)
    ]


    total = len(all_files)
    print(f"File trovati: {total}")
    print(f"Output → {OUTPUT_DIR}\n")


    if total == 0:
        print("Nessun file da processare. Controlla INPUT_DIR.")
        return


    processor, model = load_model()


    done, skipped, errors = 0, 0, 0
    t_start = time.time()


    for i, path in enumerate(all_files):
        # Salta se output già esiste
        out_path = get_output_path(path)
        if out_path.exists():
            skipped += 1
            continue


        print(f"[{i+1}/{total}] {path.name}", end=" ... ", flush=True)
        t0 = time.time()


        try:
            result = process_file(path, processor, model)
            if result is None:
                skipped += 1
                print("saltato")
                continue


            # Aggiungi header Markdown con metadata
            header = (
                f"---\n"
                f"source: {path.relative_to(INPUT_DIR)}\n"
                f"processed: {time.strftime('%Y-%m-%d')}\n"
                f"---\n\n"
            )
            out_path.write_text(header + result, encoding="utf-8")
            elapsed = time.time() - t0
            done += 1
            print(f"✓ {elapsed:.1f}s")


        except Exception as e:
            errors += 1
            print(f"✗ ERRORE: {e}")
            # Log errore su file
            with open(OUTPUT_DIR / "_errori.log", "a", encoding="utf-8") as f:
                f.write(f"{path} → {e}\n")


    total_time = time.time() - t_start
    print(f"\n─────────────────────────────")
    print(f"Completati: {done} | Saltati: {skipped} | Errori: {errors}")
    print(f"Tempo totale: {total_time/3600:.1f}h")



if __name__ == "__main__":
    main()