Spaces:
Sleeping
Sleeping
| import fitz # PyMuPDF | |
| import numpy as np | |
| import re | |
| from utils.decompose import Decomposer | |
| def clean_chinese_text(text): | |
| """Clean up text by removing spaces between Chinese characters.""" | |
| # Remove spaces between Chinese characters | |
| text = re.sub(r'([\u4e00-\u9fff])\s+([\u4e00-\u9fff])', r'\1\2', text) | |
| # Remove spaces before and after Chinese punctuation | |
| text = re.sub(r'\s+([\u3000-\u303f\uff00-\uffef])', r'\1', text) | |
| text = re.sub(r'([\u3000-\u303f\uff00-\uffef])\s+', r'\1', text) | |
| return text.strip() | |
| def categorize_highlight(color): | |
| """Categorizes highlights based on the closest color match using Euclidean distance.""" | |
| # customize the categories of highlights as you | |
| color_mapping = { | |
| (0.5608, 0.8706, 0.9765): "Ideas & Insights", # Light Blue | |
| (1.0, 0.9412, 0.4): "General Notes", # Yellow | |
| (0.4902, 0.9412, 0.4): "Action Items / To-Do", # Green | |
| (0.9686, 0.6, 0.8196): "Quotes & References", # Pink | |
| (0.9216, 0.2863, 0.2863): "Critical Issues / Warnings" # Red | |
| } | |
| # Convert color to a NumPy array for distance calculation | |
| color_array = np.array(color) | |
| # Find the closest color in the mapping using Euclidean distance | |
| best_match = min(color_mapping.keys(), key=lambda ref_color: sum((color_array[i] - ref_color[i]) ** 2 for i in range(len(color_array))) ** 0.5) | |
| return color_mapping[best_match] | |
| def clean_text_by_punctuation(text): | |
| """Clean text by removing content after the last proper punctuation mark.""" | |
| # Define proper ending punctuation marks (both Chinese and English) | |
| ending_punctuation = {'.', '。', '?', '？', '!', '！'} | |
| # Find the last occurrence of any ending punctuation | |
| last_punct_index = -1 | |
| for i, char in enumerate(text): | |
| if char in ending_punctuation: | |
| last_punct_index = i | |
| # If no proper ending punctuation found, return empty string | |
| if last_punct_index == -1: | |
| return "" | |
| # Return text up to and including the last punctuation mark | |
| return text[:last_punct_index + 1] | |
| def extract_highlights(pdf_path): | |
| """ | |
| Extract all highlights from a PDF file. | |
| Args: | |
| pdf_path (str): Path to the PDF file | |
| Returns: | |
| list: List of dictionaries containing highlight information | |
| """ | |
| highlights = [] | |
| try: | |
| pdf_document = fitz.open(pdf_path) | |
| min_y, max_y, single_y = Decomposer(pdf_document).run() | |
| highlights=[] | |
| for page_num in range(pdf_document.page_count): | |
| page = pdf_document[page_num] | |
| for annot in page.annots(): | |
| if annot.type[0] == 8: # Highlight annotation | |
| # Extract highlighted text | |
| highlight_text = page.get_text("text", clip=annot.rect, sort=True, flags=1).strip() | |
| highlight_text = highlight_text.encode("utf-8", "ignore").decode("utf-8") | |
| # Extract annotation color | |
| color_rgb = annot.colors.get("stroke", [0, 0, 0]) # Default black if undefined | |
| category = categorize_highlight(color_rgb) | |
| # Extract popup comment if it exists | |
| comment = annot.info.get("content", "").strip() if annot.has_popup else "" | |
| # Store structured highlight data | |
| if highlight_text: | |
| for chunk in highlight_text.split("\n\n"): # Handle consecutive highlights | |
| cleaned_chunk = clean_text_by_punctuation(clean_chinese_text(chunk)) | |
| if not cleaned_chunk: # Skip if no valid text after cleaning | |
| continue | |
| if (highlights and | |
| highlights[-1]['page'] == page_num and | |
| highlights[-1]['rect'][3] > (max_y - single_y) and | |
| annot.rect[1]< (min_y + single_y) | |
| ): # Handle highlights over page | |
| highlights[-1]['text'] += cleaned_chunk | |
| else: | |
| highlights.append({ | |
| "page": page_num + 1, | |
| "text": cleaned_chunk, | |
| "category": category, | |
| "comment": comment, | |
| "rect":annot.rect | |
| }) | |
| pdf_document.close() | |
| return highlights | |
| except Exception as e: | |
| print(f"Error processing PDF: {str(e)}") | |
| return [] | |
| def main(): | |
| # Example usage | |
| pdf_path = "OneDrive-2025-04-15/专题讨论_线下研讨会DeepSeek效应追踪AI产业的持续变革-12Mar2025_zho.pdf" | |
| highlights = extract_highlights(pdf_path) | |
| # Create markdown content | |
| markdown_content = "# PDF Highlights Summary\n\n" | |
| markdown_content += f"Total highlights found: {len(highlights)}\n\n" | |
| # Group highlights by category | |
| highlights_by_category = {} | |
| for highlight in highlights: | |
| category = highlight['category'] | |
| if category not in highlights_by_category: | |
| highlights_by_category[category] = [] | |
| highlights_by_category[category].append(highlight) | |
| # Add highlights grouped by category | |
| for category, category_highlights in highlights_by_category.items(): | |
| markdown_content += f"## {category}\n\n" | |
| for highlight in category_highlights: | |
| markdown_content += f"### Page {highlight['page']}\n\n" | |
| markdown_content += f"{highlight['text']}\n\n" | |
| if highlight['comment']: | |
| markdown_content += f"> {highlight['comment']}\n\n" | |
| markdown_content += "---\n\n" | |
| # Save to markdown file | |
| output_file = "highlights_summary.md" | |
| with open(output_file, "w", encoding="utf-8") as f: | |
| f.write(markdown_content) | |
| print(f"Highlights summary has been saved to {output_file}") | |
| if __name__ == "__main__": | |
| main() | |