PitGlobal/PDF_Layout_Inspector

1from pathlib import Path2import sys3 4 5LOCAL_DEPS = Path(__file__).resolve().parent / ".deps"6 7try:8    import cv29    import numpy as np10except ImportError:11    if LOCAL_DEPS.exists():12        sys.path.insert(0, str(LOCAL_DEPS))13    import cv214    import numpy as np15 16 17BASE_DIR = Path(__file__).resolve().parent18 19INPUT_DIR = Path(r"C:\Users\39334\Desktop\George!\2_Milestone\altri_file_da_testare")20OUTPUT_DIR = Path(r"C:\Users\39334\Desktop\George!\2_Milestone\output_pdf")21 22PAGE_INDEX = 023ZOOM = 2.524BAD_IDS = set()25 26 27def render_pdf_page(pdf_path: str | Path, page_index: int = 0, zoom: float = 2.5) -> np.ndarray:28    try:29        import fitz30    except ImportError:31        if LOCAL_DEPS.exists():32            sys.path.insert(0, str(LOCAL_DEPS))33        import fitz34 35    doc = fitz.open(str(pdf_path))36    page = doc[page_index]37    mat = fitz.Matrix(zoom, zoom)38    pix = page.get_pixmap(matrix=mat, alpha=False)39    img = np.frombuffer(pix.samples, dtype=np.uint8).reshape(pix.height, pix.width, pix.n)40    doc.close()41 42    if pix.n == 4:43        img = cv2.cvtColor(img, cv2.COLOR_BGRA2BGR)44    else:45        img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)46 47    return img48 49 50def preprocess_text_mask(image: np.ndarray) -> np.ndarray:51    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)52 53    bw = cv2.adaptiveThreshold(54        gray,55        255,56        cv2.ADAPTIVE_THRESH_GAUSSIAN_C,57        cv2.THRESH_BINARY_INV,58        31,59        15,60    )61 62    horiz_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (60, 1))63    horiz_lines = cv2.morphologyEx(bw, cv2.MORPH_OPEN, horiz_kernel, iterations=1)64    bw = cv2.subtract(bw, horiz_lines)65 66    vert_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (1, 60))67    vert_lines = cv2.morphologyEx(bw, cv2.MORPH_OPEN, vert_kernel, iterations=1)68    bw = cv2.subtract(bw, vert_lines)69 70    connect_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (15, 3))71    bw = cv2.morphologyEx(bw, cv2.MORPH_CLOSE, connect_kernel, iterations=1)72 73    clean_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))74    bw = cv2.morphologyEx(bw, cv2.MORPH_OPEN, clean_kernel, iterations=1)75 76    return bw77 78 79def find_text_line_boxes(mask: np.ndarray) -> list[list[int]]:80    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)81 82    boxes: list[list[int]] = []83    _h_img, w_img = mask.shape84 85    for cnt in contours:86        x, y, w, h = cv2.boundingRect(cnt)87        area = w * h88 89        if area < 20:90            continue91        if w < 20:92            continue93        if h < 6:94            continue95        if h > 120:96            continue97        if w > int(w_img * 0.98):98            continue99 100        boxes.append([x, y, x + w, y + h])101 102    boxes.sort(key=lambda b: (b[1], b[0]))103    return boxes104 105 106def draw_boxes(image: np.ndarray, boxes: list[list[int]], bad_ids: set[int] | None = None) -> np.ndarray:107    out = image.copy()108    bad_ids = bad_ids or set()109 110    for idx, (x1, y1, x2, y2) in enumerate(boxes, start=1):111        color = (0, 0, 255)112        if idx in bad_ids:113            color = (0, 180, 0)114 115        cv2.rectangle(out, (x1, y1), (x2, y2), color, 2)116        cv2.putText(117            out,118            str(idx),119            (x1, max(20, y1 - 5)),120            cv2.FONT_HERSHEY_SIMPLEX,121            0.55,122            color,123            2,124            cv2.LINE_AA,125        )126 127    return out128 129 130def extract_boxes_from_pdf(pdf_path: str | Path, page_index: int = 0, zoom: float = 2.5):131    image = render_pdf_page(pdf_path, page_index=page_index, zoom=zoom)132    mask = preprocess_text_mask(image)133    boxes = find_text_line_boxes(mask)134    return image, mask, boxes135 136 137def process_single_pdf(138    pdf_path: str | Path,139    output_dir: str | Path,140    page_index: int = 0,141    zoom: float = 2.5,142    bad_ids: set[int] | None = None,143) -> None:144    pdf_path = Path(pdf_path)145    output_dir = Path(output_dir)146    output_dir.mkdir(parents=True, exist_ok=True)147 148    image, mask, boxes = extract_boxes_from_pdf(pdf_path, page_index=page_index, zoom=zoom)149    boxed = draw_boxes(image, boxes, bad_ids=bad_ids)150 151    base_name = pdf_path.stem152    boxed_path = output_dir / f"{base_name}_boxes.png"153    mask_path = output_dir / f"{base_name}_mask.png"154 155    cv2.imwrite(str(boxed_path), boxed)156    cv2.imwrite(str(mask_path), mask)157 158    print(f"[OK] {pdf_path.name} -> {boxed_path.name}, {mask_path.name} ({len(boxes)} box)")159 160 161def process_pdf_folder(162    input_dir: str | Path,163    output_dir: str | Path,164    page_index: int = 0,165    zoom: float = 2.5,166    bad_ids: set[int] | None = None,167) -> None:168    input_dir = Path(input_dir)169    output_dir = Path(output_dir)170 171    if not input_dir.exists():172        raise FileNotFoundError(f"Cartella input non trovata: {input_dir}")173 174    output_dir.mkdir(parents=True, exist_ok=True)175 176    pdf_files = sorted(input_dir.glob("*.pdf"))177 178    if not pdf_files:179        print(f"Nessun PDF trovato in: {input_dir}")180        return181 182    for pdf_file in pdf_files:183        try:184            process_single_pdf(185                pdf_path=pdf_file,186                output_dir=output_dir,187                page_index=page_index,188                zoom=zoom,189                bad_ids=bad_ids,190            )191        except Exception as exc:192            print(f"[ERRORE] {pdf_file.name}: {exc}")193 194 195if __name__ == "__main__":196    process_pdf_folder(197        input_dir=INPUT_DIR,198        output_dir=OUTPUT_DIR,199        page_index=PAGE_INDEX,200        zoom=ZOOM,201        bad_ids=BAD_IDS,202    )203