import base64 import json import re import threading import unicodedata from collections import Counter, defaultdict from pathlib import Path from statistics import median from typing import Optional, List, Dict, Any from openai import OpenAI def build_json_prompt(seed_tag: str) -> str: return ( f"You are a precise chart data extractor. {seed_tag}\n" "Extract the underlying data as STRICT, MINIFIED JSON with this exact schema:\n" "{" "\"x_axis_label\": string, " "\"x\": [string, ...], " "\"series\": [ {\"name\": string, \"legend_index\": int, \"values\": [number|null, ...]} , ... ]" "}\n" "Rules:\n" "- Read the exact x-axis label text and use it as x_axis_label.\n" "- The x array must list x-axis tick labels in the exact visual order.\n" "- Include one series object per legend entry; use legend text as name.\n" "- For each series, legend_index must be the 0-based order of its legend entry as visually displayed (top/left to bottom/right).\n" "- For each series.values, align 1:1 with x by position; use null if missing.\n" "- Return ONLY compact JSON (no comments, no code fences, no extra text).\n" "- Use numbers (not strings) for numeric values; do not append units.\n" "- If chart has no data for a slot, use null.\n" ) def build_fallback_csv_prompt() -> str: return ( "You are a precise data extraction model. Given a chart image, extract the underlying data table.\n" "Return ONLY the CSV text with a header row and no markdown code fences.\n" "Rules:\n" "- The first column must be the x-axis values with its exact axis label as the header.\n" "- Include one column per data series using the legend labels as headers.\n" "- Preserve the original order of x-axis ticks as they appear.\n" "- Use plain CSV (comma-separated), no explanations, no extra text.\n" ) def image_to_data_uri(image_path: Path) -> str: mime = "image/png" if image_path.suffix.lower() == ".png" else "image/jpeg" data = image_path.read_bytes() b64 = base64.b64encode(data).decode("ascii") return f"data:{mime};base64,{b64}" def clean_to_csv(text: str) -> str: return text.strip() def _first_json_obj(s: str) -> Optional[str]: s = s.strip() fence = re.compile(r"^```(?:json)?\s*(.*?)\s*```$", re.S) m = fence.match(s) if m: s = m.group(1).strip() try: json.loads(s) return s except Exception: pass start = s.find("{") end = s.rfind("}") if start != -1 and end != -1 and end > start: candidate = s[start : end + 1] try: json.loads(candidate) return candidate except Exception: return None return None def _format_number(x: Optional[float]) -> str: if x is None: return "" try: v = float(x) if abs(v - round(v)) < 1e-9: return str(int(round(v))) s = f"{v:.6g}" return s except Exception: return "" def _norm_text(s: str) -> str: if not isinstance(s, str): s = str(s) s = unicodedata.normalize("NFKC", s) s = s.replace("\u2013", "-").replace("\u2014", "-") # en/em dash to hyphen s = s.replace("\u2212", "-") # minus sign to hyphen s = s.strip() s = re.sub(r"\s+", " ", s) s = s.lower() s = re.sub(r"^[\s\.,;:\|/\\]+", "", s) s = re.sub(r"[\s\.,;:\|/\\]+$", "", s) return s def _is_close(a: float, b: float) -> bool: try: a = float(a) b = float(b) except Exception: return False if a == b: return True rel = 0.05 * max(abs(a), abs(b), 1e-9) return abs(a - b) <= max(rel, 1e-6) class VLMExtractor: """Adaptive-ensemble JSON extraction -> legend-index-aware aggregation with normalized x consensus -> targeted repair -> CSV.""" def __init__(self, model: str = "gpt-5-mini", client: Optional[OpenAI] = None, ensemble_runs: int = 3) -> None: self.model = model self.client = client or OpenAI() self.ensemble_runs = max(1, ensemble_runs) self.total_prompt_tokens: int = 0 self.total_completion_tokens: int = 0 self.total_cost_usd: float = 0.0 self.num_queries: int = 0 self._usage_lock = threading.Lock() def _pricing_for_model(self) -> dict: name = (self.model or "").lower() per_million = { "gpt-5": {"in": 1.250, "in_cached": 0.125, "out": 10.000}, "gpt-5-mini": {"in": 0.250, "in_cached": 0.025, "out": 2.000}, "gpt-5-nano": {"in": 0.050, "in_cached": 0.005, "out": 0.400}, } if name.startswith("gpt-5-nano"): chosen = per_million["gpt-5-nano"] elif name.startswith("gpt-5-mini"): chosen = per_million["gpt-5-mini"] elif name.startswith("gpt-5"): chosen = per_million["gpt-5"] else: chosen = per_million["gpt-5-mini"] return {k: v / 1_000_000.0 for k, v in chosen.items()} def _accumulate_usage(self, response: Any) -> None: usage = getattr(response, "usage", None) with self._usage_lock: if usage is not None: prompt_tokens = int( (getattr(usage, "prompt_tokens", None) or getattr(usage, "input_tokens", 0)) or 0 ) completion_tokens = int( (getattr(usage, "completion_tokens", None) or getattr(usage, "output_tokens", 0)) or 0 ) details = getattr(usage, "prompt_tokens_details", None) cached_tokens = 0 if details is not None: cached_tokens = int(getattr(details, "cached_tokens", 0) or 0) else: cache_read = getattr(usage, "cache_read_input_tokens", 0) or 0 cached_tokens = int(cache_read) non_cached_prompt_tokens = max(0, prompt_tokens - cached_tokens) rates = self._pricing_for_model() cost = ( non_cached_prompt_tokens * rates["in"] + cached_tokens * rates["in_cached"] + completion_tokens * rates["out"] ) self.total_prompt_tokens += prompt_tokens self.total_completion_tokens += completion_tokens self.total_cost_usd += cost self.num_queries += 1 else: self.num_queries += 1 def _extract_json_once(self, image_uri: str, seed_tag: str) -> Optional[Dict[str, Any]]: prompt = build_json_prompt(seed_tag) response = self.client.chat.completions.create( model=self.model, messages=[ { "role": "user", "content": [ {"type": "text", "text": prompt}, {"type": "image_url", "image_url": {"url": image_uri}}, ], } ], ) self._accumulate_usage(response) text = (response.choices[0].message.content or "").strip() js = _first_json_obj(text) if not js: return None try: data = json.loads(js) except Exception: return None if not isinstance(data, dict): return None if "x_axis_label" not in data or "x" not in data or "series" not in data: return None if not isinstance(data["x"], list) or not isinstance(data["series"], list): return None x = data["x"] valid_series = [] for idx, s in enumerate(data["series"]): if not isinstance(s, dict): continue name = s.get("name") vals = s.get("values") li = s.get("legend_index", None) li_int = None if isinstance(li, int): li_int = li else: # If missing or not int, leave as None; do NOT infer from JSON order to avoid injecting noise. li_int = None if not isinstance(name, str) or not isinstance(vals, list): continue cleaned = [] for v in vals: if v is None: cleaned.append(None) else: try: cleaned.append(float(v)) except Exception: cleaned.append(None) if len(cleaned) == len(x): valid_series.append({"name": name.strip(), "legend_index": li_int, "values": cleaned}) if not valid_series: return None return {"x_axis_label": str(data["x_axis_label"]).strip(), "x": [str(t) for t in x], "series": valid_series} def _cluster_consensus(self, vals: List[float]) -> float: """Return median of the largest tolerance-based cluster; fallback to global median.""" if not vals: return None if len(vals) == 1: return float(vals[0]) best_support = 1 best_indices = None for i, v in enumerate(vals): indices = [j for j, u in enumerate(vals) if _is_close(u, v)] support = len(indices) if support > best_support or (support == best_support and best_indices is None): best_support = support best_indices = indices if best_support >= 2 and best_indices is not None: cluster_vals = [float(vals[j]) for j in best_indices] return float(median(cluster_vals)) return float(median([float(v) for v in vals])) def _aggregate_candidates(self, cands: List[Dict[str, Any]]) -> Optional[Dict[str, Any]]: if not cands: return None norm_x_lists = [tuple(_norm_text(t) for t in c["x"]) for c in cands] norm_counter = Counter(norm_x_lists) if norm_counter: canonical_norm_x = list(norm_counter.most_common(1)[0][0]) aligned = [c for c, nx in zip(cands, norm_x_lists) if list(nx) == canonical_norm_x] else: canonical_norm_x = [] aligned = [] if not aligned: best_idx = max(range(len(cands)), key=lambda i: len(norm_x_lists[i])) canonical_norm_x = list(norm_x_lists[best_idx]) aligned = [cands[best_idx]] L = len(canonical_norm_x) pos_token_counters: List[Counter] = [Counter() for _ in range(L)] for c in aligned: raw_x = c["x"] if len(raw_x) != L: continue for i in range(L): if _norm_text(raw_x[i]) == canonical_norm_x[i]: pos_token_counters[i][raw_x[i]] += 1 canonical_x_raw = [] for i in range(L): if pos_token_counters[i]: canonical_x_raw.append(pos_token_counters[i].most_common(1)[0][0]) else: canonical_x_raw.append(canonical_norm_x[i]) norm_label_counter = Counter([_norm_text(c["x_axis_label"]) for c in aligned]) canonical_norm_label = norm_label_counter.most_common(1)[0][0] raw_label_counter = Counter([c["x_axis_label"] for c in aligned if _norm_text(c["x_axis_label"]) == canonical_norm_label]) canonical_label = raw_label_counter.most_common(1)[0][0] if raw_label_counter else canonical_norm_label # Legend-index-aware aggregation # Collect index presence across aligned candidates idx_counts = Counter() per_cand_idx_map = [] for c in aligned: idx_map = {} for s in c["series"]: li = s.get("legend_index", None) if isinstance(li, int): idx_map[li] = s idx_counts[li] += 1 per_cand_idx_map.append(idx_map) use_index_alignment = any(v >= 2 for v in idx_counts.values()) and len(idx_counts) >= 1 series_values_by_idx: Dict[int, List[Optional[float]]] = {} name_counter_by_idx: Dict[int, Counter] = {} if use_index_alignment: canonical_indices = sorted(idx_counts.keys()) for li in canonical_indices: # Majority name for this legend index name_ctr = Counter() for idx_map in per_cand_idx_map: s = idx_map.get(li) if s and isinstance(s.get("name", None), str): name_ctr[s["name"].strip()] += 1 if not name_ctr: continue name_counter_by_idx[li] = name_ctr # Aggregate values position-wise col_vals = [None] * L for i in range(L): bucket = [] for idx_map in per_cand_idx_map: s = idx_map.get(li) if s: vlist = s.get("values") or [] v = vlist[i] if i < len(vlist) else None if isinstance(v, (int, float)): bucket.append(float(v)) if bucket: col_vals[i] = self._cluster_consensus(bucket) else: col_vals[i] = None series_values_by_idx[li] = col_vals # Name-based fallback for series without reliable legend_index name_map_counts = defaultdict(Counter) for c in aligned: for s in c["series"]: norm = _norm_text(s.get("name", "")) if not norm: continue name_map_counts[norm][s["name"].strip()] += 1 canonical_names = {norm: casings.most_common(1)[0][0] for norm, casings in name_map_counts.items()} # Exclude names already captured via index-based alignment to avoid duplicates used_name_norms = set() if use_index_alignment: for li, ctr in name_counter_by_idx.items(): chosen = ctr.most_common(1)[0][0] used_name_norms.add(_norm_text(chosen)) series_values_by_name: Dict[str, List[Optional[float]]] = {} for i in range(L): # Precompute per position aggregation for names pass # Aggregate by name for remaining for norm, display in canonical_names.items(): if norm in used_name_norms: continue col_vals = [None] * L for i in range(L): bucket = [] for c in aligned: for s in c.get("series", []): if _norm_text(s.get("name", "")) == norm: vlist = s.get("values") or [] v = vlist[i] if i < len(vlist) else None if isinstance(v, (int, float)): bucket.append(float(v)) if bucket: col_vals[i] = self._cluster_consensus(bucket) else: col_vals[i] = None series_values_by_name[norm] = col_vals # Decide series order: indices first (ascending), then remaining by majority order across candidates out_series = [] if use_index_alignment: for li in sorted(series_values_by_idx.keys()): name_ctr = name_counter_by_idx.get(li, Counter()) if not name_ctr: continue display_name = name_ctr.most_common(1)[0][0] out_series.append({"name": display_name, "values": series_values_by_idx[li]}) # Determine name-based order using observed orders orders = [] for c in aligned: order = tuple([_norm_text(s.get("name", "")) for s in c.get("series", []) if _norm_text(s.get("name", "")) in series_values_by_name]) if order: orders.append(order) order_counter = Counter(orders) if orders: series_order_norm = list(order_counter.most_common(1)[0][0]) else: series_order_norm = sorted(series_values_by_name.keys()) for norm in series_order_norm: out_series.append({"name": canonical_names.get(norm, norm), "values": series_values_by_name[norm]}) if not out_series: # Fallback: if index alignment failed and no names aggregated, return None return None return {"x_axis_label": canonical_label, "x": canonical_x_raw, "series": out_series} def _to_csv(self, data: Dict[str, Any]) -> str: x_label = data["x_axis_label"] x = data["x"] series = data["series"] headers = [x_label] + [s["name"] for s in series] rows = [] for i, xv in enumerate(x): row = [str(xv)] for s in series: row.append(_format_number(s["values"][i] if i < len(s["values"]) else None)) rows.append(",".join(row)) csv_text = ",".join(headers) + "\n" + "\n".join(rows) return clean_to_csv(csv_text) def _fallback_direct_csv(self, image_uri: str) -> str: prompt = build_fallback_csv_prompt() response = self.client.chat.completions.create( model=self.model, messages=[ { "role": "user", "content": [ {"type": "text", "text": prompt}, {"type": "image_url", "image_url": {"url": image_uri}}, ], } ], ) self._accumulate_usage(response) text = response.choices[0].message.content or "" return clean_to_csv(text) def _repair_missing_values(self, image_uri: str, data: Dict[str, Any], max_cells: int = 12) -> Optional[Dict[str, Any]]: holes = [] x = data.get("x", []) for s in data.get("series", []): name = s.get("name", "") vals = s.get("values", []) for i, v in enumerate(vals): if v is None: holes.append({"i": i, "series": name, "x": x[i] if i < len(x) else ""}) if len(holes) >= max_cells: break if len(holes) >= max_cells: break if not holes: return None holes_json = json.dumps(holes, separators=(",", ":")) prompt = ( "You are a precise chart data extractor. Read ONLY the requested cells from the chart image.\n" "For each item in holes, look at the series name and x tick label, read the numeric value from the chart " "(prefer exact data labels if present; otherwise estimate from the axis scale/gridlines), and return STRICT, MINIFIED JSON:\n" "{\"cells\":[{\"i\":int,\"series\":string,\"value\":number|null}, ...]}\n" "Rules:\n" "- Use the same 0-based index i and exact series string provided in holes.\n" "- Return null only if the chart truly lacks a value for that cell.\n" "- Return ONLY compact JSON (no extra text, no code fences).\n" f"holes:{holes_json}\n" ) response = self.client.chat.completions.create( model=self.model, messages=[ { "role": "user", "content": [ {"type": "text", "text": prompt}, {"type": "image_url", "image_url": {"url": image_uri}}, ], } ], ) self._accumulate_usage(response) text = (response.choices[0].message.content or "").strip() js = _first_json_obj(text) if not js: return None try: out = json.loads(js) except Exception: return None if not isinstance(out, dict) or "cells" not in out or not isinstance(out["cells"], list): return None series_by_name = {s["name"]: s for s in data.get("series", [])} changed = False for cell in out["cells"]: try: i = int(cell.get("i")) sname = str(cell.get("series")) val = cell.get("value", None) except Exception: continue if sname in series_by_name and 0 <= i < len(series_by_name[sname]["values"]): if val is None: continue try: fv = float(val) except Exception: continue series_by_name[sname]["values"][i] = fv changed = True return data if changed else None def _agreement_score(self, cands: List[Dict[str, Any]]) -> float: if len(cands) < 2: return 0.0 norm_x_lists = [tuple(_norm_text(t) for t in c["x"]) for c in cands] cnt = Counter(norm_x_lists) if not cnt: return 0.0 canonical_norm_x = list(cnt.most_common(1)[0][0]) aligned_idx = [i for i, nx in enumerate(norm_x_lists) if list(nx) == canonical_norm_x] if len(aligned_idx) < 2: return 0.0 L = len(canonical_norm_x) per_cand_series = [] for i in aligned_idx: c = cands[i] smap = {} for s in c["series"]: smap[_norm_text(s["name"])] = list(s["values"][:L]) per_cand_series.append(smap) series_norms = Counter() for smap in per_cand_series: for k in smap.keys(): series_norms[k] += 1 valid_series = [k for k, v in series_norms.items() if v >= 2] if not valid_series or L == 0: return 0.0 total = 0 agreed = 0 for i in range(L): for sn in valid_series: vals = [] for smap in per_cand_series: if sn in smap: v = smap[sn][i] if i < len(smap[sn]) else None if isinstance(v, (int, float)): vals.append(float(v)) if len(vals) < 2: continue total += 1 best = 1 for v in vals: count = sum(1 for u in vals if _is_close(u, v)) if count > best: best = count if best >= 2: agreed += 1 if total == 0: return 0.0 return agreed / total def _collect_disagreement_targets(self, cands: List[Dict[str, Any]], agg: Dict[str, Any], max_cells: int = 8) -> List[Dict[str, Any]]: if not cands or not agg: return [] norm_x_lists = [tuple(_norm_text(t) for t in c["x"]) for c in cands] cnt = Counter(norm_x_lists) if not cnt: return [] canonical_norm_x = list(cnt.most_common(1)[0][0]) agg_norm_x = [_norm_text(t) for t in agg.get("x", [])] if list(agg_norm_x) != canonical_norm_x: return [] aligned = [c for c, nx in zip(cands, norm_x_lists) if list(nx) == canonical_norm_x] if len(aligned) < 2: return [] L = len(canonical_norm_x) per_cand_series = [] for c in aligned: smap = {} for s in c.get("series", []): smap[_norm_text(s.get("name", ""))] = list((s.get("values") or [])[:L]) per_cand_series.append(smap) agg_series = agg.get("series", []) norm_to_raw = {_norm_text(s["name"]): s["name"] for s in agg_series} targets = [] for snorm, raw_name in norm_to_raw.items(): for i in range(L): vals = [] for smap in per_cand_series: if snorm in smap: v = smap[snorm][i] if i < len(smap[snorm]) else None if isinstance(v, (int, float)): vals.append(float(v)) if len(vals) >= 2: best_cluster = 1 for v in vals: cnt_close = sum(1 for u in vals if _is_close(u, v)) if cnt_close > best_cluster: best_cluster = cnt_close if best_cluster < len(vals): try: vmin, vmax = min(vals), max(vals) medv = median(vals) denom = max(abs(medv), 1.0) spread = (vmax - vmin) / denom except Exception: spread = 0.0 targets.append({"i": i, "series": raw_name, "x": agg["x"][i], "spread": spread}) if not targets: return [] targets.sort(key=lambda d: d.get("spread", 0.0), reverse=True) uniq = [] seen = set() for t in targets: key = (t["series"], t["i"]) if key in seen: continue seen.add(key) uniq.append({"i": t["i"], "series": t["series"], "x": t["x"]}) if len(uniq) >= max_cells: break return uniq def _repair_uncertain_values(self, image_uri: str, data: Dict[str, Any], candidates: List[Dict[str, Any]], max_cells: int = 8) -> Optional[Dict[str, Any]]: targets = self._collect_disagreement_targets(candidates, data, max_cells=max_cells) if not targets: return None targets_json = json.dumps(targets, separators=(",", ":")) x_axis_label = data.get("x_axis_label", "") prompt = ( "You are a precise chart data extractor. Read ONLY the requested cells from the chart image.\n" "For each item in targets, look at the series name and x tick label, read the numeric value from the chart " "(prefer exact data labels if present; otherwise estimate from the axis scale/gridlines), and return STRICT, MINIFIED JSON:\n" "{\"cells\":[{\"i\":int,\"series\":string,\"value\":number|null}, ...]}\n" "Rules:\n" f"- The x-axis label is: {x_axis_label}\n" "- Use the same 0-based index i and exact series string provided in targets.\n" "- Return null only if the chart truly lacks a value for that cell or it is unreadable.\n" "- Return ONLY compact JSON (no extra text, no code fences).\n" f"targets:{targets_json}\n" ) response = self.client.chat.completions.create( model=self.model, messages=[ { "role": "user", "content": [ {"type": "text", "text": prompt}, {"type": "image_url", "image_url": {"url": image_uri}}, ], } ], ) self._accumulate_usage(response) text = (response.choices[0].message.content or "").strip() js = _first_json_obj(text) if not js: return None try: out = json.loads(js) except Exception: return None if not isinstance(out, dict) or "cells" not in out or not isinstance(out["cells"], list): return None series_by_name = {s["name"]: s for s in data.get("series", [])} changed = False for cell in out["cells"]: try: i = int(cell.get("i")) sname = str(cell.get("series")) val = cell.get("value", None) except Exception: continue if sname in series_by_name and 0 <= i < len(series_by_name[sname]["values"]): if val is None: continue try: fv = float(val) except Exception: continue series_by_name[sname]["values"][i] = fv changed = True return data if changed else None def image_to_csv(self, image_path: Path) -> str: image_uri = image_to_data_uri(image_path) seeds = ["[Variant A]", "[Variant B]", "[Variant C]", "[Variant D]", "[Variant E]"] n0 = min(self.ensemble_runs, len(seeds)) candidates: List[Dict[str, Any]] = [] for i in range(n0): cand = self._extract_json_once(image_uri, seed_tag=seeds[i]) if cand: candidates.append(cand) max_extra = 2 target_conf = 0.65 next_seed_idx = n0 extra_used = 0 try: conf = self._agreement_score(candidates) except Exception: conf = 0.0 while extra_used < max_extra and next_seed_idx < len(seeds) and conf < target_conf: cand = self._extract_json_once(image_uri, seed_tag=seeds[next_seed_idx]) next_seed_idx += 1 if cand: candidates.append(cand) extra_used += 1 try: conf = self._agreement_score(candidates) except Exception: conf = 0.0 agg = self._aggregate_candidates(candidates) if candidates else None if agg: repaired_uncertain = self._repair_uncertain_values(image_uri, agg, candidates, max_cells=8) working = repaired_uncertain or agg repaired_missing = self._repair_missing_values(image_uri, working, max_cells=12) final_data = repaired_missing or working return self._to_csv(final_data) return self._fallback_direct_csv(image_uri)

VLM Chart-to-CSV Accuracy

Step Code View