Spaces:
Sleeping
Sleeping
| # annotations.py | |
| import fitz # PyMuPDF | |
| from typing import List, Dict, Any, Tuple | |
| import language_tool_python | |
| import io | |
| def extract_pdf_text(file) -> str: | |
| """Extracts full text from a PDF file using PyMuPDF.""" | |
| try: | |
| doc = fitz.open(stream=file.read(), filetype="pdf") if not isinstance(file, str) else fitz.open(file) | |
| full_text = "" | |
| for page in doc: | |
| full_text += page.get_text("text") + "\n" | |
| doc.close() | |
| return full_text | |
| except Exception as e: | |
| print(f"Error extracting text from PDF: {e}") | |
| return "" | |
| def check_language_issues(full_text: str) -> Dict[str, Any]: | |
| """Check for language issues using LanguageTool.""" | |
| language_tool = language_tool_python.LanguageTool('en-US') | |
| matches = language_tool.check(full_text) | |
| issues = [] | |
| for match in matches: | |
| issues.append({ | |
| "message": match.message, | |
| "context": match.context, | |
| "suggestions": match.replacements[:3] if match.replacements else [], | |
| "category": match.category, | |
| "rule_id": match.ruleId, | |
| "offset": match.offset, | |
| "length": match.errorLength | |
| }) | |
| return { | |
| "total_issues": len(issues), | |
| "issues": issues | |
| } | |
| def highlight_issues_in_pdf(file, language_matches: List[Dict[str, Any]]) -> bytes: | |
| """ | |
| Highlights language issues in the PDF and returns the annotated PDF as bytes. | |
| This function maps LanguageTool matches to specific words in the PDF | |
| and highlights those words. | |
| """ | |
| try: | |
| # Open the PDF | |
| doc = fitz.open(stream=file.read(), filetype="pdf") if not isinstance(file, str) else fitz.open(file) | |
| # Extract words with positions from each page | |
| word_list = [] # List of tuples: (page_number, word, x0, y0, x1, y1) | |
| page_text = "" | |
| for page_number in range(len(doc)): | |
| page = doc[page_number] | |
| words = page.get_text("words") # List of tuples: (x0, y0, x1, y1, "word", block_no, line_no, word_no) | |
| for w in words: | |
| word_text = w[4] | |
| word_list.append((page_number, word_text, w[0], w[1], w[2], w[3])) | |
| # Concatenate all words to form the full text | |
| concatenated_text = " ".join([w[1] for w in word_list]) | |
| # Iterate over each language issue | |
| for issue in language_matches: | |
| offset = issue["offset"] | |
| length = issue["length"] | |
| error_text = concatenated_text[offset:offset+length] | |
| # Find the words that fall within the error span | |
| current_pos = 0 | |
| target_words = [] | |
| for word in word_list: | |
| word_text = word[1] | |
| word_length = len(word_text) + 1 # +1 for the space | |
| if current_pos + word_length > offset and current_pos < offset + length: | |
| target_words.append(word) | |
| current_pos += word_length | |
| # Add highlight annotations to the target words | |
| for target in target_words: | |
| page_num, word_text, x0, y0, x1, y1 = target | |
| page = doc[page_num] | |
| # Define a rectangle around the word | |
| rect = fitz.Rect(x0, y0, x1, y1) | |
| # Add a highlight annotation | |
| highlight = page.add_highlight_annot(rect) | |
| highlight.set_colors(stroke=(1, 1, 0)) # Yellow color | |
| highlight.update() | |
| # Save annotated PDF to bytes | |
| byte_stream = io.BytesIO() | |
| doc.save(byte_stream) | |
| annotated_pdf_bytes = byte_stream.getvalue() | |
| doc.close() | |
| return annotated_pdf_bytes | |
| except Exception as e: | |
| print(f"Error in highlighting PDF: {e}") | |
| return b"" | |
| def analyze_pdf(file) -> Tuple[Dict[str, Any], bytes]: | |
| """Analyzes the PDF for language issues and returns results and annotated PDF.""" | |
| try: | |
| full_text = extract_pdf_text(file) | |
| if not full_text: | |
| return {"error": "Failed to extract text from PDF."}, None | |
| language_issues = check_language_issues(full_text) | |
| issues = language_issues.get("issues", []) | |
| annotated_pdf = highlight_issues_in_pdf(file, issues) if issues else None | |
| return language_issues, annotated_pdf | |
| except Exception as e: | |
| return {"error": str(e)}, None |