File size: 4,185 Bytes
44c2f50 6267e20 44c2f50 d95323c 44c2f50 | 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 109 110 111 112 113 114 115 | """SROIE dataset loader + normalizer.
SROIE (ICDAR 2019 Robust Reading Challenge on Scanned Receipts Information
Extraction) has ~1000 scanned Singapore-region receipts with ground truth for
four fields: company, address, date, total.
Datasets change over time — we default to `darentang/sroie` on Hugging Face,
which mirrors the ICDAR test set. If that ID has moved by the time you run
this, pass `--dataset-id` to override.
Reference: https://rrc.cvc.uab.es/?ch=13
"""
from __future__ import annotations
from collections.abc import Iterator
from typing import Any
from src.data_prep.parsers import clean_text, parse_date, parse_money
from src.schemas import Address, Receipt
from src.utils.logging import logger
# Default Hugging Face dataset ID + fallbacks. Updated 2026-07 — override
# via the CLI if it has moved again.
DEFAULT_DATASET_IDS = (
"darentang/sroie",
"mychen76/invoices-and-receipts_ocr_v1",
)
def load_sroie_split(split: str = "test", dataset_id: str | None = None):
"""Load a SROIE split from Hugging Face. Requires `datasets` package.
Returns a HuggingFace Dataset object — iterable of dicts.
"""
from datasets import load_dataset
ids_to_try = (dataset_id,) if dataset_id else DEFAULT_DATASET_IDS
last_err: Exception | None = None
for ds_id in ids_to_try:
try:
logger.info(f"Loading SROIE split={split!r} from Hugging Face id={ds_id!r}")
return load_dataset(ds_id, split=split)
except Exception as e:
logger.warning(f"Failed to load {ds_id}: {e}")
last_err = e
raise RuntimeError(
f"Could not load any SROIE dataset. Last error: {last_err}. "
f"Try passing --dataset-id explicitly. Check huggingface.co for a working ID."
) from last_err
def normalize_sroie_record(record: dict[str, Any]) -> Receipt | None:
"""Convert a raw SROIE record into our Receipt schema.
Handles both common shapes:
A) flat: {"company": ..., "date": ..., "address": ..., "total": ...}
B) nested: {"parsed_data": {"company": ..., ...}}
Returns None if we can't extract the minimum required fields (merchant + total).
"""
# Unwrap common nesting patterns. HF datasets sometimes store the annotation
# as a JSON *string* (parquet doesn\'t love nested dicts), so try to parse
# string values before treating them as opaque.
import json as _json
_raw = record.get("parsed_data") or record.get("ground_truth") or record
if isinstance(_raw, str):
try:
_raw = _json.loads(_raw)
except Exception:
return None # unparseable — skip this record rather than crash
if not isinstance(_raw, dict):
return None
src = _raw
company = clean_text(src.get("company") or src.get("merchant"))
address = clean_text(src.get("address"))
date_str = src.get("date")
total_str = src.get("total") or src.get("amount")
if not company:
logger.debug(f"SROIE record missing company field; skipping. record keys={list(src.keys())}")
return None
total = parse_money(total_str)
if total is None:
logger.debug(f"SROIE record {company!r} missing parseable total ({total_str!r}); skipping.")
return None
try:
return Receipt(
merchant=company,
merchant_address=Address(line1=address) if address else None,
transaction_date=parse_date(date_str),
total=total,
# SROIE is largely Singapore/Malaysia — SGD is the safer default than USD.
currency="SGD",
)
except Exception as e:
logger.warning(f"Failed to build Receipt from SROIE record {company!r}: {e}")
return None
def iter_normalized(dataset) -> Iterator[tuple[str, Receipt]]:
"""Iterate a HF Dataset, yielding (record_id, Receipt) pairs.
Skips records that fail normalization.
"""
for idx, rec in enumerate(dataset):
record_id = str(rec.get("id") or rec.get("image_id") or f"sroie_{idx:05d}")
normalized = normalize_sroie_record(rec)
if normalized is not None:
yield record_id, normalized
|