import io import re from typing import List, Dict, Any import numpy as np import cv2 import pdfplumber import pytesseract from PIL import Image from pdf2image import convert_from_bytes from sqlalchemy.orm import Session from models.template_models import DocumentLayout, TemplateDocument from schemas.template_schemas import DocumentLayoutSchema class DocumentProcessor: def __init__(self, db: Session): self.db = db def process_file(self, document_id: int, filename: str, file_bytes: bytes) -> List[DocumentLayout]: document = self.db.query(TemplateDocument).filter(TemplateDocument.pk_document_id == document_id).first() if not document: raise ValueError("Document not found") layouts = [] if filename.lower().endswith('.pdf'): layouts = self._process_pdf(document_id, file_bytes) elif filename.lower().endswith(('.png', '.jpg', '.jpeg', '.tiff')): layouts = self._process_image(document_id, file_bytes) # Update page count if layouts: document.page_count = max([l.page_no for l in layouts]) document.status = "processed" self.db.commit() return layouts def _process_pdf(self, document_id: int, file_bytes: bytes) -> List[DocumentLayout]: layouts = [] try: with pdfplumber.open(io.BytesIO(file_bytes)) as pdf: sequence_no = 1 for page_no, page in enumerate(pdf.pages, start=1): # Check if it has text text = page.extract_text() if text and text.strip(): # Extract words with bounding boxes words = page.extract_words() # Group words into blocks using simple proximity heuristics blocks = self._group_words_to_blocks(words, page.width, page.height) for block in blocks: db_layout = DocumentLayout( fk_document_id=document_id, page_no=page_no, text_value=block['text'], block_type=block['type'], x_coordinate=block['x0'], y_coordinate=block['top'], width=block['x1'] - block['x0'], height=block['bottom'] - block['top'], confidence=100.0, sequence_no=sequence_no ) self.db.add(db_layout) layouts.append(db_layout) sequence_no += 1 else: # Scanned PDF page -> convert to image and process # Handled separately to avoid complexity in this mock pass except Exception as e: print(f"Error processing PDF with pdfplumber: {e}") if not layouts: # Fallback to image-based processing for scanned PDFs images = convert_from_bytes(file_bytes, dpi=300) sequence_no = 1 for page_no, img in enumerate(images, start=1): img_byte_arr = io.BytesIO() img.save(img_byte_arr, format='PNG') page_layouts = self._process_image(document_id, img_byte_arr.getvalue(), page_no=page_no, start_sequence=sequence_no) layouts.extend(page_layouts) sequence_no += len(page_layouts) self.db.commit() return layouts def _process_image(self, document_id: int, file_bytes: bytes, page_no: int = 1, start_sequence: int = 1) -> List[DocumentLayout]: layouts = [] try: # Decode image nparr = np.frombuffer(file_bytes, np.uint8) img = cv2.imdecode(nparr, cv2.IMREAD_COLOR) # Preprocessing gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) # Denoise & Threshold gray = cv2.medianBlur(gray, 3) gray = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)[1] # OCR using PyTesseract with bounding boxes data = pytesseract.image_to_data(gray, output_type=pytesseract.Output.DICT) sequence_no = start_sequence # Grouping words into lines/blocks simplified height, width = gray.shape blocks = self._tesseract_data_to_blocks(data, width, height) for block in blocks: db_layout = DocumentLayout( fk_document_id=document_id, page_no=page_no, text_value=block['text'], block_type=block['type'], x_coordinate=block['x'], y_coordinate=block['y'], width=block['w'], height=block['h'], confidence=block['conf'], sequence_no=sequence_no ) self.db.add(db_layout) layouts.append(db_layout) sequence_no += 1 except Exception as e: print(f"Error processing image with tesseract: {e}") self.db.commit() return layouts def _group_words_to_blocks(self, words: List[Dict], page_width: float, page_height: float) -> List[Dict]: """Group nearby words into logical blocks and classify them.""" # Extremely simplified grouping by Y coordinate proximity blocks = [] if not words: return blocks # Sort words by Y then X words.sort(key=lambda w: (w['top'], w['x0'])) current_block = { 'text': words[0]['text'], 'x0': words[0]['x0'], 'top': words[0]['top'], 'x1': words[0]['x1'], 'bottom': words[0]['bottom'] } for word in words[1:]: # If word is roughly on the same line (y diff is small) and close horizontally if abs(word['top'] - current_block['top']) < 10 and (word['x0'] - current_block['x1']) < 50: current_block['text'] += ' ' + word['text'] current_block['x1'] = word['x1'] current_block['top'] = min(current_block['top'], word['top']) current_block['bottom'] = max(current_block['bottom'], word['bottom']) else: current_block['type'] = self._classify_block(current_block, page_height) blocks.append(current_block) current_block = { 'text': word['text'], 'x0': word['x0'], 'top': word['top'], 'x1': word['x1'], 'bottom': word['bottom'] } current_block['type'] = self._classify_block(current_block, page_height) blocks.append(current_block) return blocks def _tesseract_data_to_blocks(self, data: Dict, page_width: float, page_height: float) -> List[Dict]: blocks = [] n_boxes = len(data['text']) current_line = [] for i in range(n_boxes): if int(data['conf'][i]) > 10 and data['text'][i].strip(): # We can group by line_num blocks.append({ 'text': data['text'][i], 'x': data['left'][i], 'y': data['top'][i], 'w': data['width'][i], 'h': data['height'][i], 'conf': float(data['conf'][i]), 'type': self._classify_block({'top': data['top'][i], 'text': data['text'][i]}, page_height) }) return blocks def _classify_block(self, block: Dict, page_height: float) -> str: y = block.get('top') or block.get('y') or 0 text = block.get('text', '').lower() if y < page_height * 0.15: return "HEADER" elif y > page_height * 0.85: return "FOOTER" if "total" in text: return "TOTAL" if "tax" in text or "gst" in text or "vat" in text: return "TAX" if re.search(r'\b(vendor|from)\b', text): return "VENDOR" if re.search(r'\b(bill to|sold to)\b', text): return "BILL_TO" if re.search(r'\b(ship to)\b', text): return "SHIP_TO" # Very naive fallback for table check if re.search(r'\b(qty|rate|amount|price|item)\b', text): return "TABLE_HEADER" return "TEXT"