754 lines
43 KiB
Python
754 lines
43 KiB
Python
#!/usr/bin/env python3
|
|
# =============================================================================
|
|
# ENHANCED PDF HIGHLIGHT EXTRACTOR
|
|
# Author: Perplexity AI Companion (Updated by User Feedback)
|
|
# Date: June 3, 2025
|
|
# License: MIT
|
|
#
|
|
# Extracts highlights from PDF files, with options for interactive review,
|
|
# detailed output, text cleaning, JSON export, and page image viewing.
|
|
# =============================================================================
|
|
|
|
import time
|
|
import os
|
|
import fitz # PyMuPDF
|
|
import json
|
|
from colorama import init, Fore, Back, Style
|
|
from pathlib import Path
|
|
import re
|
|
import string
|
|
import sys
|
|
import traceback
|
|
import argparse
|
|
import difflib # For text difference calculation
|
|
import tempfile # For temporary image files
|
|
import webbrowser # For opening images/PDFs
|
|
import uuid # For unique filenames
|
|
|
|
# Attempt to import readline for better input() experience on some systems
|
|
try:
|
|
import readline
|
|
READLINE_AVAILABLE = True
|
|
except ImportError:
|
|
READLINE_AVAILABLE = False # readline not available
|
|
|
|
# =============================================================================
|
|
# GLOBAL CONFIGURATION FLAGS (Defaults, can be overridden by CLI args)
|
|
# =============================================================================
|
|
DEFAULT_PDF_PATH = "/mnt/c/Users/admin/Downloads/test2.pdf" # Example, adjust if needed
|
|
DEFAULT_PAGES_TO_PROCESS = "3" # Example: "1,3-5,all"
|
|
|
|
# Default Behavior flags (can be influenced by -d or -s CLI flags)
|
|
# These are used to initialize effective_run_args
|
|
# Keep these distinct from the effective_run_args object itself
|
|
INITIAL_SHOW_TIMING = True
|
|
INITIAL_SHOW_PROGRESS = True
|
|
INITIAL_SHOW_RAW_SEGMENTS = True
|
|
INITIAL_SHOW_EXTRACTION_DETAILS = True
|
|
INITIAL_SHOW_RECT_DETAILS = True
|
|
INITIAL_SHOW_DIFF_PERCENTAGE = True
|
|
INITIAL_CLEAN_EDGES = True
|
|
|
|
# Text extraction parameters (generally fixed)
|
|
TEXT_EXTRACTION_HORIZONTAL_PADDING = 6.0
|
|
TEXT_EXTRACTION_VERTICAL_PADDING = 1.0
|
|
|
|
# Edge cleaning configuration (generally fixed)
|
|
VALID_TWO_LETTER_WORDS = {
|
|
'am', 'an', 'as', 'at', 'be', 'by', 'do', 'go', 'he', 'if', 'in', 'is', 'it', 'me', 'my',
|
|
'no', 'of', 'on', 'or', 'ox', 'so', 'to', 'up', 'us', 'we'}
|
|
VALID_SINGLE_LETTERS = {'i', 'a'}
|
|
|
|
# Image handling configuration
|
|
IMAGE_FOLDER_PATH = 'pdf_page_images' # Relative to CWD by default
|
|
CLEAR_IMAGE_FOLDER_ON_START = True
|
|
CLEAR_IMAGE_FOLDER_ON_END = False
|
|
|
|
# Initialize colorama
|
|
init(autoreset=True)
|
|
|
|
# --- Helper Functions ---
|
|
def get_text_diff_ratio(text1, text2):
|
|
if not text1 and not text2: return 1.0
|
|
if not text1 or not text2: return 0.0
|
|
return difflib.SequenceMatcher(None, str(text1), str(text2)).ratio()
|
|
|
|
def clean_segment_edges_func(text_to_clean, clean_edges_setting):
|
|
if not clean_edges_setting or not text_to_clean: return text_to_clean
|
|
text_to_clean = re.sub(r'\s+', ' ', text_to_clean.strip())
|
|
words = text_to_clean.split()
|
|
if not words: return text_to_clean
|
|
|
|
current_idx = 0
|
|
while current_idx < len(words):
|
|
token = words[current_idx]
|
|
core_token = token.rstrip(string.punctuation)
|
|
trailing_punctuation = token[len(core_token):]
|
|
if not core_token: words.pop(current_idx); continue
|
|
core_should_be_removed = (len(core_token) == 1 and core_token.isalpha() and core_token.lower() not in VALID_SINGLE_LETTERS) or \
|
|
(len(core_token) == 2 and core_token.isalpha() and core_token.lower() not in VALID_TWO_LETTER_WORDS)
|
|
if core_should_be_removed:
|
|
if trailing_punctuation: words[current_idx] = trailing_punctuation
|
|
else: words.pop(current_idx)
|
|
continue
|
|
break
|
|
while words:
|
|
token = words[-1]
|
|
core_token = token.lstrip(string.punctuation)
|
|
leading_punctuation = token[:-len(core_token)] if core_token else ""
|
|
if not core_token: words.pop(); continue
|
|
core_should_be_removed = (len(core_token) == 1 and core_token.isalpha() and core_token.lower() not in VALID_SINGLE_LETTERS) or \
|
|
(len(core_token) == 2 and core_token.isalpha() and core_token.lower() not in VALID_TWO_LETTER_WORDS)
|
|
if core_should_be_removed:
|
|
if leading_punctuation: words[-1] = leading_punctuation
|
|
else: words.pop()
|
|
continue
|
|
break
|
|
return ' '.join(words)
|
|
|
|
def input_with_prefill(prompt, text):
|
|
if READLINE_AVAILABLE:
|
|
def hook():
|
|
readline.insert_text(text)
|
|
readline.redisplay()
|
|
readline.set_pre_input_hook(hook)
|
|
result = input(prompt)
|
|
readline.set_pre_input_hook()
|
|
return result
|
|
else:
|
|
print(Fore.MAGENTA + "Current text (edit below):\n" + Style.RESET_ALL + f"{text}")
|
|
return input(prompt)
|
|
|
|
def _clear_png_files_in_folder(folder_path_str, run_args_for_print_control):
|
|
# This function CLEARS files if folder exists. It DOES NOT CREATE the folder.
|
|
if not folder_path_str: return
|
|
|
|
folder = Path(folder_path_str) # Path relative to CWD if not absolute
|
|
abs_folder_path = folder.resolve()
|
|
|
|
if run_args_for_print_control.debug:
|
|
print(Fore.CYAN + f" [Debug] _clear_png_files_in_folder: Checking {abs_folder_path} (Specified as: '{folder_path_str}')")
|
|
|
|
if abs_folder_path.is_dir():
|
|
if run_args_for_print_control.show_progress:
|
|
print(Fore.BLUE + f"Clearing *.png files from {abs_folder_path}...")
|
|
cleared_count = 0
|
|
try:
|
|
for file_path in abs_folder_path.glob("*.png"):
|
|
if file_path.is_file():
|
|
file_path.unlink()
|
|
cleared_count +=1
|
|
except Exception as e:
|
|
if run_args_for_print_control.show_progress: # Also show error if progress is on
|
|
print(Fore.RED + f"Error during file deletion in {abs_folder_path}: {e}")
|
|
|
|
if run_args_for_print_control.show_progress:
|
|
if cleared_count > 0:
|
|
print(Fore.BLUE + f"Cleared {cleared_count} *.png files from {abs_folder_path}.")
|
|
else:
|
|
print(Fore.BLUE + f"No *.png files found to clear in {abs_folder_path}.")
|
|
else:
|
|
if run_args_for_print_control.show_progress:
|
|
print(Fore.YELLOW + f"Image folder {abs_folder_path} not found, skipping clear.")
|
|
elif run_args_for_print_control.debug: # Still log if not found in debug, even if not show_progress
|
|
print(Fore.CYAN + f" [Debug] _clear_png_files_in_folder: Folder {abs_folder_path} does not exist. Nothing to clear.")
|
|
|
|
|
|
class EnhancedPDFHighlightExtractor:
|
|
def __init__(self, pdf_path, effective_run_args, main_doc_for_image_view=None):
|
|
self.pdf_path = Path(pdf_path)
|
|
self.run_args = effective_run_args
|
|
self.pdf_filename_stem = self.pdf_path.stem
|
|
self.highlights_data = []
|
|
self.main_doc_for_image_view = main_doc_for_image_view
|
|
|
|
def _get_highlight_color_from_rgb_tuple(self, rgb_tuple_floats_or_ints):
|
|
if not rgb_tuple_floats_or_ints or len(rgb_tuple_floats_or_ints) < 3 : return 'unknown_color'
|
|
r, g, b = [int(x * 255) if isinstance(x, float) and 0.0 <= x <= 1.0 else int(x) for x in rgb_tuple_floats_or_ints[:3]]
|
|
if r == 142 and g == 221 and b == 249: return 'blue'
|
|
if r > 200 and g > 200 and b < 150: return 'yellow'
|
|
if r < 150 and g > 180 and b < 150: return 'green'
|
|
if r < 150 and g < 180 and b > 180: return 'blue'
|
|
if r > 180 and g < 180 and b > 180: return 'pink'
|
|
return 'other_color'
|
|
|
|
def _get_highlight_color_from_annot_colors_dict(self, colors_dict):
|
|
if not colors_dict: return 'unknown_color', None
|
|
rgb_tuple = colors_dict.get('stroke') or colors_dict.get('fill')
|
|
if not rgb_tuple: return 'unknown_color', None
|
|
return self._get_highlight_color_from_rgb_tuple(rgb_tuple), rgb_tuple[:3]
|
|
|
|
def _extract_text_from_multi_segment_highlight(self, page, annot, page_num, hl_id):
|
|
overall_highlight_color_name, _ = self._get_highlight_color_from_annot_colors_dict(annot.colors)
|
|
color_code_for_segment_print = self._get_color_display_codes(overall_highlight_color_name)
|
|
quads_vertices = annot.vertices
|
|
if not quads_vertices:
|
|
if self.run_args.show_extraction_details: print(Fore.YELLOW + f" No quads for HL {hl_id} on page {page_num}")
|
|
return None, 0, []
|
|
|
|
processed_quads_as_points_list = []
|
|
if len(quads_vertices) % 4 == 0:
|
|
for i in range(0, len(quads_vertices), 4):
|
|
try:
|
|
quad_points = [fitz.Point(p) for p in quads_vertices[i:i+4]]
|
|
processed_quads_as_points_list.append(quad_points)
|
|
except Exception as e:
|
|
if self.run_args.show_extraction_details: print(Fore.YELLOW + f" Skipping malformed quad points: {e}")
|
|
continue
|
|
|
|
try:
|
|
sorted_quad_points_list = sorted(processed_quads_as_points_list, key=lambda qp_list: (fitz.Quad(qp_list).rect.y0, fitz.Quad(qp_list).rect.x0))
|
|
except Exception as e:
|
|
if self.run_args.show_extraction_details: print(Fore.RED + f" Error sorting quads for HL {hl_id}: {e}. Using original order.")
|
|
sorted_quad_points_list = processed_quads_as_points_list
|
|
|
|
if self.run_args.show_extraction_details:
|
|
print(color_code_for_segment_print + Fore.CYAN + f" Processing {len(sorted_quad_points_list)} segments for HL {hl_id} (Color: {overall_highlight_color_name.upper()}) on page {page_num}" + Style.RESET_ALL)
|
|
|
|
segment_texts_final = []
|
|
raw_segment_texts_for_diff = []
|
|
for seg_idx, quad_points in enumerate(sorted_quad_points_list):
|
|
try:
|
|
bounds = fitz.Quad(quad_points).rect
|
|
padded_rect = fitz.Rect(bounds.x0 - TEXT_EXTRACTION_HORIZONTAL_PADDING, bounds.y0 - TEXT_EXTRACTION_VERTICAL_PADDING,
|
|
bounds.x1 + TEXT_EXTRACTION_HORIZONTAL_PADDING, bounds.y1 + TEXT_EXTRACTION_VERTICAL_PADDING)
|
|
padded_rect.intersect(page.rect)
|
|
if padded_rect.is_empty:
|
|
if self.run_args.show_extraction_details: print(Fore.YELLOW + f" Segment {seg_idx+1} empty padded_rect for HL {hl_id}")
|
|
continue
|
|
raw_text_from_pdf_segment = page.get_text("text", clip=padded_rect, sort=True).strip()
|
|
raw_segment_texts_for_diff.append(raw_text_from_pdf_segment)
|
|
cleaned_text_segment = re.sub(r'\s+', ' ', raw_text_from_pdf_segment).strip()
|
|
cleaned_text_segment = re.sub(r'[\x00-\x08\x0B\x0C\x0E-\x1F\x7F-\x9F]', '', cleaned_text_segment)
|
|
final_text_segment = clean_segment_edges_func(cleaned_text_segment, self.run_args.clean_edges)
|
|
|
|
if final_text_segment:
|
|
segment_texts_final.append(final_text_segment)
|
|
if self.run_args.show_raw_segments and not self.run_args.interactive:
|
|
print(color_code_for_segment_print + Fore.LIGHTBLUE_EX + f" Segment {seg_idx+1} (P{page_num}, HL{hl_id}, Color: {overall_highlight_color_name.upper()}):" + Style.RESET_ALL)
|
|
if self.run_args.show_diff_percentage:
|
|
similarity = get_text_diff_ratio(raw_text_from_pdf_segment, final_text_segment)
|
|
diff_percent = (1 - similarity) * 100
|
|
print(Fore.LIGHTMAGENTA_EX + f" Raw PDF : \"{raw_text_from_pdf_segment}\"")
|
|
print(Fore.LIGHTBLUE_EX + f" Final Seg: \"{final_text_segment}\"")
|
|
print(Fore.YELLOW + f" Diff: {diff_percent:.2f}%")
|
|
else: print(Fore.LIGHTBLUE_EX + f" Final Seg: \"{final_text_segment}\"")
|
|
except Exception as e:
|
|
if self.run_args.show_extraction_details: print(Fore.RED + f" Error processing segment {seg_idx+1} for HL {hl_id}: {e}")
|
|
raw_segment_texts_for_diff.append("")
|
|
continue
|
|
|
|
if not segment_texts_final: return None, len(sorted_quad_points_list), raw_segment_texts_for_diff
|
|
combined_text = segment_texts_final[0]
|
|
for i in range(1, len(segment_texts_final)):
|
|
prev_text = combined_text; current_text = segment_texts_final[i]
|
|
if prev_text.endswith('-') or prev_text.endswith('¬'): combined_text = prev_text.rstrip('-¬') + current_text
|
|
else: combined_text += ' ' + current_text
|
|
|
|
if self.run_args.clean_edges: combined_text = clean_segment_edges_func(combined_text, self.run_args.clean_edges)
|
|
combined_text = re.sub(r'\s+', ' ', combined_text).strip()
|
|
return combined_text if combined_text else None, len(sorted_quad_points_list), raw_segment_texts_for_diff
|
|
|
|
def extract_highlights(self, doc):
|
|
all_extracted_highlights = []
|
|
try:
|
|
if self.run_args.show_progress and not self.run_args.interactive:
|
|
print(Fore.BLUE + f"\n🎨 Processing highlights for PDF: {self.pdf_path.name}")
|
|
|
|
pages_str_to_parse = self.run_args.pages if self.run_args.pages else DEFAULT_PAGES_TO_PROCESS
|
|
pages_to_process = self._parse_specific_pages(pages_str_to_parse, doc.page_count)
|
|
if not pages_to_process:
|
|
if self.run_args.show_progress: print(Fore.YELLOW + "No valid pages selected.")
|
|
return []
|
|
|
|
highlight_id_counter_on_page = {}
|
|
for page_num in pages_to_process:
|
|
page = doc.load_page(page_num - 1)
|
|
highlight_id_counter_on_page.setdefault(page_num, 0)
|
|
if self.run_args.show_progress and not self.run_args.interactive:
|
|
print(Fore.CYAN + f" 📄 Processing Page {page_num}...")
|
|
try: page_annotations = list(page.annots())
|
|
except Exception as e:
|
|
if self.run_args.show_progress: print(Fore.RED + f" ⚠️ Error loading annots: {e}")
|
|
continue
|
|
|
|
highlight_annotations = [a for a in page_annotations if hasattr(a, 'type') and a.type[0] == fitz.PDF_ANNOT_HIGHLIGHT and hasattr(a, 'rect') and a.rect]
|
|
if not highlight_annotations:
|
|
if self.run_args.show_progress and not self.run_args.interactive: print(Fore.WHITE + f" No highlights on page {page_num}.")
|
|
continue
|
|
|
|
if self.run_args.show_rect_details:
|
|
print(Fore.YELLOW + f"--- Annotations before sorting (Page {page_num}) ---")
|
|
temp_debug_list = []
|
|
for annot_debug in highlight_annotations:
|
|
debug_text_snippet = page.get_text("text", clip=annot_debug.rect).strip().replace("\n", " ")
|
|
color_name_debug, rgb_values_debug = self._get_highlight_color_from_annot_colors_dict(annot_debug.colors)
|
|
rgb_display = f"RGB: {tuple(int(c*255) if isinstance(c,float) else int(c) for c in rgb_values_debug[:3])}" if rgb_values_debug else "RGB: N/A"
|
|
temp_debug_list.append({
|
|
"rect": annot_debug.rect, "text_snippet": debug_text_snippet, "color_name": color_name_debug,
|
|
"rgb_display": rgb_display, "vertices_count": len(annot_debug.vertices) if annot_debug.vertices else 0 })
|
|
temp_debug_list.sort(key=lambda item: (item["rect"].y0, item["rect"].x0))
|
|
for item_idx, item_val in enumerate(temp_debug_list):
|
|
print(f" {item_idx+1}. Rect: {item_val['rect']}, Vertices: {item_val['vertices_count']}, Color: {item_val['color_name'].upper()} ({item_val['rgb_display']}), Text: \"{item_val['text_snippet']}\"")
|
|
print(Fore.YELLOW + "----------------------------------------------------")
|
|
|
|
highlight_annotations.sort(key=lambda a: (a.rect.y0, a.rect.x0))
|
|
for annot in highlight_annotations:
|
|
try:
|
|
highlight_id_counter_on_page[page_num] += 1; current_hl_id_on_page = highlight_id_counter_on_page[page_num]
|
|
color_name, raw_rgb_floats = self._get_highlight_color_from_annot_colors_dict(annot.colors)
|
|
extracted_text, num_segments, _ = self._extract_text_from_multi_segment_highlight(page, annot, page_num, current_hl_id_on_page)
|
|
if extracted_text and extracted_text.strip():
|
|
if self.run_args.show_extraction_details and not self.run_args.interactive:
|
|
print(Fore.GREEN + f" ✅ Final (P{page_num}, HL{current_hl_id_on_page}): \"{extracted_text[:100]}\"")
|
|
all_extracted_highlights.append({
|
|
'page': page_num, 'highlight_id_on_page': current_hl_id_on_page, 'text': extracted_text,
|
|
'color': color_name, 'raw_rgb_values': raw_rgb_floats, 'type': 'highlight',
|
|
'y_position': annot.rect.y0, 'x_position': annot.rect.x0,
|
|
'rect_details': (annot.rect.x0, annot.rect.y0, annot.rect.x1, annot.rect.y1),
|
|
'num_segments': num_segments })
|
|
elif self.run_args.show_progress and not self.run_args.interactive:
|
|
print(Fore.YELLOW + f" ⚠️ No text for HL {current_hl_id_on_page} on page {page_num}")
|
|
except Exception as e:
|
|
if self.run_args.show_progress and not self.run_args.interactive:
|
|
print(Fore.RED + f" 🔴 Error processing annot on page {page_num}: {e}")
|
|
if self.run_args.debug: traceback.print_exc()
|
|
continue
|
|
|
|
if self.run_args.interactive:
|
|
print(Fore.MAGENTA + "\nEntering interactive review session...")
|
|
self.highlights_data = self._interactive_review_session(all_extracted_highlights)
|
|
else: self.highlights_data = all_extracted_highlights
|
|
|
|
if self.run_args.show_progress and not self.run_args.interactive and not self.run_args.silent:
|
|
print(Fore.MAGENTA + f" 📊 Total highlights extracted: {len(self.highlights_data)}")
|
|
return self.highlights_data
|
|
except Exception as e:
|
|
print(Fore.RED + f"❌ Major error during highlight extraction: {e}")
|
|
if self.run_args.debug: traceback.print_exc()
|
|
return []
|
|
|
|
def _view_page_image_interactively(self, page_num_to_view):
|
|
if not self.main_doc_for_image_view:
|
|
print(Fore.RED + "Error: PDF document not available for image rendering. This should not happen.")
|
|
return
|
|
|
|
tmp_image_path_obj = None
|
|
image_created_in_managed_folder = False
|
|
image_successfully_saved = False
|
|
|
|
if self.run_args.show_progress:
|
|
print(Fore.BLUE + f"Preparing to view image for page {page_num_to_view}...")
|
|
|
|
try:
|
|
page_index = page_num_to_view - 1
|
|
page = self.main_doc_for_image_view.load_page(page_index)
|
|
if self.run_args.debug:
|
|
print(Fore.CYAN + f" [Debug] Loaded page object for index {page_index}: {page}")
|
|
|
|
pix = page.get_pixmap(dpi=150)
|
|
if self.run_args.debug:
|
|
print(Fore.CYAN + f" [Debug] Created pixmap: {pix}. Alpha: {pix.alpha}, Colorspace: {pix.colorspace.name}")
|
|
|
|
if IMAGE_FOLDER_PATH:
|
|
img_dir_path_obj = Path(IMAGE_FOLDER_PATH) # Path relative to CWD if not absolute
|
|
abs_img_dir = img_dir_path_obj.resolve()
|
|
|
|
if self.run_args.debug:
|
|
print(Fore.CYAN + f" [Debug] Using IMAGE_FOLDER_PATH: '{IMAGE_FOLDER_PATH}' (Absolute: {abs_img_dir})")
|
|
|
|
try:
|
|
abs_img_dir.mkdir(parents=True, exist_ok=True)
|
|
if self.run_args.debug:
|
|
print(Fore.CYAN + f" [Debug] Ensured image directory exists: {abs_img_dir} (Status: {abs_img_dir.is_dir()})")
|
|
except Exception as e_mkdir:
|
|
print(Fore.RED + f" ERROR: Could not create directory {abs_img_dir}: {e_mkdir}")
|
|
if self.run_args.debug: traceback.print_exc()
|
|
# Do not proceed if directory creation fails
|
|
input(Fore.CYAN + "Press Enter to acknowledge and continue...")
|
|
return
|
|
|
|
|
|
unique_id = uuid.uuid4().hex[:8]
|
|
tmp_image_path_obj = abs_img_dir / f"page_{page_num_to_view}_{unique_id}.png"
|
|
image_created_in_managed_folder = True
|
|
else:
|
|
fd, temp_path_str = tempfile.mkstemp(suffix=".png", prefix="pdf_page_img_")
|
|
os.close(fd)
|
|
tmp_image_path_obj = Path(temp_path_str)
|
|
if self.run_args.debug:
|
|
print(Fore.CYAN + f" [Debug] Using system temporary file: {tmp_image_path_obj.resolve()}")
|
|
|
|
resolved_save_path = tmp_image_path_obj.resolve()
|
|
if self.run_args.debug:
|
|
print(Fore.CYAN + f" [Debug] Attempting to save image to: {resolved_save_path}")
|
|
|
|
pix.save(str(resolved_save_path))
|
|
|
|
if resolved_save_path.exists() and resolved_save_path.is_file():
|
|
image_successfully_saved = True
|
|
if self.run_args.show_progress: # Print for normal progress too, not just debug
|
|
print(Fore.GREEN + f" Image for page {page_num_to_view} successfully saved to: {resolved_save_path}")
|
|
if self.run_args.debug:
|
|
print(Fore.CYAN + f" [Debug] File size: {resolved_save_path.stat().st_size} bytes")
|
|
else:
|
|
if self.run_args.show_progress:
|
|
print(Fore.RED + f" ERROR: Failed to save image to {resolved_save_path}. File does not exist after save attempt.")
|
|
|
|
except Exception as e_render_save:
|
|
if self.run_args.show_progress:
|
|
print(Fore.RED + f" Error during image rendering or saving: {e_render_save}")
|
|
if self.run_args.debug:
|
|
traceback.print_exc()
|
|
|
|
if image_successfully_saved and tmp_image_path_obj:
|
|
if self.run_args.show_progress:
|
|
print(Fore.CYAN + f"Attempting to open image with default application...")
|
|
try:
|
|
file_uri = tmp_image_path_obj.resolve().as_uri()
|
|
if self.run_args.debug:
|
|
print(Fore.CYAN + f" [Debug] Opening URI: {file_uri}")
|
|
|
|
opened_successfully = webbrowser.open(file_uri)
|
|
|
|
if self.run_args.debug: # More detailed feedback in debug mode
|
|
print(Fore.CYAN + f" [Debug] webbrowser.open() returned: {opened_successfully}")
|
|
|
|
if not opened_successfully:
|
|
if self.run_args.show_progress:
|
|
print(Fore.YELLOW + " webbrowser.open() reported failure (returned False or None).")
|
|
print(Fore.YELLOW + f" This often means no default application is configured for PNG files or your browser.")
|
|
elif self.run_args.show_progress:
|
|
print(Fore.GREEN + " Image hopefully opened. Check your applications.")
|
|
|
|
if self.run_args.show_progress:
|
|
print(Fore.YELLOW + f" If the image did not open, please manually open: {tmp_image_path_obj.resolve()}")
|
|
input(Fore.CYAN + "Press Enter after viewing image to continue...")
|
|
|
|
except Exception as e_open:
|
|
if self.run_args.show_progress:
|
|
print(Fore.RED + f" Could not open image using webbrowser: {e_open}")
|
|
print(Fore.YELLOW + " This could be due to your system's environment (e.g., missing 'xdg-utils' on Linux, no default PNG viewer).")
|
|
print(Fore.YELLOW + f" Please try opening the image manually: {tmp_image_path_obj.resolve()}")
|
|
if self.run_args.debug:
|
|
traceback.print_exc()
|
|
input(Fore.CYAN + "Press Enter to acknowledge and continue...")
|
|
elif tmp_image_path_obj :
|
|
if self.run_args.show_progress:
|
|
print(Fore.YELLOW + " Skipping attempt to open image as it was not saved successfully.")
|
|
input(Fore.CYAN + "Press Enter to continue...")
|
|
else:
|
|
if self.run_args.show_progress:
|
|
print(Fore.RED + " Cannot attempt to open image as image path was not determined.")
|
|
input(Fore.CYAN + "Press Enter to continue...")
|
|
|
|
finally:
|
|
if tmp_image_path_obj and tmp_image_path_obj.exists():
|
|
if image_created_in_managed_folder:
|
|
if self.run_args.debug:
|
|
print(Fore.CYAN + f" [Debug] Image '{tmp_image_path_obj.name}' remains in managed folder '{IMAGE_FOLDER_PATH}'.")
|
|
print(Fore.CYAN + f" [Debug] It will be cleared based on CLEAR_IMAGE_FOLDER_ON_END ({CLEAR_IMAGE_FOLDER_ON_END}).")
|
|
else:
|
|
try:
|
|
tmp_image_path_obj.unlink()
|
|
if self.run_args.debug:
|
|
print(Fore.CYAN + f" [Debug] Deleted system temporary image: {tmp_image_path_obj.resolve()}")
|
|
except Exception as e_unlink:
|
|
if self.run_args.debug:
|
|
print(Fore.YELLOW + f" Warning: Could not delete system temp image {tmp_image_path_obj.resolve()}: {e_unlink}")
|
|
elif tmp_image_path_obj and not tmp_image_path_obj.exists() and image_successfully_saved:
|
|
if self.run_args.debug:
|
|
print(Fore.RED + f" [Debug] Inconsistency: Image was marked saved, but {tmp_image_path_obj.resolve()} does not exist at cleanup (and wasn't a system temp explicitly deleted here).")
|
|
|
|
|
|
def _interactive_review_session(self, highlights_list):
|
|
if not highlights_list:
|
|
if self.run_args.show_progress : print(Fore.YELLOW + "No highlights to review.")
|
|
return []
|
|
reviewed_highlights = [dict(h) for h in highlights_list]
|
|
idx, num_highlights = 0, len(reviewed_highlights)
|
|
AVAILABLE_COLORS = ['yellow', 'green', 'blue', 'pink', 'other_color', 'unknown_color']
|
|
|
|
while 0 <= idx < num_highlights:
|
|
item = reviewed_highlights[idx]
|
|
print(Style.RESET_ALL + "\n" + "="*15 + f" Review HL {idx+1}/{num_highlights} (Page {item['page']}) " + "="*15)
|
|
|
|
current_color_display = self._get_color_display_codes(item['color'])
|
|
print(f"Color: {current_color_display}{item['color'].upper()}{Style.RESET_ALL}", end="")
|
|
if item['color'] == 'other_color' and item.get('raw_rgb_values'):
|
|
rgb = item['raw_rgb_values'][:3]
|
|
rgb_disp = tuple(int(c*255) if isinstance(c,float) else int(c) for c in rgb)
|
|
print(f" (RGB: {rgb_disp})", end="")
|
|
print()
|
|
|
|
print(f"Text: {item['text']}")
|
|
|
|
prompt_options = ["[N]ext", "[P]rev", "[U]p", "[M]ove Down", "[C]olor", "[E]dit", "[D]elete", "[O]pen Img", "[S]ave&Exit", "[Q]uit"]
|
|
action_prompt_str = Fore.CYAN + ", ".join(prompt_options) + "? > " + Style.RESET_ALL
|
|
action = input(action_prompt_str).lower().strip()
|
|
|
|
if action == 'n': idx = (idx + 1) % num_highlights if num_highlights > 0 else 0
|
|
elif action == 'p': idx = (idx - 1 + num_highlights) % num_highlights if num_highlights > 0 else 0
|
|
elif action == 'u':
|
|
if idx > 0:
|
|
reviewed_highlights.insert(idx - 1, reviewed_highlights.pop(idx))
|
|
idx -= 1
|
|
print(Fore.GREEN + "Moved up.")
|
|
else: print(Fore.YELLOW + "Already at the top.")
|
|
elif action == 'm':
|
|
if idx < num_highlights - 1:
|
|
reviewed_highlights.insert(idx + 1, reviewed_highlights.pop(idx))
|
|
idx += 1
|
|
print(Fore.GREEN + "Moved down.")
|
|
else: print(Fore.YELLOW + "Already at the bottom.")
|
|
elif action == 'c':
|
|
print("Available colors:", ", ".join(f"{i+1}.{self._get_color_display_codes(co)}{co.upper()}{Style.RESET_ALL}" for i,co in enumerate(AVAILABLE_COLORS)))
|
|
try:
|
|
choice_str = input(Fore.YELLOW + "Enter number for new color: " + Style.RESET_ALL)
|
|
if not choice_str: print(Fore.BLUE + "Color change cancelled (no input)."); continue
|
|
choice = int(choice_str) - 1
|
|
if 0 <= choice < len(AVAILABLE_COLORS):
|
|
item['color'] = AVAILABLE_COLORS[choice]
|
|
print(Fore.GREEN + f"Color changed to {AVAILABLE_COLORS[choice].upper()}.")
|
|
else: print(Fore.RED + "Invalid color choice.")
|
|
except ValueError: print(Fore.RED + "Invalid input. Please enter a number.")
|
|
elif action == 'e':
|
|
edit_prompt = Fore.YELLOW + "New text (blank=keep, 'CLEAR'=empty): > " + Style.RESET_ALL
|
|
new_text = input_with_prefill(edit_prompt, item['text'])
|
|
|
|
if new_text.strip().upper() == 'CLEAR':
|
|
item['text'] = ""
|
|
print(Fore.GREEN + "Text cleared.")
|
|
elif new_text == item['text'] or not new_text.strip() :
|
|
print(Fore.BLUE + "Text kept as is.")
|
|
else:
|
|
item['text'] = new_text
|
|
print(Fore.GREEN + "Text updated.")
|
|
elif action == 'd':
|
|
if input(Fore.RED + "Are you sure you want to delete this highlight? [y/N]: " + Style.RESET_ALL).lower() == 'y':
|
|
reviewed_highlights.pop(idx)
|
|
num_highlights = len(reviewed_highlights)
|
|
print(Fore.GREEN + "Highlight deleted.")
|
|
if num_highlights == 0:
|
|
print(Fore.YELLOW + "No more highlights to review."); break
|
|
if idx >= num_highlights: idx = num_highlights - 1
|
|
else: print(Fore.BLUE + "Deletion cancelled.")
|
|
elif action == 'o': self._view_page_image_interactively(item['page'])
|
|
elif action == 's':
|
|
print(Fore.GREEN + "Saving changes and exiting review session.")
|
|
break
|
|
elif action == 'q':
|
|
if input(Fore.RED+"Are you sure you want to quit review? Changes will not be saved. [y/N]: " + Style.RESET_ALL).lower()=='y':
|
|
print(Fore.YELLOW+"Quitting review session. Changes made in this session are DISCARDED.")
|
|
return highlights_list
|
|
else:
|
|
print(Fore.BLUE + "Quit cancelled.")
|
|
else: print(Fore.RED + "Invalid action. Please choose from the list.")
|
|
return reviewed_highlights
|
|
|
|
def _parse_specific_pages(self, pages_str, total_pages):
|
|
if not pages_str or pages_str.lower() == "all": return list(range(1, total_pages + 1))
|
|
parsed_pages = set()
|
|
try:
|
|
for part in pages_str.split(','):
|
|
part = part.strip();
|
|
if not part: continue
|
|
if '-' in part:
|
|
start_str, end_str = part.split('-', 1); start = int(start_str); end = int(end_str)
|
|
start = max(1, start); end = min(total_pages, end)
|
|
if start <= end: parsed_pages.update(range(start, end + 1))
|
|
else:
|
|
page_val = int(part)
|
|
if 1 <= page_val <= total_pages: parsed_pages.add(page_val)
|
|
return sorted(list(parsed_pages)) if parsed_pages else []
|
|
except ValueError as e:
|
|
if self.run_args.show_progress: print(Fore.YELLOW + f"⚠️ Invalid page range: {pages_str}. Error: {e}.")
|
|
return []
|
|
|
|
def _get_color_display_codes(self, color_name_str):
|
|
return {'yellow': Back.YELLOW + Fore.BLACK, 'green': Back.GREEN + Fore.BLACK,
|
|
'blue': Back.BLUE + Fore.WHITE, 'pink': Back.MAGENTA + Fore.WHITE,
|
|
'other_color': Back.WHITE + Fore.BLACK, 'unknown_color': Back.LIGHTBLACK_EX + Fore.WHITE
|
|
}.get(color_name_str.lower(), Back.LIGHTBLACK_EX + Fore.WHITE)
|
|
|
|
def display_results(self):
|
|
if not self.run_args.show_progress: return # Don't display if progress is off (e.g. silent)
|
|
|
|
print("\n" + Fore.CYAN + Style.BRIGHT + "="*30 + " EXTRACTED HIGHLIGHTS " + "="*30 + Style.RESET_ALL)
|
|
if not self.highlights_data: print("\n❌ No highlights extracted or all were deleted."); return
|
|
current_page = None
|
|
for item in self.highlights_data:
|
|
if item.get('page') != current_page:
|
|
current_page = item.get('page'); print(f"\n📄 {Style.BRIGHT}Page {current_page}{Style.RESET_ALL}\n" + "-"*25)
|
|
color_name = item.get('color', 'unknown_color')
|
|
color_code = self._get_color_display_codes(color_name)
|
|
num_segments = item.get('num_segments', 0)
|
|
segment_info = f" [{num_segments} segments]" if num_segments > 1 else ""
|
|
text_content = item.get('text', "*NO TEXT*")
|
|
display_color_name = color_name.upper()
|
|
if color_name == 'other_color':
|
|
raw_rgb = item.get('raw_rgb_values')
|
|
if raw_rgb and len(raw_rgb) >=3:
|
|
rgb_disp = tuple(int(c*255) if isinstance(c,float) else int(c) for c in raw_rgb[:3])
|
|
display_color_name += f" (RGB: {rgb_disp})"
|
|
print(f"🎨 {color_code}{display_color_name}{Style.RESET_ALL}{segment_info}")
|
|
print(f" \"{text_content}\""); print()
|
|
|
|
def save_to_json(self, output_path_str):
|
|
output_path = Path(output_path_str).resolve() # Resolve to absolute path for clarity
|
|
try:
|
|
output_path.parent.mkdir(parents=True, exist_ok=True)
|
|
if self.run_args.debug:
|
|
print(Fore.CYAN + f" [Debug] Ensured parent directory for JSON exists: {output_path.parent}")
|
|
except Exception as e_mkdir:
|
|
if self.run_args.show_progress: # Also show error if progress is on
|
|
print(Fore.RED + f"❌ Error creating directory for JSON output {output_path.parent}: {e_mkdir}")
|
|
if self.run_args.debug: traceback.print_exc()
|
|
return # Cannot save if directory cannot be made
|
|
|
|
data_to_save = {
|
|
'pdf_file_processed': str(self.pdf_path.name), 'pdf_full_path': str(self.pdf_path.resolve()),
|
|
'pages_processed_spec': self.run_args.pages if self.run_args.pages else DEFAULT_PAGES_TO_PROCESS,
|
|
'extraction_timestamp': time.strftime("%Y-%m-%d %H:%M:%S %Z"),
|
|
'total_highlights_extracted': len(self.highlights_data),
|
|
'settings_used': {
|
|
'clean_edges': self.run_args.clean_edges,
|
|
'show_diff_percentage': self.run_args.show_diff_percentage
|
|
},
|
|
'highlights_data': self.highlights_data }
|
|
try:
|
|
with open(output_path, 'w', encoding='utf-8') as f: json.dump(data_to_save, f, indent=2, ensure_ascii=False)
|
|
if self.run_args.show_progress: print(Fore.GREEN + f"💾 Data saved to {output_path}")
|
|
except IOError as e:
|
|
if self.run_args.show_progress: print(Fore.RED + f"❌ Error saving JSON to {output_path}: {e}")
|
|
if self.run_args.debug: traceback.print_exc()
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(
|
|
description="Enhanced PDF Highlight Extractor.",
|
|
formatter_class=argparse.RawTextHelpFormatter,
|
|
epilog=f"""Examples:
|
|
{sys.argv[0]} mydoc.pdf
|
|
{sys.argv[0]} mydoc.pdf -p "1,5-7" -i
|
|
{sys.argv[0]} -t -s --output-json results/test.json
|
|
{sys.argv[0]} doc.pdf -d
|
|
|
|
If interactive image viewing ('O' option) fails, try running with the -d (debug)
|
|
flag. This will print detailed information about image paths and creation steps.
|
|
Common issues include missing default PNG viewers or OS-level permission problems.
|
|
The IMAGE_FOLDER_PATH ('{IMAGE_FOLDER_PATH}') is relative to where you run the script.
|
|
""")
|
|
parser.add_argument("pdf_path_arg", nargs='?', default=None, help="Path to PDF. Prompts if not in test/silent mode & not provided.")
|
|
parser.add_argument("-p", "--pages", type=str, default=None, help=f"Pages (e.g., \"1,3-5\", \"all\"). Default: \"{DEFAULT_PAGES_TO_PROCESS}\".")
|
|
parser.add_argument("-i", "--interactive", action="store_true", help="Enable interactive review mode.")
|
|
parser.add_argument("-t", "--test", action="store_true", help=f"Test mode. Uses default PDF ('{DEFAULT_PDF_PATH}'), auto-saves JSON.")
|
|
parser.add_argument("-s", "--silent", action="store_true", help="Silent mode. Minimal output. Auto-saves JSON. Implies -t if no PDF path.")
|
|
parser.add_argument("-d", "--debug", action="store_true", help="Debug mode. Enables all detailed SHOW flags and prints more internal details.")
|
|
parser.add_argument("--output-json", type=str, default=None, help="Custom output JSON filename/path.")
|
|
|
|
cli_args = parser.parse_args()
|
|
|
|
effective_run_args = argparse.Namespace()
|
|
effective_run_args.debug = cli_args.debug
|
|
effective_run_args.silent = cli_args.silent
|
|
|
|
# Initialize based on global defaults
|
|
effective_run_args.show_timing = INITIAL_SHOW_TIMING
|
|
effective_run_args.show_progress = INITIAL_SHOW_PROGRESS
|
|
effective_run_args.show_raw_segments = INITIAL_SHOW_RAW_SEGMENTS
|
|
effective_run_args.show_extraction_details = INITIAL_SHOW_EXTRACTION_DETAILS
|
|
effective_run_args.show_rect_details = INITIAL_SHOW_RECT_DETAILS
|
|
effective_run_args.show_diff_percentage = INITIAL_SHOW_DIFF_PERCENTAGE
|
|
effective_run_args.clean_edges = INITIAL_CLEAN_EDGES
|
|
|
|
# Override show flags based on debug or silent
|
|
if effective_run_args.debug:
|
|
for key in ['show_timing', 'show_progress', 'show_raw_segments', 'show_extraction_details', 'show_rect_details', 'show_diff_percentage']:
|
|
setattr(effective_run_args, key, True) # Debug enables all these
|
|
|
|
if effective_run_args.silent:
|
|
for key in ['show_timing', 'show_progress', 'show_raw_segments', 'show_extraction_details', 'show_rect_details', 'show_diff_percentage']:
|
|
setattr(effective_run_args, key, False) # Silent disables all these
|
|
effective_run_args.interactive = False
|
|
else: # Not silent
|
|
effective_run_args.interactive = cli_args.interactive
|
|
|
|
effective_run_args.pages = cli_args.pages
|
|
|
|
start_time = time.time()
|
|
if effective_run_args.show_progress: print(Fore.MAGENTA + Style.BRIGHT + "🎨 PDF Highlight Extractor 🎨" + Style.RESET_ALL)
|
|
if effective_run_args.debug:
|
|
print(Fore.CYAN + f" [Debug] Current Working Directory: {Path.cwd()}")
|
|
print(Fore.CYAN + f" [Debug] Effective Run Arguments: {effective_run_args}")
|
|
|
|
|
|
if IMAGE_FOLDER_PATH and CLEAR_IMAGE_FOLDER_ON_START:
|
|
_clear_png_files_in_folder(IMAGE_FOLDER_PATH, effective_run_args)
|
|
|
|
pdf_path_to_use = None
|
|
if cli_args.test: pdf_path_to_use = DEFAULT_PDF_PATH
|
|
elif cli_args.pdf_path_arg: pdf_path_to_use = cli_args.pdf_path_arg
|
|
elif cli_args.silent: pdf_path_to_use = DEFAULT_PDF_PATH
|
|
else:
|
|
pdf_path_input = input(f"📄 PDF path (Enter for default '{DEFAULT_PDF_PATH}'): ").strip().strip('"')
|
|
pdf_path_to_use = pdf_path_input if pdf_path_input else DEFAULT_PDF_PATH
|
|
|
|
if not pdf_path_to_use:
|
|
if effective_run_args.show_progress: print(Fore.RED + "❌ No PDF path specified. Exiting.")
|
|
sys.exit(1)
|
|
|
|
resolved_path = Path(pdf_path_to_use).resolve()
|
|
if not resolved_path.exists() or not resolved_path.is_file():
|
|
if effective_run_args.show_progress: print(Fore.RED + f"❌ PDF not found or is not a file: {resolved_path}")
|
|
sys.exit(1)
|
|
|
|
doc_for_processing = None
|
|
try:
|
|
doc_for_processing = fitz.open(str(resolved_path))
|
|
extractor = EnhancedPDFHighlightExtractor(resolved_path, effective_run_args, main_doc_for_image_view=doc_for_processing)
|
|
extractor.extract_highlights(doc_for_processing)
|
|
|
|
if not effective_run_args.interactive and effective_run_args.show_progress:
|
|
extractor.display_results()
|
|
elif effective_run_args.interactive and effective_run_args.show_progress:
|
|
if input(Fore.CYAN+"Interactive session ended. Display final results? [Y/n]: " + Style.RESET_ALL).lower().strip()!='n':
|
|
extractor.display_results()
|
|
|
|
json_output_path_str = cli_args.output_json if cli_args.output_json else str(resolved_path.parent / f"{resolved_path.stem}_highlights.json")
|
|
|
|
perform_save = False
|
|
if cli_args.test or cli_args.silent:
|
|
perform_save = True
|
|
elif effective_run_args.show_progress: # Only prompt if not silent
|
|
save_prompt_input = input(f"💾 Save to JSON? (Enter for default '{json_output_path_str}', type 'skip' to not save, or enter a custom path): " + Style.RESET_ALL).strip()
|
|
if save_prompt_input.lower() != 'skip':
|
|
perform_save = True
|
|
if save_prompt_input:
|
|
json_output_path_str = save_prompt_input
|
|
|
|
if perform_save:
|
|
if extractor.highlights_data:
|
|
extractor.save_to_json(json_output_path_str)
|
|
elif effective_run_args.show_progress:
|
|
print(Fore.YELLOW + "No highlights were extracted or kept, so JSON file was not saved.")
|
|
elif effective_run_args.show_progress:
|
|
print(Fore.BLUE + "Skipped saving highlights to JSON.")
|
|
|
|
except Exception as e:
|
|
if effective_run_args.show_progress: print(Fore.RED+Style.BRIGHT+f"💥 An critical error occurred in the main execution: {e}")
|
|
if effective_run_args.debug:
|
|
traceback.print_exc()
|
|
finally:
|
|
if doc_for_processing:
|
|
doc_for_processing.close()
|
|
|
|
if IMAGE_FOLDER_PATH and CLEAR_IMAGE_FOLDER_ON_END:
|
|
_clear_png_files_in_folder(IMAGE_FOLDER_PATH, effective_run_args)
|
|
|
|
if effective_run_args.show_timing:
|
|
print(Fore.CYAN + f"\n⏱️ Total execution time: {time.time()-start_time:.2f} seconds")
|
|
|
|
if __name__ == '__main__':
|
|
main()
|