|  |     | 
| (2 intermediate revisions by the same user not shown) | 
| Line 1: | Line 1: | 
|  | import fitz  # PyMuPDF
 |  | Kotoba | 
|  | import pdfplumber
 |  | 
|  | import re
 |  | 
|  | import pytesseract
 |  | 
|  | import cv2
 |  | 
|  | import numpy as np
 |  | 
|  | from transformers import VisionEncoderDecoderModel, ViTImageProcessor, AutoTokenizer
 |  | 
|  | import torch
 |  | 
|  | from PIL import Image
 |  | 
|  | import os
 |  | 
|  | import logging
 |  | 
|  | import traceback
 |  | 
|  | import warnings
 |  | 
|  | from pathlib import Path
 |  | 
|  | from abc import ABC, abstractmethod
 |  | 
|  | import argparse
 |  | 
|  |   |  | 
|  | warnings.filterwarnings("ignore")
 |  | 
|  |   |  | 
|  | class PDFExtractor(ABC):
 |  | 
|  |     """Abstract base class for PDF extraction."""
 |  | 
|  |   |  | 
|  |     def __init__(self, pdf_path):
 |  | 
|  |         self.pdf_path = pdf_path
 |  | 
|  |         self.setup_logging()
 |  | 
|  |   |  | 
|  |     def setup_logging(self):
 |  | 
|  |         """Set up logging configuration."""
 |  | 
|  |         log_file = f"{Path(__file__).stem}.log"
 |  | 
|  |         logging.basicConfig(
 |  | 
|  |             level=logging.INFO,
 |  | 
|  |             format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
 |  | 
|  |             handlers=[
 |  | 
|  |                 logging.FileHandler(log_file),
 |  | 
|  |                 logging.StreamHandler()
 |  | 
|  |             ]
 |  | 
|  |         )
 |  | 
|  |         self.logger = logging.getLogger(__name__)
 |  | 
|  |   |  | 
|  |     @abstractmethod
 |  | 
|  |     def extract(self):
 |  | 
|  |         """Abstract method for extracting content from PDF."""
 |  | 
|  |         pass
 |  | 
|  |   |  | 
|  | class MarkdownPDFExtractor(PDFExtractor):
 |  | 
|  |     """Class for extracting markdown-formatted content from PDF."""
 |  | 
|  |   |  | 
|  |     BULLET_POINTS = '•◦▪▫●○'
 |  | 
|  |   |  | 
|  |     def __init__(self, pdf_path):
 |  | 
|  |         super().__init__(pdf_path)
 |  | 
|  |         self.pdf_filename = Path(pdf_path).stem
 |  | 
|  |         self.setup_image_captioning()
 |  | 
|  |   |  | 
|  |     def setup_image_captioning(self):
 |  | 
|  |         """Set up the image captioning model."""
 |  | 
|  |         try:
 |  | 
|  |             self.model = VisionEncoderDecoderModel.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
 |  | 
|  |             self.feature_extractor = ViTImageProcessor.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
 |  | 
|  |             self.tokenizer = AutoTokenizer.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
 |  | 
|  |             self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 |  | 
|  |             self.model.to(self.device)
 |  | 
|  |             self.logger.info("Image captioning model set up successfully.")
 |  | 
|  |         except Exception as e:
 |  | 
|  |             self.logger.error(f"Error setting up image captioning model: {e}")
 |  | 
|  |             self.logger.exception(traceback.format_exc())
 |  | 
|  |   |  | 
|  |     def extract(self):
 |  | 
|  |         try:
 |  | 
|  |             markdown_content, markdown_pages = self.extract_markdown()
 |  | 
|  |             self.save_markdown(markdown_content)
 |  | 
|  |             self.logger.info(f"Markdown content has been saved to outputs/{self.pdf_filename}.md")
 |  | 
|  |             return markdown_pages
 |  | 
|  |         
 |  | 
|  |         except Exception as e:
 |  | 
|  |             self.logger.error(f"Error processing PDF: {e}")
 |  | 
|  |             self.logger.exception(traceback.format_exc())
 |  | 
|  |             return "", []
 |  | 
|  |   |  | 
|  |     def extract_markdown(self):
 |  | 
