Spaces:
Sleeping
Sleeping
File size: 4,206 Bytes
c1b1f3a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 | import os
import re
import numpy as np
import cv2
from typing import List, Dict, Any
from sklearn.base import BaseEstimator, TransformerMixin
from sklearn.feature_extraction.text import TfidfVectorizer
CLASS_INDICATIVE_KEYWORDS = [
"invoice", "receipt", "application", "applicant", "bill", "total", "form",
"fatura", "tax", "cashier", "subtotal", "change", "signature"
]
class LayoutFeatureExtractor(BaseEstimator, TransformerMixin):
"""
Extracts layout and statistical features from OCR tokens and document records:
- token count
- digit ratio
- keyword presence indicators
- aspect ratio
- line count estimate
"""
def __init__(self):
self.invoice_kw = ["invoice", "tax", "po", "bill", "due", "balance", "vendor"]
self.receipt_kw = ["receipt", "change", "cash", "cashier", "subtotal", "card", "store"]
self.form_kw = ["applicant", "signature", "employee", "title", "department", "date", "form"]
def fit(self, X, y=None):
return self
def transform(self, records: List[Any]) -> np.ndarray:
feats = []
for r in records:
words = r.ocr_words if hasattr(r, "ocr_words") else r.get("ocr_words", [])
text_tokens = [w.text if hasattr(w, "text") else w.get("text", "") for w in words]
full_text = " ".join(text_tokens).lower()
num_tokens = len(text_tokens)
total_chars = max(len(full_text), 1)
num_digits = sum(c.isdigit() for c in full_text)
digit_ratio = num_digits / total_chars
has_inv = float(any(kw in full_text for kw in self.invoice_kw))
has_rec = float(any(kw in full_text for kw in self.receipt_kw))
has_frm = float(any(kw in full_text for kw in self.form_kw))
# Bounding box layout estimation
y_coords = [w.bbox[1] if hasattr(w, "bbox") else w.get("bbox", [0, 0, 0, 0])[1] for w in words]
unique_lines = len(set([y // 20 for y in y_coords])) if y_coords else 0
# Aspect ratio from image if available
aspect_ratio = 1.4 # standard portrait default
img_path = r.image_path if hasattr(r, "image_path") else r.get("image_path", "")
if img_path and os.path.exists(img_path):
try:
from PIL import Image
with Image.open(img_path) as im:
w_img, h_img = im.size
aspect_ratio = float(w_img / max(h_img, 1))
except Exception:
pass
feats.append([
float(num_tokens),
float(digit_ratio),
has_inv,
has_rec,
has_frm,
float(unique_lines),
float(aspect_ratio)
])
return np.array(feats, dtype=np.float32)
def mask_class_keywords(text: str) -> str:
"""Masks obvious class-indicative keywords for robustness ablation."""
pattern = re.compile(r'\b(' + '|'.join(CLASS_INDICATIVE_KEYWORDS) + r')\b', re.IGNORECASE)
return pattern.sub("[MASKED]", text)
def apply_image_augmentations(img: np.ndarray, aug_type: str = "all") -> np.ndarray:
"""
Applies image distortions for robustness testing:
- blur (Gaussian blur)
- rotation (+/- 5 degrees)
- compression (JPEG encoding at low quality)
- noise (Gaussian random noise)
"""
out = img.copy()
h, w = out.shape[:2]
if aug_type in ["rotation", "all"]:
angle = float(np.random.choice([-5.0, 5.0]))
M = cv2.getRotationMatrix2D((w // 2, h // 2), angle, 1.0)
out = cv2.warpAffine(out, M, (w, h), borderMode=cv2.BORDER_REPLICATE)
if aug_type in ["blur", "all"]:
out = cv2.GaussianBlur(out, (5, 5), 0)
if aug_type in ["compression", "all"]:
encode_param = [int(cv2.IMWRITE_JPEG_QUALITY), 30]
_, enc = cv2.imencode('.jpg', out, encode_param)
out = cv2.imdecode(enc, 1)
if aug_type in ["noise", "all"]:
noise = np.random.normal(0, 15, out.shape).astype(np.float32)
out = np.clip(out.astype(np.float32) + noise, 0, 255).astype(np.uint8)
return out
|