Spaces:
Sleeping
Sleeping
File size: 3,138 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 | 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
|