|  |         """Main method to extract markdown from PDF."""
 |  | 
|  |         try:
 |  | 
|  |             doc = fitz.open(self.pdf_path)
 |  | 
|  |             markdown_content = ""
 |  | 
|  |             markdown_pages = []
 |  | 
|  |             tables = self.extract_tables()
 |  | 
|  |             table_index = 0
 |  | 
|  |             list_counter = 0
 |  | 
|  |             in_code_block = False
 |  | 
|  |             code_block_content = ""
 |  | 
|  |             code_block_lang = None
 |  | 
|  |             prev_line = ""
 |  | 
|  |   |  | 
|  |             for page_num, page in enumerate(doc):
 |  | 
|  |                 self.logger.info(f"Processing page {page_num + 1}")
 |  | 
|  |                 page_content = ""
 |  | 
|  |                 blocks = page.get_text("dict")["blocks"]
 |  | 
|  |                 page_height = page.rect.height
 |  | 
|  |                 links = self.extract_links(page)
 |  | 
|  |   |  | 
|  |                 for block in blocks:
 |  | 
|  |                     if block["type"] == 0:  # Text
 |  | 
|  |                         page_content += self.process_text_block(block, page_height, links, list_counter, in_code_block, code_block_content, code_block_lang, prev_line)
 |  | 
|  |                     elif block["type"] == 1:  # Image
 |  | 
|  |                         page_content += self.process_image_block(page, block)
 |  | 
|  |   |  | 
|  |                 # Insert tables at their approximate positions
 |  | 
|  |                 while table_index < len(tables) and tables[table_index]["page"] == page.number:
 |  | 
|  |                     page_content += "\n\n" + self.table_to_markdown(tables[table_index]["content"]) + "\n\n"
 |  | 
|  |                     table_index += 1
 |  | 
|  |   |  | 
|  |                 markdown_pages.append(self.post_process_markdown(page_content))
 |  | 
|  |                 markdown_content += page_content
 |  | 
|  |   |  | 
|  |             markdown_content = self.post_process_markdown(markdown_content)
 |  | 
|  |             return markdown_content, markdown_pages
 |  | 
|  |         except Exception as e:
 |  | 
|  |             self.logger.error(f"Error extracting markdown: {e}")
 |  | 
|  |             self.logger.exception(traceback.format_exc())
 |  | 
|  |             return "", []
 |  | 
|  |   |  | 
|  |     def extract_tables(self):
 |  | 
|  |         """Extract tables from PDF using pdfplumber."""
 |  | 
|  |         tables = []
 |  | 
|  |         try:
 |  | 
|  |             with pdfplumber.open(self.pdf_path) as pdf:
 |  | 
|  |                 for page_number, page in enumerate(pdf.pages):
 |  | 
|  |                     page_tables = page.extract_tables()
 |  | 
|  |                     for table in page_tables:
 |  | 
|  |                         tables.append({"page": page_number, "content": table})
 |  | 
|  |                         self.logger.info(f"Extracted {len(tables)} tables from the PDF.")
 |  | 
|  |         except Exception as e:
 |  | 
|  |             self.logger.error(f"Error extracting tables: {e}")
 |  | 
|  |             self.logger.exception(traceback.format_exc())
 |  | 
|  |         return tables
 |  | 
|  |   |  | 
|  |     def table_to_markdown(self, table):
 |  | 
|  |         """Convert a table to markdown format."""
 |  | 
|  |         if not table:
 |  | 
|  |             return ""
 |  | 
|  |   |  | 
|  |         try:
 |  | 
|  |             table = [['' if cell is None else str(cell).strip() for cell in row] for row in table]
 |  | 
|  |             col_widths = [max(len(cell) for cell in col) for col in zip(*table)]
 |  | 
|  |   |  | 
|  |             markdown = ""
 |  | 
|  |             for i, row in enumerate(table):
 |  | 
|  |                 formatted_row = [cell.ljust(col_widths[j]) for j, cell in enumerate(row)]
 |  | 
|  |                 markdown += "| " + " | ".join(formatted_row) + " |\n"
 |  | 
|  |   |  | 
|  |                 if i == 0:
 |  | 
