File size: 3,334 Bytes
29e0671 | 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 116 117 118 119 120 121 122 123 124 | import json
import os
import torch
from PIL import Image
from tqdm import tqdm
DEFAULT_SAVE_DIR = "./data/processed"
def clean_price(price_str: str) -> str:
return str(price_str).strip().rstrip('.')
def parse_ground_truth(gt_string: str) -> dict:
gt = json.loads(gt_string)
parsed = gt.get("gt_parse", {})
items = []
menu = parsed.get("menu", [])
if not isinstance(menu, list):
menu = []
for item in menu:
if not isinstance(item, dict):
continue
items.append({
"name" : str(item.get("nm", "")).strip(),
"count": str(item.get("cnt", "")).strip(),
"price": clean_price(item.get("price", "0")),
})
sub = parsed.get("sub_total", {})
if not isinstance(sub, dict):
sub = {}
total_obj = parsed.get("total", {})
if not isinstance(total_obj, dict):
total_obj = {}
return {
"vendor" : "unknown",
"items" : items,
"subtotal": clean_price(sub.get("subtotal_price", "0")),
"tax" : clean_price(sub.get("tax_price", "0")),
"service" : clean_price(sub.get("service_price", "0")),
"total" : clean_price(total_obj.get("total_price", "0")),
}
def format_instruction(sample: dict) -> dict:
img = sample["image"].convert("RGB")
parsed = parse_ground_truth(sample["ground_truth"])
prompt = (
"Extract the receipt information and return ONLY valid JSON.\n"
'Schema: {"vendor":"","items":[{"name":"","count":"","price":""}],"subtotal":"","tax":"","service":"","total":""}'
)
target = json.dumps(parsed, ensure_ascii=False)
return {
"image" : img,
"prompt": prompt,
"target": target,
}
def save_splits(splits: dict, save_dir: str = DEFAULT_SAVE_DIR) -> None:
os.makedirs(save_dir, exist_ok=True)
for split, data in splits.items():
path = os.path.join(save_dir, f"{split}.pt")
torch.save(data, path)
print(f" Saved {len(data):>5} samples → {path}")
def load_splits(save_dir: str = DEFAULT_SAVE_DIR) -> dict:
splits = {}
for split in ["train", "validation", "test"]:
path = os.path.join(save_dir, f"{split}.pt")
if not os.path.exists(path):
raise FileNotFoundError(
f"Expected '{path}' — run the pipeline first to generate it."
)
splits[split] = torch.load(path)
print(f" Loaded {len(splits[split]):>5} samples ← {path}")
return splits
def preprocess_dataset(dataset: dict, save_dir: str = None) -> dict:
splits = {}
for split in ["train", "validation", "test"]:
processed = []
skipped = 0
for sample in tqdm(dataset[split], desc=f" {split:<12}", unit="sample"):
try:
formatted = format_instruction(sample)
processed.append(formatted)
except Exception as e:
skipped += 1
tqdm.write(f" [skip] {split} sample: {e}")
continue
if skipped:
print(f" {split}: skipped {skipped} samples")
splits[split] = processed
if save_dir is not None:
print("\nSaving processed splits...")
save_splits(splits, save_dir)
return splits |