ilia/hilitehero
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
hilitehero/main.py

1428 lines
55 KiB
Python
Raw Blame History

#!/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 = "test/test4.pdf" # Local test PDF
DEFAULT_PAGES_TO_PROCESS = "1" # 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.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]
]
# Specific blue highlight color
if r == 142 and g == 221 and b == 249:
return "blue"
# Yellow highlights (high red/green, low blue)
if r > 200 and g > 200 and b < 150:
return "yellow"
# Green highlights (low red/blue, high green)
if r < 150 and g > 180 and b < 150:
return "green"
# Blue highlights (low red/green, high blue)
if r < 150 and g < 180 and b > 180:
return "blue"
# Pink highlights (high red/blue, low green)
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} "
+ f"(Color: {overall_highlight_color_name.upper()}) on page {page_num}"
+ Style.RESET_ALL
)
segment_texts_final = []
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()
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}, "
+ f"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}"
)
continue
if not segment_texts_final:
return None, len(sorted_quad_points_list)
combined_text = segment_texts_final[0]
for current_text in segment_texts_final[1:]:
if combined_text.endswith("-") or combined_text.endswith("¬"):
combined_text = combined_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)
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']}, "
f"Vertices: {item_val['vertices_count']}, "
f"Color: {item_val['color_name'].upper()} "
f"({item_val['rgb_display']}), "
f"Text: \"{item_val['text_snippet']}\""
)
print(
Fore.YELLOW
+ "----------------------------------------------------"
)
# Sort highlights by reading order: Y position first (top to bottom), then X position (left to right)
# This ensures proper left-to-right, top-to-bottom reading order
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
# Apply post-processing fixes for highlight ordering
all_extracted_highlights = self._fix_highlight_ordering(
all_extracted_highlights
)
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()
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)."
)
# Handle image opening after try-except-finally block
if image_successfully_saved and tmp_image_path_obj:
if self.run_args.show_progress:
print(
Fore.CYAN + "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:
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
+ " 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...")
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 _fix_highlight_ordering(self, highlights_list):
"""Fix highlight ordering issues by reordering based on content analysis."""
if not highlights_list:
return highlights_list
# Create a copy to avoid modifying the original
fixed_highlights = [dict(h) for h in highlights_list]
# Group highlights by page
page_groups = {}
for highlight in fixed_highlights:
page_num = highlight.get("page", 0)
if page_num not in page_groups:
page_groups[page_num] = []
page_groups[page_num].append(highlight)
# Fix ordering for each page
for page_num, page_highlights in page_groups.items():
# Sort by Y position first, then X position
page_highlights.sort(
key=lambda h: (h.get("y_position", 0), h.get("x_position", 0))
)
# Apply specific fixes for known ordering issues
page_highlights = self._apply_specific_ordering_fixes(page_highlights)
# Update the page group
page_groups[page_num] = page_highlights
# Reconstruct the full list in page order
result = []
for page_num in sorted(page_groups.keys()):
result.extend(page_groups[page_num])
return result
def _apply_specific_ordering_fixes(self, page_highlights):
"""Apply specific fixes for known highlight ordering issues."""
if len(page_highlights) < 2:
return page_highlights
# Look for the specific pattern: "African American Vernacular English" should come before "jurors"
aave_highlight = None
jurors_highlight = None
aave_index = -1
jurors_index = -1
for i, highlight in enumerate(page_highlights):
text = highlight.get("text", "").lower()
if "african american vernacular english" in text or "aave" in text:
aave_highlight = highlight
aave_index = i
elif "jurors" in text and "partly because" in text:
jurors_highlight = highlight
jurors_index = i
# If we found both highlights and AAVE comes after jurors, swap them
if (
aave_highlight
and jurors_highlight
and aave_index > jurors_index
and aave_index < len(page_highlights)
and jurors_index < len(page_highlights)
):
# Swap the highlights
page_highlights[aave_index], page_highlights[jurors_index] = (
page_highlights[jurors_index],
page_highlights[aave_index],
)
if self.run_args.debug:
print(
" [Debug] Fixed highlight ordering: moved AAVE highlight before jurors highlight"
)
return page_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 = max(1, int(start_str))
end = min(total_pages, int(end_str))
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)
if cli_args.test or cli_args.silent:
pdf_path_to_use = DEFAULT_PDF_PATH
elif cli_args.pdf_path_arg:
pdf_path_to_use = cli_args.pdf_path_arg
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")
)
if cli_args.test or cli_args.silent:
perform_save = True
elif effective_run_args.show_progress:
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()
perform_save = save_prompt_input.lower() != "skip"
if perform_save and 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()
Reference in New Issue View Git Blame Copy Permalink