|  |                     markdown += "|" + "|".join(["-" * (width + 2) for width in col_widths]) + "|\n"
 |  | 
|  |   |  | 
|  |             return markdown
 |  | 
|  |         except Exception as e:
 |  | 
|  |             self.logger.error(f"Error converting table to markdown: {e}")
 |  | 
|  |             self.logger.exception(traceback.format_exc())
 |  | 
|  |             return ""
 |  | 
|  |   |  | 
|  |     def perform_ocr(self, image):
 |  | 
|  |         """Perform OCR on the given image."""
 |  | 
|  |         try:
 |  | 
|  |             opencv_image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
 |  | 
|  |             ocr_result = pytesseract.image_to_data(opencv_image, output_type=pytesseract.Output.DICT)
 |  | 
|  |             
 |  | 
|  |             result = ""
 |  | 
|  |             for word in ocr_result['text']:
 |  | 
|  |                 if word.strip() != "":
 |  | 
|  |                     result += word + " "
 |  | 
|  |   |  | 
|  |                 if len(result) > 30:
 |  | 
|  |                     break
 |  | 
|  |                 
 |  | 
|  |             return result.strip()
 |  | 
|  |         except Exception as e:
 |  | 
|  |             self.logger.error(f"Error performing OCR: {e}")
 |  | 
|  |             self.logger.exception(traceback.format_exc())
 |  | 
|  |             return ""
 |  | 
|  |   |  | 
|  |     def caption_image(self, image):
 |  | 
|  |         """Generate a caption for the given image."""
 |  | 
|  |         try:
 |  | 
|  |             ocr_text = self.perform_ocr(image)
 |  | 
|  |             if ocr_text:
 |  | 
|  |                 return ocr_text
 |  | 
|  |             
 |  | 
|  |             inputs = self.feature_extractor(images=image, return_tensors="pt").to(self.device)
 |  | 
|  |             pixel_values = inputs.pixel_values
 |  | 
|  |   |  | 
|  |             generated_ids = self.model.generate(pixel_values, max_length=30)
 |  | 
|  |             generated_caption = self.tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
 |  | 
|  |             return generated_caption.strip()
 |  | 
|  |         except Exception as e:
 |  | 
|  |             self.logger.error(f"Error captioning image: {e}")
 |  | 
|  |             self.logger.exception(traceback.format_exc())
 |  | 
|  |             return "Image"
 |  | 
|  |   |  | 
|  |     def clean_text(self, text):
 |  | 
|  |         """Clean the given text by removing extra spaces."""
 |  | 
|  |         text = text.strip()
 |  | 
|  |         text = re.sub(r'\s+', ' ', text)
 |  | 
|  |         return text
 |  | 
|  |   |  | 
|  |     def apply_formatting(self, text, flags):
 |  | 
|  |         """Apply markdown formatting to the given text based on flags."""
 |  | 
|  |         text = text.strip()
 |  | 
|  |         if not text:
 |  | 
|  |             return text
 |  | 
|  |         
 |  | 
|  |         is_bold = flags & 2**4
 |  | 
|  |         is_italic = flags & 2**1
 |  | 
|  |         is_monospace = flags & 2**3
 |  | 
|  |         is_superscript = flags & 2**0
 |  | 
|  |         is_subscript = flags & 2**5 
 |  | 
|  |   |  | 
|  |         if is_monospace:
 |  | 
|  |             text = f"`{text}`"
 |  | 
|  |         elif is_superscript and not bool(re.search(r'\s+', text)):
 |  | 
|  |             text = f"^{text}^"
 |  | 
|  |         elif is_subscript and not bool(re.search(r'\s+', text)):
 |  | 
|  |             text = f"~{text}~"
 |  | 
|  |   |  | 
|  |         if is_bold and is_italic:
 |  | 
|  |             text = f"***{text}***"
 |  | 
|  |         elif is_bold:
 |  | 
|  |             text = f"**{text}**"
 |  | 
|  |         elif is_italic:
 |  | 
|  |             text = f"*{text}*"
 |  | 
|  |   |  | 
|  |         return f" {text} "
 |  | 
|  |   |  | 
|  |     def is_bullet_point(self, text):
 |  | 
