pdf_split.py

import fitz  # PyMuPDF
import os
import re

def find_pdf_in_directory(directory):
    pdf_files = [f for f in os.listdir(directory) if f.lower().endswith('.pdf')]
    if pdf_files:
        return os.path.join(directory, pdf_files[0])
    else:
        print("No PDF files found in the directory.")
        return None

def extract_chapters_from_toc(doc):
    chapters = []
    toc = doc.get_toc()  # Get the Table of Contents (TOC)
    if not toc:
        print("No TOC found in the PDF.")
        return chapters

    inside_chapters = False  # Flag to detect when we enter the "CHAPTERS" section
    for entry in toc:
        level, title, page = entry

        # Detect start of "CHAPTERS" section
        if title.strip().lower() == "chapters":
            inside_chapters = True
            continue

        if inside_chapters and re.match(r"^\d+$", title.strip()):  # Match simple chapter numbers
            chapters.append((f"Chapter {title.strip()}", page))

    return chapters

def extract_chapters_by_keywords(doc):
    chapters = []
    chapter_pattern = re.compile(r'\bChapter\s+\d+|\b\d+\b', re.IGNORECASE)  # Match "Chapter X" or just numbers
    min_page_gap = 5  # Smaller gap for closely placed chapters
    last_chapter_page = -min_page_gap  # Ensures the first chapter is detected

    for page_num in range(doc.page_count):
        page = doc.load_page(page_num)
        text = page.get_text()

        # Look for chapter headings using regex
        match = chapter_pattern.search(text)
        if match and page_num >= last_chapter_page + min_page_gap:
            chapter_title = match.group(0)
            chapters.append((chapter_title, page_num + 1))  # Store as (title, page)
            last_chapter_page = page_num

    return chapters

def split_pdf_by_chapters(directory):
    pdf_path = find_pdf_in_directory(directory)
    if not pdf_path:
        return

    doc = fitz.open(pdf_path)

    # Attempt to extract chapters using the TOC first
    chapters = extract_chapters_from_toc(doc)
    if not chapters:
        # Fallback to keyword-based chapter extraction
        chapters = extract_chapters_by_keywords(doc)

    if not chapters:
        print("No chapters found in the document.")
        return

    # Extract page ranges for each chapter and save as a separate PDF
    output_dir = os.path.dirname(pdf_path)
    chapter_ranges = list(zip(chapters, chapters[1:]))  # Pair each chapter with the next one for page ranges

    for (current_chapter, current_page), (_, next_page) in chapter_ranges:
        # Adjust page ranges to be zero-indexed for PyMuPDF
        start_page = current_page - 1
        end_page = next_page - 1

        # Extract the page range for the current chapter
        new_pdf = fitz.open()  # Create a new, empty PDF
        for page_num in range(start_page, end_page):
            new_pdf.insert_pdf(doc, from_page=page_num, to_page=page_num)

        # Save the new chapter PDF
        output_path = os.path.join(output_dir, f"{current_chapter.replace(' ', '_')}.pdf")
        new_pdf.save(output_path)
        new_pdf.close()
        print(f"Created: {output_path}")

    # Handle the last chapter
    last_chapter_title, last_chapter_start = chapters[-1]
    last_start_page = last_chapter_start - 1

    # Extract pages from the last chapter to the end of the document
    new_pdf = fitz.open()
    for page_num in range(last_start_page, doc.page_count):
        new_pdf.insert_pdf(doc, from_page=page_num, to_page=page_num)

    # Save the last chapter
    last_output_path = os.path.join(output_dir, f"{last_chapter_title.replace(' ', '_')}.pdf")
    new_pdf.save(last_output_path)
    new_pdf.close()
    print(f"Created: {last_output_path}")

    doc.close()

if __name__ == "__main__":
    current_directory = os.getcwd()
    split_pdf_by_chapters(current_directory)
    print("PDF chapters extracted and saved successfully.")