Text Extraction Done, UI Fixes Done, Extracted Template Layout Saved in DB
This commit is contained in:
Binary file not shown.
217
backend/engine/ocr/document_processor.py
Normal file
217
backend/engine/ocr/document_processor.py
Normal file
@@ -0,0 +1,217 @@
|
||||
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"
|
||||
Reference in New Issue
Block a user