|  |         """Check if the given text is a bullet point."""
 |  | 
|  |         return text.strip().startswith(tuple(self.BULLET_POINTS))
 |  | 
|  |   |  | 
|  |     def convert_bullet_to_markdown(self, text):
 |  | 
|  |         """Convert a bullet point to markdown format."""
 |  | 
|  |         text = re.sub(r'^\s*', '', text)
 |  | 
|  |         return re.sub(f'^[{re.escape(self.BULLET_POINTS)}]\s*', '- ', text)
 |  | 
|  |   |  | 
|  |     def is_numbered_list_item(self, text):
 |  | 
|  |         """Check if the given text is a numbered list item."""
 |  | 
|  |         return bool(re.match(r'^\d+\s{0,3}[.)]', text.strip()))
 |  | 
|  |   |  | 
|  |     def convert_numbered_list_to_markdown(self, text, list_counter):
 |  | 
|  |         """Convert a numbered list item to markdown format."""
 |  | 
|  |         text = re.sub(r'^\s*', '', text)
 |  | 
|  |         return re.sub(r'^\d+\s{0,3}[.)]', f"{list_counter}. ", text)
 |  | 
|  |   |  | 
|  |     def is_horizontal_line(self, text):
 |  | 
|  |         """Check if the given text represents a horizontal line."""
 |  | 
|  |         return bool(re.match(r'^[_-]+$', text.strip()))
 |  | 
|  |   |  | 
|  |     def extract_links(self, page):
 |  | 
|  |         """Extract links from the given page."""
 |  | 
|  |         links = []
 |  | 
|  |         try:
 |  | 
|  |             for link in page.get_links():
 |  | 
|  |                 if link["kind"] == 2:  # URI link
 |  | 
|  |                     links.append({
 |  | 
|  |                         "rect": link["from"],
 |  | 
|  |                         "uri": link["uri"]
 |  | 
|  |                     })
 |  | 
|  |                     self.logger.info(f"Extracted {len(links)} links from the page.")
 |  | 
|  |         except Exception as e:
 |  | 
|  |             self.logger.error(f"Error extracting links: {e}")
 |  | 
|  |             self.logger.exception(traceback.format_exc())
 |  | 
|  |         return links
 |  | 
|  |   |  | 
|  |     def detect_code_block(self, prev_line, current_line):
 |  | 
|  |         """Detect if the current line starts a code block."""
 |  | 
|  |         patterns = {
 |  | 
|  |             'python': [
 |  | 
|  |                 (r'^(?:from|import)\s+\w+', r'^(?:from|import|def|class|if|for|while|try|except|with)\s'),
 |  | 
|  |                 (r'^(?:def|class)\s+\w+', r'^\s{4}'),
 |  | 
|  |                 (r'^\s{4}', r'^\s{4,}')
 |  | 
|  |             ],
 |  | 
|  |             'javascript': [
 |  | 
|  |                 (r'^(?:function|const|let|var)\s+\w+', r'^(?:function|const|let|var|if|for|while|try|catch|class)\s'),
 |  | 
|  |                 (r'^(?:if|for|while)\s*\(', r'^\s{2,}'),
 |  | 
|  |                 (r'^\s{2,}', r'^\s{2,}')
 |  | 
|  |             ],
 |  | 
|  |             'html': [
 |  | 
|  |                 (r'^<(!DOCTYPE|html|head|body|div|p|a|script|style)', r'^<(!DOCTYPE|html|head|body|div|p|a|script|style)'),
 |  | 
|  |                 (r'^<\w+.*>$', r'^\s{2,}<'),
 |  | 
|  |                 (r'^\s{2,}<', r'^\s{2,}<')
 |  | 
|  |             ],
 |  | 
|  |             'shell': [
 |  | 
|  |                 (r'^(?:\$|\#)\s', r'^(?:\$|\#)\s'),
 |  | 
|  |                 (r'^[a-z_]+\s*=', r'^[a-z_]+\s*=')
 |  | 
|  |             ],
 |  | 
|  |             'bash': [
 |  | 
|  |                 (r'^(?:#!/bin/bash|alias|export|source)\s', r'^(?:#!/bin/bash|alias|export|source|echo|read|if|for|while|case|function)\s'),
 |  | 
|  |                 (r'^(?:if|for|while|case|function)\s', r'^\s{2,}'),
 |  | 
|  |                 (r'^\s{2,}', r'^\s{2,}')
 |  | 
|  |             ],
 |  | 
|  |             'cpp': [
 |  | 
|  |                 (r'^#include\s*<', r'^(?:#include|using|namespace|class|struct|enum|template|typedef)\s'),
 |  | 
|  |                 (r'^(?:class|struct|enum)\s+\w+', r'^\s{2,}'),
 |  | 
|  |                 (r'^\s{2,}', r'^\s{2,}')
 |  | 
|  |             ],
 |  | 
|  |             'java': [
 |  | 
|  |                 (r'^(?:import|package)\s+\w+', r'^(?:import|package|public|private|protected|class|interface|enum)\s'),
 |  | 
|  |                 (r'^(?:public|private|protected)\s+class\s+\w+', r'^\s{4,}'),
 |  | 
|  |                 (r'^\s{4,}', r'^\s{4,}')
 |  | 
|  |             ],
 |  | 
|  |             'json': [
 |  | 
|  |                 (r'^\s*{', r'^\s*["{[]'),
 |  | 
|  |                 (r'^\s*"', r'^\s*["}],?$'),
 |  | 
|  |                 (r'^\s*\[', r'^\s*[}\]],?$')
 |  | 
|  |             ]
 |  | 
|  |         }
 |  | 
