Spaces:
Sleeping
Sleeping
| import os | |
| import json | |
| import shutil | |
| from typing import List | |
| from .schema import DocumentRecord, OCRWord, CLASS_FIELD_SCHEMAS | |
| def convert_inv_cdip(raw_base_dir: str, target_dir: str, max_samples: int = None) -> List[DocumentRecord]: | |
| """ | |
| Converts INV-CDIP labeled invoice dataset into DocumentRecord format. | |
| """ | |
| records = [] | |
| anno_dir = os.path.join(raw_base_dir, "annotation") | |
| images_dir = os.path.join(raw_base_dir, "data", "labeled") | |
| if not os.path.exists(anno_dir) or not os.path.exists(images_dir): | |
| return records | |
| os.makedirs(target_dir, exist_ok=True) | |
| anno_files = sorted([f for f in os.listdir(anno_dir) if f.endswith(".json")]) | |
| if max_samples: | |
| anno_files = anno_files[:max_samples] | |
| for f in anno_files: | |
| doc_id = os.path.splitext(f)[0] | |
| anno_path = os.path.join(anno_dir, f) | |
| img_src = os.path.join(images_dir, f"{doc_id}.png") | |
| if not os.path.exists(img_src): | |
| continue | |
| # Copy image to target_dir | |
| 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: | |
| anno_data = json.load(jf) | |
| # Map fields to unified schema | |
| fields = {k: None for k in CLASS_FIELD_SCHEMAS["invoice"]} | |
| ocr_words: List[OCRWord] = [] | |
| for fd in anno_data.get("Fields", []): | |
| val = fd.get("value", {}) | |
| key = fd.get("key", {}) | |
| lbl = val.get("label", "").lower() | |
| tag = val.get("tag", "") | |
| bbox_dict = val.get("bbox", {}) | |
| val_bbox = [ | |
| bbox_dict.get("xmin", 0), | |
| bbox_dict.get("ymin", 0), | |
| bbox_dict.get("xmax", 0), | |
| bbox_dict.get("ymax", 0) | |
| ] | |
| if "number" in lbl: | |
| fields["invoice_number"] = tag | |
| elif "date" in lbl: | |
| fields["invoice_date"] = tag | |
| elif "total" in lbl and "amount" in lbl: | |
| fields["total_amount"] = tag | |
| elif "order" in lbl: | |
| fields["buyer_name"] = tag | |
| # Add key words | |
| if key and key.get("tag"): | |
| k_bbox_dict = key.get("bbox", {}) | |
| k_box = [ | |
| k_bbox_dict.get("xmin", 0), | |
| k_bbox_dict.get("ymin", 0), | |
| k_bbox_dict.get("xmax", 0), | |
| k_bbox_dict.get("ymax", 0) | |
| ] | |
| for w in key["tag"].split(): | |
| ocr_words.append(OCRWord(text=w, bbox=k_box, conf=1.0)) | |
| # Add value words | |
| if tag: | |
| for w in tag.split(): | |
| ocr_words.append(OCRWord(text=w, bbox=val_bbox, conf=1.0)) | |
| rec = DocumentRecord( | |
| doc_id=doc_id, | |
| source_dataset="INV-CDIP", | |
| doc_class="invoice", | |
| image_path=dest_img_path, | |
| ocr_words=ocr_words, | |
| fields=fields | |
| ) | |
| records.append(rec) | |
| return records | |