Text-to-Document-Generation/PDF-Redaction-API
0
1"""2PDF Redaction module using NER3"""4from pdf2image import convert_from_path5import pytesseract6from pypdf import PdfReader, PdfWriter7from pypdf.generic import DictionaryObject, ArrayObject, NameObject, NumberObject8from transformers import pipeline, AutoTokenizer, AutoModelForTokenClassification9from typing import List, Dict, Optional10import logging11 12logger = logging.getLogger(__name__)13 14 15class PDFRedactor:16 """PDF Redaction using Named Entity Recognition"""17 18 def __init__(self, model_name: str = "./model"):19 """20 Initialize the PDF Redactor21 22 Args:23 model_name: HuggingFace model name for NER24 """25 self.model_name = model_name26 self.ner_pipeline = None27 self._load_model()28 29 def _load_model(self):30 """Load the NER model"""31 try:32 logger.info(f"Loading NER model: {self.model_name}")33 tokenizer = AutoTokenizer.from_pretrained(self.model_name)34 model = AutoModelForTokenClassification.from_pretrained(self.model_name)35 36 self.ner_pipeline = pipeline("token-classification", model=model, 37 tokenizer=tokenizer)38 logger.info("NER model loaded successfully")39 except Exception as e:40 logger.error(f"Error loading NER model: {str(e)}")41 raise42 43 def is_model_loaded(self) -> bool:44 """Check if the model is loaded"""45 return self.ner_pipeline is not None46 47 def perform_ocr(self, pdf_path: str, dpi: int = 300) -> List[Dict]:48 """49 Perform OCR on PDF and extract word bounding boxes50 51 Args:52 pdf_path: Path to the PDF file53 dpi: DPI for PDF to image conversion54 55 Returns:56 List of word data with bounding boxes and image dimensions57 """58 logger.info(f"Starting OCR on {pdf_path} at {dpi} DPI")59 all_words_data = []60 61 try:62 images = convert_from_path(pdf_path, dpi=dpi)63 logger.info(f"Converted PDF to {len(images)} images")64 65 for page_num, image in enumerate(images):66 # Get image dimensions67 image_width, image_height = image.size68 69 # Perform OCR70 data = pytesseract.image_to_data(image, output_type=pytesseract.Output.DICT)71 logger.info(f"OCR data: {data['text']}")72 73 num_words = len(data['text'])74 for i in range(num_words):75 word_text = data['text'][i].strip()76 confidence = int(data['conf'][i])77 78 # Filter out empty or low-confidence words79 if word_text and confidence > 0:80 all_words_data.append({81 'text': word_text,82 'box': (data['left'][i], data['top'][i], 83 data['width'][i], data['height'][i]),84 'page': page_num + 1,85 'confidence': confidence,86 'image_width': image_width,87 'image_height': image_height88 })89 90 logger.info(f"Processed page {page_num + 1}: {len([w for w in all_words_data if w['page'] == page_num + 1])} words")91 92 logger.info(f"OCR complete: {len(all_words_data)} total words extracted")93 return all_words_data94 95 except Exception as e:96 logger.error(f"Error during OCR: {str(e)}")97 raise98 99 def run_ner(self, text: str) -> List[Dict]:100 """101 Run NER on text102 103 Args:104 text: Input text105 106 Returns:107 List of identified entities108 """109 if not self.ner_pipeline:110 raise RuntimeError("NER model not loaded")111 112 logger.info(f"Running NER on text of length {len(text)}")113 114 try:115 results = self.ner_pipeline(text)116 logger.info(f"NER identified {len(results)} entities")117 return results118 except Exception as e:119 logger.error(f"Error during NER: {str(e)}")120 raise121 122 def map_entities_to_boxes(self, ner_results: List[Dict], 123 ocr_data: List[Dict]) -> List[Dict]:124 """125 Map NER entities to OCR bounding boxes126 127 Args:128 ner_results: List of NER entities129 ocr_data: List of OCR word data130 131 Returns:132 List of mapped entities with bounding boxes133 """134 logger.info("Mapping NER entities to OCR bounding boxes")135 mapped_entities = []136 137 # Create character span mapping138 ocr_word_char_spans = []139 current_char_index = 0140 141 for ocr_data_idx, word_info in enumerate(ocr_data):142 word_text = word_info['text']143 length = len(word_text)144 145 ocr_word_char_spans.append({146 'ocr_data_idx': ocr_data_idx,147 'start_char': current_char_index,148 'end_char': current_char_index + length149 })150 current_char_index += length + 1151 152 # Map each NER entity to OCR words153 for ner_entity in ner_results:154 ner_entity_type = ner_entity['entity']155 ner_start = ner_entity['start']156 ner_end = ner_entity['end']157 ner_word = ner_entity['word']158 159 matching_ocr_words = []160 161 for ocr_word_span in ocr_word_char_spans:162 ocr_start = ocr_word_span['start_char']163 ocr_end = ocr_word_span['end_char']164 165 # Check for overlap166 if max(ocr_start, ner_start) < min(ocr_end, ner_end):167 matching_ocr_words.append(ocr_data[ocr_word_span['ocr_data_idx']])168 169 if matching_ocr_words:170 mapped_entities.append({171 'entity_type': ner_entity_type,172 'entity_text': ner_word,173 'words': matching_ocr_words174 })175 176 logger.info(f"Mapped {len(mapped_entities)} entities to bounding boxes")177 return mapped_entities178 179 def create_redacted_pdf(self, original_pdf_path: str, 180 mapped_entities: List[Dict],181 output_path: str) -> str:182 """183 Create redacted PDF with black rectangles over entities184 185 Args:186 original_pdf_path: Path to original PDF187 mapped_entities: List of entities with bounding boxes188 output_path: Path for output PDF189 190 Returns:191 Path to redacted PDF192 """193 logger.info(f"Creating redacted PDF: {output_path}")194 195 try:196 reader = PdfReader(original_pdf_path)197 writer = PdfWriter()198 199 for page_num in range(len(reader.pages)):200 page = reader.pages[page_num]201 media_box = page.mediabox202 page_width = float(media_box.width)203 page_height = float(media_box.height)204 205 writer.add_page(page)206 207 page_entities = 0208 for entity_info in mapped_entities:209 for word_info in entity_info['words']:210 if word_info['page'] == page_num + 1:211 x, y, w, h = word_info['box']212 213 # Get image dimensions214 image_width = word_info['image_width']215 image_height = word_info['image_height']216 217 # Scale coordinates218 scale_x = page_width / image_width219 scale_y = page_height / image_height220 221 x_scaled = x * scale_x222 y_scaled = y * scale_y223 w_scaled = w * scale_x224 h_scaled = h * scale_y225 226 # Convert to PDF coordinates227 llx = x_scaled228 lly = page_height - (y_scaled + h_scaled)229 urx = x_scaled + w_scaled230 ury = page_height - y_scaled231 232 # Create redaction annotation233 redaction_annotation = DictionaryObject()234 redaction_annotation.update({235 NameObject("/Type"): NameObject("/Annot"),236 NameObject("/Subtype"): NameObject("/Square"),237 NameObject("/Rect"): ArrayObject([238 NumberObject(llx),239 NumberObject(lly),240 NumberObject(urx),241 NumberObject(ury),242 ]),243 NameObject("/C"): ArrayObject([244 NumberObject(0), NumberObject(0), NumberObject(0)245 ]),246 NameObject("/IC"): ArrayObject([247 NumberObject(0), NumberObject(0), NumberObject(0)248 ]),249 NameObject("/BS"): DictionaryObject({250 NameObject("/W"): NumberObject(0)251 })252 })253 254 writer.add_annotation(page_number=page_num, 255 annotation=redaction_annotation)256 page_entities += 1257 258 logger.info(f"Page {page_num + 1}: Added {page_entities} redactions")259 260 # Write output261 with open(output_path, "wb") as output_file:262 writer.write(output_file)263 264 logger.info(f"Redacted PDF created successfully: {output_path}")265 return output_path266 267 except Exception as e:268 logger.error(f"Error creating redacted PDF: {str(e)}")269 raise270 271 def redact_document(self, pdf_path: str, output_path: str,272 dpi: int = 300,273 entity_filter: Optional[List[str]] = None) -> Dict:274 """275 Complete redaction pipeline276 277 Args:278 pdf_path: Path to input PDF279 output_path: Path for output PDF280 dpi: DPI for OCR281 entity_filter: List of entity types to redact (None = all)282 283 Returns:284 Dictionary with redaction results285 """286 logger.info(f"Starting redaction pipeline for {pdf_path}")287 288 # Step 1: OCR289 ocr_data = self.perform_ocr(pdf_path, dpi)290 291 # Step 2: Extract text292 full_text = " ".join([word['text'] for word in ocr_data])293 294 # Step 3: NER295 ner_results = self.run_ner(full_text)296 297 # Step 4: Map entities to boxes298 mapped_entities = self.map_entities_to_boxes(ner_results, ocr_data)299 300 # Step 5: Filter entities if requested301 if entity_filter:302 mapped_entities = [303 e for e in mapped_entities 304 if e['entity_type'] in entity_filter305 ]306 logger.info(f"Filtered to {len(mapped_entities)} entities of types: {entity_filter}")307 308 # Step 6: Create redacted PDF309 self.create_redacted_pdf(pdf_path, mapped_entities, output_path)310 311 return {312 'output_path': output_path,313 'total_words': len(ocr_data),314 'total_entities': len(ner_results),315 'redacted_entities': len(mapped_entities),316 'entities': mapped_entities317 }318