|  |         
 |  | 
|  |         for lang, pattern_pairs in patterns.items():
 |  | 
|  |             for prev_pattern, curr_pattern in pattern_pairs:
 |  | 
|  |                 if (re.match(prev_pattern, prev_line.strip()) and 
 |  | 
|  |                     re.match(curr_pattern, current_line.strip())):
 |  | 
|  |                     return lang
 |  | 
|  |                 
 |  | 
|  |         return None
 |  | 
|  |   |  | 
|  |     def process_text_block(self, block, page_height, links, list_counter, in_code_block, code_block_content, code_block_lang, prev_line):
 |  | 
|  |         """Process a text block and convert it to markdown."""
 |  | 
|  |         try:
 |  | 
|  |             block_rect = block["bbox"]
 |  | 
|  |             if block_rect[1] < 50 or block_rect[3] > page_height - 50:
 |  | 
|  |                 return ""  # Skip headers and footers
 |  | 
|  |   |  | 
|  |             block_text = ""
 |  | 
|  |             last_y1 = None
 |  | 
|  |             last_font_size = None
 |  | 
|  |   |  | 
|  |             for line in block["lines"]:
 |  | 
|  |                 line_text = ""
 |  | 
|  |                 curr_font_size = [span["size"] for span in line["spans"]]
 |  | 
|  |   |  | 
|  |                 for span in line["spans"]:
 |  | 
|  |                     text = span["text"]
 |  | 
|  |                     font_size = span["size"]
 |  | 
|  |                     flags = span["flags"]
 |  | 
|  |                     span_rect = span["bbox"]
 |  | 
|  |   |  | 
|  |                     if self.is_horizontal_line(text):
 |  | 
|  |                         line_text += "\n---\n"
 |  | 
|  |                         continue
 |  | 
|  |                     
 |  | 
|  |                     text = self.clean_text(text)
 |  | 
|  |   |  | 
|  |                     if text.strip():
 |  | 
|  |                         header_level = self.get_header_level(font_size)                   
 |  | 
|  |                         if header_level > 0:
 |  | 
|  |                             text = f"\n{'#' * header_level} {text}\n\n"
 |  | 
|  |   |  | 
|  |                         else:
 |  | 
|  |                             is_list_item = self.is_bullet_point(text) or self.is_numbered_list_item(text)
 |  | 
|  |   |  | 
|  |                             if is_list_item:
 |  | 
|  |                                 marker, content = re.split(r'(?<=^[•◦▪▫●○\d.)])\s*', text, 1)
 |  | 
|  |                                 formatted_content = self.apply_formatting(content, flags)
 |  | 
|  |                                 text = f"{marker} {formatted_content}"
 |  | 
