PitGlobal/PDF_Layout_Inspector

1from __future__ import annotations2 3import argparse4import csv5import json6from pathlib import Path7import sys8 9 10LOCAL_DEPS = Path(__file__).resolve().parent / ".deps"11 12try:13    import cv214    import numpy as np15except ImportError:16    if LOCAL_DEPS.exists():17        sys.path.insert(0, str(LOCAL_DEPS))18    import cv219    import numpy as np20 21from box_source import draw_boxes, extract_boxes_from_pdf22from features import crop_box, preprocess_text_mask_fine, structural_embedding23from model import SimilarityModel, load_model_artifact24 25INPUT_PDF_PATH = Path(r"C:\Users\39334\Desktop\George!\2_Milestone\altri_file_da_testare\book3.pdf")26 27ROOT_DIR = Path(__file__).resolve().parent28ARTIFACTS_DIR = ROOT_DIR / "artifacts"29MODEL_INPUT = ARTIFACTS_DIR / "dataset" / "embeddings_dataset.npz"30MODEL_ARTIFACT = ARTIFACTS_DIR / "models" / "logistic_regression_model.npz"31PRED_DIR = ARTIFACTS_DIR / "predictions"32 33 34def load_model(npz_path: Path, model_path: Path | None = None):35    if model_path and model_path.exists():36        return load_model_artifact(model_path)37 38    data = np.load(npz_path, allow_pickle=True)39    embeddings = data["embeddings"].astype(np.float32)40    labels = data["labels"].astype(np.int32)41    return SimilarityModel.from_arrays(embeddings, labels)42 43 44def predict_pdf(pdf_path: str | Path, out_dir: str | Path | None = None):45    pdf_path = Path(pdf_path)46    out_dir = Path(out_dir) if out_dir else (PRED_DIR / pdf_path.stem)47    out_dir.mkdir(parents=True, exist_ok=True)48 49    image, _mask, boxes = extract_boxes_from_pdf(pdf_path)50    fine_mask = preprocess_text_mask_fine(image)51    model = load_model(MODEL_INPUT, MODEL_ARTIFACT)52 53    predictions = []54    bad_ids: set[int] = set()55 56    for idx, box in enumerate(boxes, start=1):57        mask_crop = crop_box(fine_mask, box, pad=2)58        embedding, debug = structural_embedding(mask_crop)59        scores = model.predict_scores(embedding)60 61        if scores["label"] == "bad_overlap":62            bad_ids.add(idx)63 64        x1, y1, x2, y2 = box65        predictions.append(66            {67                "line_id": idx,68                "x1": x1,69                "y1": y1,70                "x2": x2,71                "y2": y2,72                "w": x2 - x1,73                "h": y2 - y1,74                "predicted_label": scores["label"],75                "score_good": round(float(scores["score_good"]), 6),76                "score_bad": round(float(scores["score_bad"]), 6),77                "margin": round(float(scores["margin"]), 6),78                "band_count": round(float(debug["band_count"]), 4),79                "band_gap": round(float(debug["band_gap"]), 4),80                "char_count": round(float(debug["char_count"]), 4),81                "tall_char_count": round(float(debug["tall_char_count"]), 4),82                "wide_char_count": round(float(debug["wide_char_count"]), 4),83                "max_cy_gap": round(float(debug["max_cy_gap"]), 4),84            }85        )86 87    result = draw_boxes(image, boxes, bad_ids=bad_ids)88    cv2.imwrite(str(out_dir / f"{pdf_path.stem}_prediction.png"), result)89 90    with open(out_dir / f"{pdf_path.stem}_prediction.json", "w", encoding="utf-8") as f:91        json.dump(predictions, f, indent=2)92 93    with open(out_dir / f"{pdf_path.stem}_prediction.csv", "w", newline="", encoding="utf-8") as f:94        writer = csv.DictWriter(f, fieldnames=list(predictions[0].keys()) if predictions else ["line_id"])95        writer.writeheader()96        writer.writerows(predictions)97 98    print(f"[OK] PDF: {pdf_path}")99    print(f"[OK] Box totali: {len(boxes)}")100    print(f"[OK] Box predetti bad_overlap: {sorted(bad_ids)}")101    print(f"[OK] Output dir: {out_dir}")102    return predictions103 104 105def main():106    parser = argparse.ArgumentParser(description="Predice i box con sovrapposizione su un PDF.")107    parser.add_argument(108        "pdf_path",109        nargs="?",110        default=str(INPUT_PDF_PATH),111        help=f"Percorso del PDF da analizzare (default: {INPUT_PDF_PATH})",112    )113    parser.add_argument("--out-dir", default=None, help="Cartella di output")114    args = parser.parse_args()115    predict_pdf(args.pdf_path, args.out_dir)116 117 118if __name__ == "__main__":119    main()120