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"
|
||||
Binary file not shown.
174
backend/engine/recognition/template_engine.py
Normal file
174
backend/engine/recognition/template_engine.py
Normal file
@@ -0,0 +1,174 @@
|
||||
import hashlib
|
||||
import json
|
||||
from typing import List, Dict, Any, Tuple
|
||||
from sqlalchemy.orm import Session
|
||||
from models.template_models import Template, TemplateDocument, DocumentLayout, TemplateRecognitionHistory
|
||||
from schemas.template_schemas import TemplateRecognitionResult
|
||||
|
||||
class TemplateRecognitionEngine:
|
||||
def __init__(self, db: Session):
|
||||
self.db = db
|
||||
|
||||
def generate_fingerprint(self, layouts: List[DocumentLayout]) -> str:
|
||||
"""
|
||||
Generate a structural fingerprint based on the document's layout.
|
||||
We use relative positions of HEADER, VENDOR, TABLE_HEADER blocks.
|
||||
"""
|
||||
key_blocks = []
|
||||
for layout in layouts:
|
||||
if layout.block_type in ["HEADER", "VENDOR", "TABLE_HEADER", "TOTAL"]:
|
||||
key_blocks.append({
|
||||
"type": layout.block_type,
|
||||
"text": layout.text_value[:50] if layout.text_value else "", # first 50 chars
|
||||
"rx": round(float(layout.x_coordinate) / 100) if layout.x_coordinate else 0, # relative bucket x
|
||||
"ry": round(float(layout.y_coordinate) / 100) if layout.y_coordinate else 0 # relative bucket y
|
||||
})
|
||||
|
||||
# Sort top to bottom, left to right
|
||||
key_blocks.sort(key=lambda b: (b['ry'], b['rx']))
|
||||
|
||||
fingerprint_data = json.dumps(key_blocks)
|
||||
return hashlib.sha256(fingerprint_data.encode('utf-8')).hexdigest()
|
||||
|
||||
def recognize_template(self, document_id: int) -> TemplateRecognitionResult:
|
||||
"""
|
||||
Matches a document against existing templates using a 4-level weighted score:
|
||||
Level 1: Vendor Match (40%)
|
||||
Level 2: Header Similarity (20%)
|
||||
Level 3: Layout/Fingerprint Similarity (20%)
|
||||
Level 4: Coordinate Similarity (20%)
|
||||
Threshold: 85%
|
||||
"""
|
||||
layouts = self.db.query(DocumentLayout).filter(DocumentLayout.fk_document_id == document_id).all()
|
||||
if not layouts:
|
||||
return TemplateRecognitionResult(templateMatched=False)
|
||||
|
||||
doc_fingerprint = self.generate_fingerprint(layouts)
|
||||
|
||||
# Extract features for scoring
|
||||
doc_vendor = self._extract_vendor_name(layouts)
|
||||
doc_headers = self._extract_headers(layouts)
|
||||
|
||||
templates = self.db.query(Template).filter(Template.active_flag == True).all()
|
||||
best_match = None
|
||||
highest_score = 0
|
||||
|
||||
for template in templates:
|
||||
score = 0.0
|
||||
|
||||
# Level 1: Fingerprint Exact Match (counts for Layout + Vendor + Header if exact)
|
||||
if template.template_fingerprint == doc_fingerprint:
|
||||
score += 100.0
|
||||
else:
|
||||
# Need to load template sample mapping to compare heuristics
|
||||
# In a real system, we'd compare against the mapped fields' coordinates
|
||||
score += self._calculate_heuristic_score(template, doc_vendor, doc_headers, layouts)
|
||||
|
||||
if score > highest_score:
|
||||
highest_score = score
|
||||
best_match = template
|
||||
|
||||
# Store history
|
||||
if best_match:
|
||||
history = TemplateRecognitionHistory(
|
||||
fk_template_id=best_match.pk_template_id,
|
||||
fk_document_id=document_id,
|
||||
recognition_score=highest_score,
|
||||
matched_flag=(highest_score >= 85.0)
|
||||
)
|
||||
self.db.add(history)
|
||||
self.db.commit()
|
||||
|
||||
if best_match and highest_score >= 85.0:
|
||||
extracted_fields = self._auto_map_fields(best_match, layouts)
|
||||
return TemplateRecognitionResult(
|
||||
templateMatched=True,
|
||||
templateId=best_match.pk_template_id,
|
||||
confidence=highest_score,
|
||||
extractedFields=extracted_fields
|
||||
)
|
||||
|
||||
return TemplateRecognitionResult(templateMatched=False)
|
||||
|
||||
def _extract_vendor_name(self, layouts: List[DocumentLayout]) -> str:
|
||||
for layout in layouts:
|
||||
if layout.block_type == "VENDOR":
|
||||
return layout.text_value.lower()
|
||||
return ""
|
||||
|
||||
def _extract_headers(self, layouts: List[DocumentLayout]) -> List[str]:
|
||||
headers = []
|
||||
for layout in layouts:
|
||||
if layout.block_type == "TABLE_HEADER" and layout.text_value:
|
||||
headers.append(layout.text_value.lower())
|
||||
return headers
|
||||
|
||||
def _calculate_heuristic_score(self, template: Template, doc_vendor: str, doc_headers: List[str], layouts: List[DocumentLayout]) -> float:
|
||||
score = 0.0
|
||||
|
||||
# This requires the template to have some stored metadata or we check its fields
|
||||
# E.g., if template has a field "Vendor Name" mapped to a specific text
|
||||
|
||||
# 1. Vendor Match (40%)
|
||||
# For this mockup, we check if the template name matches the vendor
|
||||
if template.template_name.lower() in doc_vendor or doc_vendor in template.template_name.lower():
|
||||
score += 40.0
|
||||
|
||||
# 2. Header Similarity (20%)
|
||||
# Check if template fields exist that match doc_headers
|
||||
field_labels = [f.field_label.lower() for f in template.fields]
|
||||
header_matches = sum(1 for h in doc_headers if any(h in fl or fl in h for fl in field_labels))
|
||||
if len(doc_headers) > 0:
|
||||
score += (header_matches / len(doc_headers)) * 20.0
|
||||
|
||||
# 3. Layout / Coordinate Similarity
|
||||
# We can check the template's previous mappings coordinates against current document
|
||||
# If coordinates are within a tolerance, we add score
|
||||
mapping_count = 0
|
||||
match_count = 0
|
||||
for mapping in template.mappings:
|
||||
mapping_count += 1
|
||||
# Find a layout block in current document that is near mapping coordinates
|
||||
for layout in layouts:
|
||||
if layout.page_no == mapping.page_no:
|
||||
# check distance
|
||||
if mapping.x_coordinate and mapping.y_coordinate and layout.x_coordinate and layout.y_coordinate:
|
||||
dx = abs(float(mapping.x_coordinate) - float(layout.x_coordinate))
|
||||
dy = abs(float(mapping.y_coordinate) - float(layout.y_coordinate))
|
||||
if dx < 50 and dy < 20: # arbitrary tolerance
|
||||
match_count += 1
|
||||
break
|
||||
|
||||
if mapping_count > 0:
|
||||
coord_score = (match_count / mapping_count) * 40.0 # Layout 20% + Coord 20%
|
||||
score += coord_score
|
||||
|
||||
return score
|
||||
|
||||
def _auto_map_fields(self, template: Template, layouts: List[DocumentLayout]) -> List[dict]:
|
||||
extracted = []
|
||||
for mapping in template.mappings:
|
||||
# Find closest block in new document
|
||||
best_block = None
|
||||
min_dist = float('inf')
|
||||
|
||||
for layout in layouts:
|
||||
if layout.page_no == mapping.page_no:
|
||||
if mapping.x_coordinate and mapping.y_coordinate and layout.x_coordinate and layout.y_coordinate:
|
||||
dx = abs(float(mapping.x_coordinate) - float(layout.x_coordinate))
|
||||
dy = abs(float(mapping.y_coordinate) - float(layout.y_coordinate))
|
||||
dist = dx**2 + dy**2
|
||||
if dist < min_dist and dist < 5000: # tolerance squared
|
||||
min_dist = dist
|
||||
best_block = layout
|
||||
|
||||
if best_block:
|
||||
extracted.append({
|
||||
"field_id": mapping.fk_template_field_id,
|
||||
"field_label": mapping.field.field_label if mapping.field else "",
|
||||
"value": best_block.text_value,
|
||||
"confidence": 95.0, # arbitrary high confidence for coordinate match
|
||||
"layout_id": best_block.pk_document_data_id
|
||||
})
|
||||
|
||||
return extracted
|
||||
Reference in New Issue
Block a user