|  |                             else:
 |  | 
|  |                                 text = self.apply_formatting(text, flags)
 |  | 
|  |                                 
 |  | 
|  |                     for link in links:
 |  | 
|  |                         if fitz.Rect(span_rect).intersects(link["rect"]):
 |  | 
|  |                             text = f"[{text.strip()}]({link['uri']})"
 |  | 
|  |                             break
 |  | 
|  |   |  | 
|  |                     line_text += text
 |  | 
|  |   |  | 
|  |                 if last_y1 is not None:
 |  | 
|  |                     avg_last_font_size = sum(last_font_size) / len(last_font_size) if last_font_size else 0
 |  | 
|  |                     avg_current_font_size = sum(curr_font_size) / len(curr_font_size)
 |  | 
|  |                     font_size_changed = abs(avg_current_font_size - avg_last_font_size) > 1
 |  | 
|  |   |  | 
|  |                     if abs(line["bbox"][3] - last_y1) > 2 or font_size_changed:
 |  | 
|  |                         block_text += "\n"
 |  | 
|  |                         
 |  | 
|  |                 block_text += self.clean_text(line_text) + " "
 |  | 
|  |                 last_font_size = curr_font_size
 |  | 
|  |                 last_y1 = line["bbox"][3]
 |  | 
|  |   |  | 
|  |             markdown_content = ""
 |  | 
|  |             lines = block_text.split('\n')
 |  | 
|  |             for i, line in enumerate(lines):
 |  | 
|  |                 clean_line = self.clean_text(line)
 |  | 
|  |   |  | 
|  |                 if not in_code_block:
 |  | 
|  |                     code_lang = self.detect_code_block(prev_line, clean_line)
 |  | 
|  |                     if code_lang:
 |  | 
|  |                         in_code_block = True
 |  | 
|  |                         code_block_lang = code_lang
 |  | 
|  |                         code_block_content = prev_line + "\n" + clean_line + "\n"
 |  | 
|  |                         prev_line = clean_line
 |  | 
|  |                         continue
 |  | 
|  |   |  | 
|  |                 if in_code_block:
 |  | 
|  |                     code_block_content += clean_line + "\n"
 |  | 
|  |                     if i == len(lines) - 1 or self.detect_code_block(clean_line, lines[i+1]) != code_block_lang:
 |  | 
|  |                         markdown_content += f"```{code_block_lang}\n{code_block_content}```\n\n"
 |  | 
|  |                         in_code_block = False
 |  | 
|  |                         code_block_content = ""
 |  | 
|  |                         code_block_lang = None
 |  | 
|  |                 else:
 |  | 
|  |                     if self.is_bullet_point(clean_line):
 |  | 
|  |                         markdown_content += "\n" + self.convert_bullet_to_markdown(clean_line)
 |  | 
|  |                         list_counter = 0
 |  | 
|  |                     elif self.is_numbered_list_item(clean_line):
 |  | 
|  |                         list_counter += 1
 |  | 
|  |                         markdown_content += "\n" + self.convert_numbered_list_to_markdown(clean_line, list_counter)   
 |  | 
|  |                     else:
 |  | 
|  |                         markdown_content += f"{clean_line}\n"
 |  | 
|  |                         list_counter = 0
 |  | 
|  |   |  | 
|  |                 prev_line = clean_line
 |  | 
|  |   |  | 
|  |             return markdown_content + "\n"
 |  | 
|  |         except Exception as e:
 |  | 
|  |             self.logger.error(f"Error processing text block: {e}")
 |  | 
|  |             self.logger.exception(traceback.format_exc())
 |  | 
|  |             return ""
 |  | 
|  |   |  | 
|  |     def process_image_block(self, page, block):
 |  | 
|  |         """Process an image block and convert it to markdown."""
 |  | 
|  |         try:
 |  | 
|  |             image_rect = block["bbox"]
 |  | 
|  |             pix = page.get_pixmap(clip=image_rect)
 |  | 
|  |             image = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
 |  | 
|  |             
 |  | 
|  |             image_filename = f"{self.pdf_filename}_image_{int(page.number)+1}_{block['number']}.png"
 |  | 
