Spaces:
Sleeping
Sleeping
| 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 | |