Spaces:
Sleeping
Sleeping
File size: 3,924 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 | import os
import json
import shutil
from typing import List
from .schema import DocumentRecord, OCRWord, CLASS_FIELD_SCHEMAS
def convert_funsd(funsd_base_dir: str, target_dir: str, max_samples: int = None) -> List[DocumentRecord]:
"""
Converts FUNSD form dataset into DocumentRecord format.
"""
records = []
os.makedirs(target_dir, exist_ok=True)
# Search in training_data and testing_data
subdirs = ["training_data", "testing_data"]
total_loaded = 0
for sdir in subdirs:
anno_dir = os.path.join(funsd_base_dir, sdir, "annotations")
img_dir = os.path.join(funsd_base_dir, sdir, "images")
if not os.path.exists(anno_dir) or not os.path.exists(img_dir):
continue
for fname in sorted(os.listdir(anno_dir)):
if not fname.endswith(".json"):
continue
if max_samples and total_loaded >= max_samples:
break
doc_id = f"funsd_{os.path.splitext(fname)[0]}"
anno_path = os.path.join(anno_dir, fname)
img_src = os.path.join(img_dir, os.path.splitext(fname)[0] + ".png")
if not os.path.exists(img_src):
continue
dest_img_path = os.path.join(target_dir, f"{doc_id}.png")
if not os.path.exists(dest_img_path):
shutil.copy2(img_src, dest_img_path)
with open(anno_path, "r", encoding="utf-8") as jf:
fdata = json.load(jf)
fields = {k: None for k in CLASS_FIELD_SCHEMAS["application_form"]}
ocr_words: List[OCRWord] = []
qa_pairs = []
# Extract words and detect key-value associations
items = fdata.get("form", [])
# Map item id to item
id_to_item = {item.get("id"): item for item in items}
for item in items:
label = item.get("label", "")
text = item.get("text", "").strip()
words = item.get("words", [])
for w in words:
w_text = w.get("text", "")
w_box = w.get("box", [0, 0, 0, 0])
if w_text:
ocr_words.append(OCRWord(text=w_text, bbox=w_box, conf=1.0))
# Check question-answer linking
if label == "question":
linked_ids = []
for link in item.get("linking", []):
if len(link) == 2 and link[0] == item.get("id"):
linked_ids.append(link[1])
answer_texts = [id_to_item[aid].get("text", "") for aid in linked_ids if aid in id_to_item]
ans_str = " ".join(answer_texts).strip() if answer_texts else None
q_lower = text.lower()
if ans_str:
qa_pairs.append(f"{text}: {ans_str}")
if any(k in q_lower for k in ["name", "applicant", "employee", "signature", "person"]) and not fields["applicant_name"]:
fields["applicant_name"] = ans_str
elif "date" in q_lower and not fields["date"]:
fields["date"] = ans_str
elif any(k in q_lower for k in ["address", "street", "city", "state", "location"]) and not fields["address"]:
fields["address"] = ans_str
fields["other_key_value_pairs"] = len(qa_pairs)
rec = DocumentRecord(
doc_id=doc_id,
source_dataset="FUNSD",
doc_class="application_form",
image_path=dest_img_path,
ocr_words=ocr_words,
fields=fields
)
records.append(rec)
total_loaded += 1
if max_samples and total_loaded >= max_samples:
break
return records
|