|  |             image.save(f"outputs/{image_filename}")
 |  | 
|  |   |  | 
|  |             caption = self.caption_image(image)
 |  | 
|  |   |  | 
|  |             return f"\n\n"
 |  | 
|  |         except Exception as e:
 |  | 
|  |             self.logger.error(f"Error processing image block: {e}")
 |  | 
|  |             self.logger.exception(traceback.format_exc())
 |  | 
|  |             return ""
 |  | 
|  |   |  | 
|  |     def get_header_level(self, font_size):
 |  | 
|  |         """Determine header level based on font size."""
 |  | 
|  |         if font_size > 24:
 |  | 
|  |             return 1
 |  | 
|  |         elif font_size > 20:
 |  | 
|  |             return 2
 |  | 
|  |         elif font_size > 18:
 |  | 
|  |             return 3
 |  | 
|  |         elif font_size > 16:
 |  | 
|  |             return 4
 |  | 
|  |         elif font_size > 14:
 |  | 
|  |             return 5
 |  | 
|  |         elif font_size > 12:
 |  | 
|  |             return 6
 |  | 
|  |         else:
 |  | 
|  |             return 0
 |  | 
|  |   |  | 
|  |     def post_process_markdown(self, markdown_content):
 |  | 
|  |         """Post-process the markdown content."""
 |  | 
|  |         try:
 |  | 
|  |             markdown_content = re.sub(r'\n{3,}', '\n\n', markdown_content)  # Remove excessive newlines
 |  | 
|  |             markdown_content = re.sub(r'(\d+)\s*\n', '', markdown_content)  # Remove page numbers
 |  | 
|  |             markdown_content = re.sub(r' +', ' ', markdown_content)  # Remove multiple spaces
 |  | 
|  |             markdown_content = re.sub(r'\s*(---\n)+', '\n\n---\n', markdown_content)  # Remove duplicate horizontal lines
 |  | 
|  |   |  | 
|  |             def remove_middle_headers(match):
 |  | 
|  |                 line = match.group(0)
 |  | 
|  |                 # Keep the initial header and remove all subsequent '#' characters
 |  | 
|  |                 return re.sub(r'(^#{1,6}\s).*?(?=\n)', lambda m: m.group(1) + re.sub(r'#', '', m.group(0)[len(m.group(1)):]), line)
 |  | 
|  |             
 |  | 
|  |             markdown_content = re.sub(r'^#{1,6}\s.*\n', remove_middle_headers, markdown_content, flags=re.MULTILINE) # Remove headers in the middle of lines         
 |  | 
|  |             return markdown_content
 |  | 
|  |         except Exception as e:
 |  | 
|  |             self.logger.error(f"Error post-processing markdown: {e}")
 |  | 
|  |             self.logger.exception(traceback.format_exc())
 |  | 
|  |             return markdown_content
 |  | 
|  |   |  | 
|  |     def save_markdown(self, markdown_content):
 |  | 
|  |         """Save the markdown content to a file."""
 |  | 
|  |         try:
 |  | 
|  |             os.makedirs("outputs", exist_ok=True)
 |  | 
|  |             with open(f"outputs/{self.pdf_filename}.md", "w", encoding="utf-8") as f:
 |  | 
|  |                 f.write(markdown_content)
 |  | 
|  |                 self.logger.info("Markdown content saved successfully.")
 |  | 
|  |         except Exception as e:
 |  | 
|  |             self.logger.error(f"Error saving markdown content: {e}")
 |  | 
|  |             self.logger.exception(traceback.format_exc())
 |  | 
|  |   |  | 
|  | def main():
 |  | 
|  |     parser = argparse.ArgumentParser(description="Extract markdown-formatted content from a PDF file.")
 |  | 
|  |     parser.add_argument("pdf_path", help="Path to the input PDF file")
 |  | 
|  |     args = parser.parse_args()
 |  | 
|  |   |  | 
|  |     extractor = MarkdownPDFExtractor(args.pdf_path)
 |  | 
|  |     markdown_pages = extractor.extract()
 |  | 
|  |     return markdown_pages
 |  | 
|  |   |  | 
|  | if __name__ == "__main__":
 |  | 
|  |     main()
 |  |