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