218 lines
8.8 KiB
Python
218 lines
8.8 KiB
Python
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"
|