Spaces:
Configuration error
Configuration error
| import base64 | |
| import json | |
| import os | |
| import re | |
| import shutil | |
| import threading | |
| import traceback | |
| import urllib.error | |
| import urllib.parse | |
| import urllib.request | |
| import zipfile | |
| from datetime import datetime, timezone | |
| from uuid import uuid4 | |
| from flask import Flask, jsonify, render_template, request, send_from_directory | |
| from PIL import Image, ImageOps | |
| import anthropic | |
| try: | |
| from openai import OpenAI | |
| except Exception: | |
| OpenAI = None | |
| APP_ROOT = os.path.dirname(os.path.abspath(__file__)) | |
| OUTPUT_DIR = os.path.join(APP_ROOT, "outputs") | |
| GEMINI_API_BASE = "https://generativelanguage.googleapis.com/v1beta" | |
| app = Flask(__name__, static_folder="static", template_folder="templates") | |
| PROMPTS = { | |
| "extract": """The image depicts a figure from a research paper. Please carefully analyze the figure and extract the horizontal (\"x\") and vertical (\"y\") coordinates of each individual data point from each curve. Pay extremely close attention to the axes: | |
| - Identify the tick marks and labels on both axes, and use them to determine the numeric values for every point. | |
| - Read the labels of both axes (the text or symbols denoting what each axis represents). | |
| Data Reporting Requirements: | |
| 1. Separate each curve's data so that every curve has its own two columns of x and y. | |
| 2. The first row of your output must contain the axis labels (for example, \"x-axis label\" for the x-column and \"y-axis label (Curve X)\" for the y-column). Include the curve label in the y-axis header if it's available. | |
| 3. Subsequent rows must present the extracted numeric data in comma-separated values (CSV) format, with: | |
| - No additional words or commentary. | |
| - Exactly one row per data point in each curve. | |
| - For multiple curves, produce additional pairs of columns (two columns per curve: one for x-values, one for y-values). | |
| 4. If, for any reason, you are unable to extract the data, respond with only the word \"None\". | |
| 5. Make sure that both axis are labeled and there are individual values assigned to each tickmark, and that they increase in a logical monotonius fashion. If they are not it may be impossible to extract the data. In that case respond with only the word \"None\". | |
| Key Instructions: | |
| - Look at each point individually; do not assume a pattern or formula. Each data point must be read precisely from its plotted location. | |
| - Verify the numerical axis labels carefully so the extracted data is as accurate as possible. | |
| - Do not include any explanations or text other than the required CSV table. | |
| Output Format Example (if you have two curves): | |
| x-axis label,y-axis label (Curve 1),x-axis label,y-axis label (Curve 2) | |
| 0.1,0.5,0.1,0.75 | |
| 0.2,0.52,0.2,0.77 | |
| ... | |
| (Or \"None\" if you cannot extract the data.) | |
| Remember: The sole output should be either the CSV table (with all columns for the curves) or \"None\". Nothing else. | |
| """, | |
| "code_fix": "The text above is an error produced by your code, please fix the code so that this error does not appear. Use only matplotlib and numpy (plus the Python standard library). Repeat the whole code and only the code so that your whole response can be directly copied and executed. Do not explain and do not say anything else, respond with just the code. Do not output the markdown code fence.", | |
| "compare_x": "You are provided with two images of research plots extracted from academic papers. Do these two plots have the same x-axis (horizontal)? Do they have the same ranges, labels, etc.? Answer with a single word, \"yes\" or \"no\" only.", | |
| "compare_y": "You are provided with two images of research plots extracted from academic papers. Do these two plots have the same y-axis (vertical)? Do they have the same ranges, labels, etc.? Answer with a single word, \"yes\" or \"no\" only.", | |
| "compare_number": "You are provided with two images of research plots extracted from academic papers. Do these two plots have the same number of points (for point plots)? Do the curves look like they connect the same amount of pointis (for line plots)?. Answer with a single word, \"yes\" or \"no\" only.", | |
| "compare_trend": "You are provided with two images of research plots extracted from academic papers. Do these sets of points or curves on these two plots represent the same trends? Do they follow the same patterns? Are points distributed in the same way? Answer with a single word, \"yes\" or \"no\" only.", | |
| } | |
| class GeminiClient: | |
| def __init__(self, api_key): | |
| self.api_key = api_key | |
| def _request(self, method, path, payload=None, params=None): | |
| query = dict(params or {}) | |
| query["key"] = self.api_key | |
| url = f"{GEMINI_API_BASE}{path}?{urllib.parse.urlencode(query)}" | |
| data = None | |
| headers = {} | |
| if payload is not None: | |
| data = json.dumps(payload).encode("utf-8") | |
| headers["Content-Type"] = "application/json" | |
| request_obj = urllib.request.Request(url, data=data, headers=headers, method=method) | |
| try: | |
| with urllib.request.urlopen(request_obj, timeout=120) as response: | |
| body = response.read().decode("utf-8") | |
| except urllib.error.HTTPError as exc: | |
| body = exc.read().decode("utf-8", errors="replace") | |
| try: | |
| payload = json.loads(body) | |
| message = payload.get("error", {}).get("message") or body | |
| except Exception: | |
| message = body | |
| raise RuntimeError(message) from exc | |
| except urllib.error.URLError as exc: | |
| raise RuntimeError(str(exc.reason)) from exc | |
| return json.loads(body) if body else {} | |
| def list_models(self, page_token=""): | |
| params = {"pageSize": 1000} | |
| if page_token: | |
| params["pageToken"] = page_token | |
| return self._request("GET", "/models", params=params) | |
| def generate_content(self, model, contents, max_output_tokens): | |
| model_name = gemini_model_name(model) | |
| model_path = urllib.parse.quote(model_name, safe="/") | |
| return self._request( | |
| "POST", | |
| f"/{model_path}:generateContent", | |
| payload={ | |
| "contents": contents, | |
| "generationConfig": {"maxOutputTokens": max_output_tokens}, | |
| }, | |
| ) | |
| def detect_provider(api_key): | |
| key = (api_key or "").strip() | |
| if key.startswith("sk-ant-"): | |
| return "anthropic" | |
| if key.startswith("sk-"): | |
| return "openai" | |
| return "gemini" if key else "" | |
| def init_provider_client(api_key, provider): | |
| if provider == "anthropic": | |
| return anthropic.Anthropic(api_key=api_key) | |
| if provider == "openai": | |
| if OpenAI is None: | |
| raise RuntimeError("OpenAI SDK is not installed on the server.") | |
| return OpenAI(api_key=api_key) | |
| if provider == "gemini": | |
| return GeminiClient(api_key=api_key) | |
| raise ValueError("Unsupported API key format. Use an Anthropic, OpenAI, or Google Gemini API key.") | |
| def provider_label(provider): | |
| if provider == "anthropic": | |
| return "Anthropic" | |
| if provider == "openai": | |
| return "OpenAI" | |
| if provider == "gemini": | |
| return "Google Gemini" | |
| return "Unknown" | |
| def app_path(path): | |
| prefix = (request.script_root or "").rstrip("/") | |
| return f"{prefix}{path}" | |
| def index(): | |
| return render_template("index.html") | |
| def outputs(filename): | |
| return send_from_directory(OUTPUT_DIR, filename) | |
| def run_extraction(): | |
| api_key = (request.form.get("api_key") or "").strip() | |
| model = (request.form.get("model") or "").strip() | |
| image_file = request.files.get("image") | |
| provider = detect_provider(api_key) | |
| if not api_key: | |
| return jsonify({"error": "Missing API key."}), 400 | |
| if not provider: | |
| return jsonify({"error": "Unsupported API key format. Use an Anthropic, OpenAI, or Google Gemini API key."}), 400 | |
| if not image_file: | |
| return jsonify({"error": "Missing image upload."}), 400 | |
| run_id = uuid4().hex[:10] | |
| run_dir = os.path.join(OUTPUT_DIR, run_id) | |
| os.makedirs(run_dir, exist_ok=True) | |
| ext = os.path.splitext(image_file.filename)[-1].lower() | |
| if ext not in [".png", ".jpg", ".jpeg", ".zip"]: | |
| return jsonify({"error": "Please upload a PNG, JPG, or ZIP file."}), 400 | |
| try: | |
| client = init_provider_client(api_key, provider) | |
| except Exception as exc: | |
| return jsonify({ | |
| "error": f"Failed to initialize {provider_label(provider)} client. Check your API key.", | |
| "details": str(exc), | |
| }), 400 | |
| if not model: | |
| try: | |
| models = fetch_all_models(client, provider) | |
| model = choose_default_model(models, provider) | |
| except Exception as exc: | |
| return jsonify({ | |
| "error": f"Failed to fetch available models from {provider_label(provider)}.", | |
| "details": str(exc), | |
| }), 500 | |
| if not model: | |
| return jsonify({"error": "No available vision-capable models found for this API key."}), 400 | |
| try: | |
| if ext == ".zip": | |
| zip_path = os.path.join(run_dir, "input.zip") | |
| image_file.save(zip_path) | |
| result = prepare_zip_run(zip_path, run_dir, run_id) | |
| if "error" in result: | |
| return jsonify({"error": result["error"]}), 400 | |
| result["base_path"] = request.script_root or "" | |
| start_zip_worker(api_key, provider, model, result) | |
| return jsonify({ | |
| "run_id": run_id, | |
| "status": "zip_pending", | |
| "progress_url": app_path(f"/api/progress/{run_id}"), | |
| "results_zip_url": app_path(f"/outputs/{run_id}/results.zip"), | |
| }) | |
| input_name = f"input{ext if ext != '.jpeg' else '.jpg'}" | |
| input_path = os.path.join(run_dir, input_name) | |
| image_file.save(input_path) | |
| result = process_single_image(client, provider, model, input_path, run_dir) | |
| if result["status"] == "none": | |
| return jsonify({ | |
| "run_id": run_id, | |
| "data_csv": result.get("data_csv", "None"), | |
| "status": "none", | |
| "input_url": app_path(f"/outputs/{run_id}/{input_name}"), | |
| }) | |
| if result["status"] == "error": | |
| return jsonify({ | |
| "error": "Replot code failed after retries.", | |
| "details": result.get("details", ""), | |
| "data_csv": result.get("data_csv", ""), | |
| "run_id": run_id, | |
| }), 500 | |
| return jsonify({ | |
| "run_id": run_id, | |
| "status": "ok", | |
| "data_csv": result["data_csv"], | |
| "code": result["code"], | |
| "validation": result["validation"], | |
| "validated": result["validated"], | |
| "input_url": app_path(f"/outputs/{run_id}/{input_name}"), | |
| "replot_url": app_path(f"/outputs/{run_id}/replot.png"), | |
| "stacked_url": app_path(f"/outputs/{run_id}/stacked.png"), | |
| "comparison_url": app_path(f"/outputs/{run_id}/comparison.png"), | |
| }) | |
| except Exception as exc: | |
| return jsonify({ | |
| "error": "Unexpected error.", | |
| "details": str(exc), | |
| "trace": traceback.format_exc(), | |
| }), 500 | |
| def progress(run_id): | |
| progress_path = os.path.join(OUTPUT_DIR, run_id, "progress.json") | |
| if not os.path.exists(progress_path): | |
| return jsonify({"error": "Progress not found."}), 404 | |
| try: | |
| with open(progress_path, "r", encoding="utf-8") as handle: | |
| return jsonify(json.loads(handle.read())) | |
| except Exception: | |
| return jsonify({"error": "Failed to read progress."}), 500 | |
| def list_models(): | |
| payload = request.get_json(silent=True) or {} | |
| api_key = (payload.get("api_key") or request.form.get("api_key") or "").strip() | |
| provider = detect_provider(api_key) | |
| if not api_key: | |
| return jsonify({"error": "Missing API key."}), 400 | |
| if not provider: | |
| return jsonify({"error": "Unsupported API key format. Use an Anthropic, OpenAI, or Google Gemini API key."}), 400 | |
| try: | |
| client = init_provider_client(api_key, provider) | |
| except Exception as exc: | |
| return jsonify({ | |
| "error": f"Failed to initialize {provider_label(provider)} client. Check your API key.", | |
| "details": str(exc), | |
| }), 400 | |
| try: | |
| models = fetch_all_models(client, provider) | |
| return jsonify({ | |
| "models": models, | |
| "default_model": choose_default_model(models, provider), | |
| "provider": provider, | |
| }) | |
| except Exception as exc: | |
| return jsonify({ | |
| "error": f"Failed to fetch models from {provider_label(provider)}.", | |
| "details": str(exc), | |
| }), 500 | |
| def encode_image(image_path): | |
| with open(image_path, "rb") as image_file: | |
| return base64.b64encode(image_file.read()).decode("utf-8") | |
| def parse_model_list_response(items): | |
| models = [] | |
| seen = set() | |
| for item in items or []: | |
| if isinstance(item, dict): | |
| model_id = item.get("id") | |
| display_name = item.get("display_name") or item.get("name") | |
| created_at = item.get("created_at") or item.get("created") | |
| else: | |
| model_id = getattr(item, "id", None) | |
| display_name = getattr(item, "display_name", None) or getattr(item, "name", None) | |
| created_at = getattr(item, "created_at", None) or getattr(item, "created", None) | |
| if not model_id or model_id in seen: | |
| continue | |
| seen.add(model_id) | |
| created_text = "" | |
| if created_at is not None: | |
| if isinstance(created_at, (int, float)): | |
| try: | |
| created_text = datetime.fromtimestamp(created_at, tz=timezone.utc).isoformat() | |
| except Exception: | |
| created_text = str(created_at) | |
| elif hasattr(created_at, "isoformat"): | |
| created_text = created_at.isoformat() | |
| else: | |
| created_text = str(created_at) | |
| models.append({ | |
| "id": str(model_id), | |
| "display_name": str(display_name or model_id), | |
| "created_at": created_text, | |
| }) | |
| return models | |
| def is_anthropic_vision_model(model): | |
| model_id = (model.get("id") or "").lower() | |
| if not model_id.startswith("claude"): | |
| return False | |
| return any(family in model_id for family in ["opus", "sonnet", "haiku", "fable", "mythos"]) | |
| def is_openai_vision_model(model): | |
| model_id = (model.get("id") or "").lower() | |
| if any(excluded in model_id for excluded in [ | |
| "audio", | |
| "dall-e", | |
| "embedding", | |
| "gpt-image", | |
| "moderation", | |
| "realtime", | |
| "search", | |
| "transcribe", | |
| "tts", | |
| "whisper", | |
| ]): | |
| return False | |
| if model_id.startswith("computer-use-preview"): | |
| return True | |
| if model_id.startswith(("gpt-5", "gpt-4o", "gpt-4.1", "gpt-4.5")): | |
| return True | |
| return bool(re.match(r"^o(1|3|4)(?:-|$)", model_id)) | |
| def is_gemini_vision_model(model): | |
| model_id = (model.get("id") or "").lower() | |
| if "embedding" in model_id: | |
| return False | |
| return "gemini" in model_id | |
| def fetch_all_anthropic_models(client): | |
| page = client.models.list(limit=1000) | |
| collected = [] | |
| while True: | |
| page_data = getattr(page, "data", None) or [] | |
| collected.extend(page_data) | |
| if hasattr(page, "has_next_page") and page.has_next_page(): | |
| page = page.get_next_page() | |
| continue | |
| break | |
| models = parse_model_list_response(collected) | |
| return [model for model in models if is_anthropic_vision_model(model)] | |
| def fetch_all_openai_models(client): | |
| page = client.models.list() | |
| collected = [] | |
| while True: | |
| page_data = getattr(page, "data", None) or [] | |
| collected.extend(page_data) | |
| if hasattr(page, "has_next_page") and page.has_next_page(): | |
| page = page.get_next_page() | |
| continue | |
| break | |
| models = parse_model_list_response(collected) | |
| return [model for model in models if is_openai_vision_model(model)] | |
| def gemini_model_name(model): | |
| value = (model or "").strip() | |
| if value.startswith("models/"): | |
| return value | |
| return f"models/{value}" | |
| def parse_gemini_model_list_response(items): | |
| models = [] | |
| seen = set() | |
| for item in items or []: | |
| model_name = item.get("name") or item.get("baseModelId") | |
| model_id = str(model_name or "").strip() | |
| if model_id.startswith("models/"): | |
| model_id = model_id[len("models/"):] | |
| methods = item.get("supportedGenerationMethods") or [] | |
| if methods and "generateContent" not in methods: | |
| continue | |
| if not model_id or model_id in seen: | |
| continue | |
| seen.add(model_id) | |
| display_name = item.get("displayName") or item.get("baseModelId") or model_id | |
| models.append({ | |
| "id": model_id, | |
| "display_name": str(display_name), | |
| "created_at": "", | |
| }) | |
| return models | |
| def fetch_all_gemini_models(client): | |
| collected = [] | |
| page_token = "" | |
| while True: | |
| page = client.list_models(page_token=page_token) | |
| collected.extend(page.get("models") or []) | |
| page_token = page.get("nextPageToken") or "" | |
| if not page_token: | |
| break | |
| models = parse_gemini_model_list_response(collected) | |
| return [model for model in models if is_gemini_vision_model(model)] | |
| def gemini_sort_key(model): | |
| # Gemini lists carry no release date: the version number tracks release | |
| # order, then pro > flash > flash-lite, then stable before preview. | |
| # "-latest" aliases float to the top. | |
| model_id = (model.get("id") or "").lower() | |
| version = re.search(r"-(\d+(?:\.\d+)?)(?:-|$)", model_id) | |
| if "lite" in model_id: | |
| tier = 1 | |
| elif "flash" in model_id: | |
| tier = 2 | |
| elif "pro" in model_id: | |
| tier = 3 | |
| else: | |
| tier = 0 | |
| return ( | |
| "latest" in model_id, | |
| float(version.group(1)) if version else 0.0, | |
| tier, | |
| not any(tag in model_id for tag in ["preview", "exp"]), | |
| model_id, | |
| ) | |
| def sort_models_newest_first(models, provider): | |
| if provider == "gemini": | |
| return sorted(models, key=gemini_sort_key, reverse=True) | |
| # Anthropic and OpenAI report a creation time (UTC ISO strings sort correctly) | |
| return sorted(models, key=lambda model: (model.get("created_at") or "", model.get("id") or ""), reverse=True) | |
| def fetch_all_models(client, provider): | |
| if provider == "anthropic": | |
| return sort_models_newest_first(fetch_all_anthropic_models(client), provider) | |
| if provider == "openai": | |
| return sort_models_newest_first(fetch_all_openai_models(client), provider) | |
| if provider == "gemini": | |
| return sort_models_newest_first(fetch_all_gemini_models(client), provider) | |
| return [] | |
| def choose_default_model(models, provider): | |
| # fetch_all_models returns lists sorted newest-first, so the default is | |
| # simply the top entry. | |
| return models[0]["id"] if models else "" | |
| def media_type_for_path(image_path): | |
| ext = os.path.splitext(image_path)[-1].lower() | |
| if ext in [".jpg", ".jpeg"]: | |
| return "image/jpeg" | |
| return "image/png" | |
| def write_progress(run_dir, data): | |
| progress_path = os.path.join(run_dir, "progress.json") | |
| try: | |
| with open(progress_path, "w", encoding="utf-8") as handle: | |
| handle.write(json.dumps(data)) | |
| except Exception: | |
| pass | |
| def process_single_image(client, provider, model, input_path, output_dir): | |
| os.makedirs(output_dir, exist_ok=True) | |
| replot_path = os.path.join(output_dir, "replot.png") | |
| stacked_path = os.path.join(output_dir, "stacked.png") | |
| comparison_path = os.path.join(output_dir, "comparison.png") | |
| data_path = os.path.join(output_dir, "data.csv") | |
| code_path = os.path.join(output_dir, "replot_code.py") | |
| validation_path = os.path.join(output_dir, "validation.json") | |
| base64_image = encode_image(input_path) | |
| context, data_csv = prompt_extract(client, provider, model, base64_image, media_type_for_path(input_path)) | |
| if not data_csv or "none" in data_csv.strip().lower(): | |
| try: | |
| with open(data_path, "w", encoding="utf-8") as handle: | |
| handle.write(data_csv.strip() if data_csv else "None") | |
| except Exception: | |
| pass | |
| return {"status": "none", "data_csv": data_csv.strip() if data_csv else "None"} | |
| code, convo = prompt_replot_code(client, provider, model, data_csv, replot_path, context) | |
| code, error_output = exec_replot_code(client, provider, model, code, replot_path, convo) | |
| if error_output: | |
| error_log = os.path.join(output_dir, "replot_error.txt") | |
| try: | |
| with open(error_log, "w", encoding="utf-8") as handle: | |
| handle.write(error_output) | |
| except Exception: | |
| pass | |
| return {"status": "error", "details": error_output, "data_csv": data_csv} | |
| try: | |
| with open(data_path, "w", encoding="utf-8") as handle: | |
| handle.write(data_csv) | |
| with open(code_path, "w", encoding="utf-8") as handle: | |
| handle.write(code) | |
| except Exception: | |
| pass | |
| stacked_path = stack_images_vertically( | |
| input_path, | |
| replot_path, | |
| stacked_path, | |
| border_color=None, | |
| border_size=0, | |
| ) | |
| validation = run_comparisons(client, provider, model, stacked_path) | |
| validated = "yes" if not validation["wrong"] else "no" | |
| comparison_path = stack_images_vertically( | |
| input_path, | |
| replot_path, | |
| comparison_path, | |
| border_color=validated, | |
| border_size=24, | |
| ) | |
| try: | |
| with open(validation_path, "w", encoding="utf-8") as handle: | |
| handle.write(json.dumps(validation, indent=2)) | |
| except Exception: | |
| pass | |
| return { | |
| "status": "ok", | |
| "data_csv": data_csv, | |
| "code": code, | |
| "validation": validation, | |
| "validated": validated, | |
| } | |
| def prepare_zip_run(zip_path, run_dir, run_id): | |
| extract_dir = os.path.join(run_dir, "zip_in") | |
| results_dir = os.path.join(run_dir, "results") | |
| os.makedirs(extract_dir, exist_ok=True) | |
| os.makedirs(results_dir, exist_ok=True) | |
| extract_zip_safe(zip_path, extract_dir) | |
| image_paths = [] | |
| for root, _, files in os.walk(extract_dir): | |
| for name in files: | |
| lower = name.lower() | |
| if lower.endswith((".png", ".jpg", ".jpeg")): | |
| image_paths.append(os.path.join(root, name)) | |
| image_paths.sort() | |
| if not image_paths: | |
| return {"error": "No PNG or JPG files found in the ZIP."} | |
| used_names = set() | |
| jobs = [] | |
| for idx, image_path in enumerate(image_paths, 1): | |
| base = os.path.splitext(os.path.basename(image_path))[0] | |
| safe_base = "".join(ch if ch.isalnum() else "_" for ch in base).strip("_") or f"image_{idx}" | |
| name = safe_base | |
| suffix = 1 | |
| while name in used_names: | |
| suffix += 1 | |
| name = f"{safe_base}_{suffix}" | |
| used_names.add(name) | |
| item_dir = os.path.join(results_dir, name) | |
| os.makedirs(item_dir, exist_ok=True) | |
| ext = os.path.splitext(image_path)[-1].lower() | |
| input_name = f"input{ext if ext != '.jpeg' else '.jpg'}" | |
| input_path = os.path.join(item_dir, input_name) | |
| shutil.copy(image_path, input_path) | |
| jobs.append({ | |
| "name": name, | |
| "input_path": input_path, | |
| "item_dir": item_dir, | |
| }) | |
| write_progress(run_dir, { | |
| "status": "running", | |
| "total": len(jobs), | |
| "processed": 0, | |
| "current": "", | |
| "items": [], | |
| }) | |
| return { | |
| "run_id": run_id, | |
| "run_dir": run_dir, | |
| "results_dir": results_dir, | |
| "jobs": jobs, | |
| } | |
| def start_zip_worker(api_key, provider, model, payload): | |
| thread = threading.Thread( | |
| target=process_zip_async, | |
| args=(api_key, provider, model, payload), | |
| daemon=True, | |
| ) | |
| thread.start() | |
| def process_zip_async(api_key, provider, model, payload): | |
| run_dir = payload["run_dir"] | |
| results_dir = payload["results_dir"] | |
| jobs = payload["jobs"] | |
| try: | |
| client = init_provider_client(api_key, provider) | |
| except Exception as exc: | |
| write_progress(run_dir, { | |
| "status": "error", | |
| "error": f"Failed to initialize {provider_label(provider)} client.", | |
| "details": str(exc), | |
| }) | |
| return | |
| items = [] | |
| total = len(jobs) | |
| for idx, job in enumerate(jobs, 1): | |
| write_progress(run_dir, { | |
| "status": "running", | |
| "total": total, | |
| "processed": idx - 1, | |
| "current": job["name"], | |
| "items": items, | |
| }) | |
| result = process_single_image(client, provider, model, job["input_path"], job["item_dir"]) | |
| items.append({ | |
| "name": job["name"], | |
| "status": result.get("status"), | |
| "validated": result.get("validated", ""), | |
| "wrong_why": result.get("validation", {}).get("wrong_why", "") if result.get("validation") else "", | |
| }) | |
| write_progress(run_dir, { | |
| "status": "running", | |
| "total": total, | |
| "processed": idx, | |
| "current": job["name"], | |
| "items": items, | |
| }) | |
| summary_path = os.path.join(results_dir, "summary.json") | |
| try: | |
| with open(summary_path, "w", encoding="utf-8") as handle: | |
| handle.write(json.dumps({"items": items}, indent=2)) | |
| except Exception: | |
| pass | |
| zip_out = os.path.join(run_dir, "results.zip") | |
| with zipfile.ZipFile(zip_out, "w", compression=zipfile.ZIP_DEFLATED) as zf: | |
| for root, _, files in os.walk(results_dir): | |
| for file_name in files: | |
| full_path = os.path.join(root, file_name) | |
| rel_path = os.path.relpath(full_path, results_dir) | |
| zf.write(full_path, arcname=rel_path) | |
| write_progress(run_dir, { | |
| "status": "done", | |
| "total": total, | |
| "processed": total, | |
| "current": "", | |
| "items": items, | |
| "results_zip_url": f"{(payload.get('base_path') or '').rstrip('/')}/outputs/{payload['run_id']}/results.zip", | |
| }) | |
| def extract_zip_safe(zip_path, dest_dir): | |
| with zipfile.ZipFile(zip_path, "r") as zf: | |
| for member in zf.infolist(): | |
| if member.filename.endswith("/"): | |
| continue | |
| extracted_path = os.path.normpath(os.path.join(dest_dir, member.filename)) | |
| if not extracted_path.startswith(os.path.abspath(dest_dir) + os.sep): | |
| continue | |
| os.makedirs(os.path.dirname(extracted_path), exist_ok=True) | |
| with zf.open(member) as source, open(extracted_path, "wb") as target: | |
| shutil.copyfileobj(source, target) | |
| def extract_openai_text(response): | |
| if isinstance(response, dict): | |
| output_text = response.get("output_text") | |
| if output_text: | |
| return str(output_text) | |
| outputs = response.get("output") or [] | |
| else: | |
| output_text = getattr(response, "output_text", None) | |
| outputs = getattr(response, "output", None) or [] | |
| if output_text: | |
| return output_text | |
| parts = [] | |
| for item in outputs: | |
| content = item.get("content") if isinstance(item, dict) else getattr(item, "content", None) | |
| for block in content or []: | |
| block_type = block.get("type") if isinstance(block, dict) else getattr(block, "type", None) | |
| if block_type not in ["output_text", "text"]: | |
| continue | |
| text = block.get("text") if isinstance(block, dict) else getattr(block, "text", None) | |
| if not text and not isinstance(block, dict): | |
| text = getattr(block, "value", None) | |
| if text: | |
| parts.append(str(text)) | |
| return "\n".join(parts).strip() | |
| def extract_gemini_text(response): | |
| parts = [] | |
| for candidate in response.get("candidates") or []: | |
| content = candidate.get("content") or {} | |
| for part in content.get("parts") or []: | |
| text = part.get("text") | |
| if text: | |
| parts.append(str(text)) | |
| return "\n".join(parts).strip() | |
| def prompt_extract(client, provider, model, base64_image, media_type): | |
| if provider == "anthropic": | |
| messages = [ | |
| { | |
| "role": "user", | |
| "content": [ | |
| { | |
| "type": "image", | |
| "source": { | |
| "type": "base64", | |
| "media_type": media_type, | |
| "data": base64_image, | |
| }, | |
| }, | |
| {"type": "text", "text": PROMPTS["extract"]}, | |
| ], | |
| } | |
| ] | |
| response = client.messages.create( | |
| model=model, | |
| max_tokens=2048, | |
| #temperature=0, | |
| messages=messages, | |
| ) | |
| return messages, response.content[0].text | |
| if provider == "gemini": | |
| messages = [ | |
| { | |
| "role": "user", | |
| "parts": [ | |
| { | |
| "inline_data": { | |
| "mime_type": media_type, | |
| "data": base64_image, | |
| }, | |
| }, | |
| {"text": PROMPTS["extract"]}, | |
| ], | |
| } | |
| ] | |
| response = client.generate_content(model, messages, max_output_tokens=2048) | |
| return messages, extract_gemini_text(response) | |
| image_data_url = f"data:{media_type};base64,{base64_image}" | |
| messages = [ | |
| { | |
| "role": "user", | |
| "content": [ | |
| {"type": "input_image", "image_url": image_data_url}, | |
| {"type": "input_text", "text": PROMPTS["extract"]}, | |
| ], | |
| } | |
| ] | |
| response = client.responses.create( | |
| model=model, | |
| input=messages, | |
| max_output_tokens=2048, | |
| ) | |
| return messages, extract_openai_text(response) | |
| def prompt_replot_code(client, provider, model, data_csv, replot_path, messages): | |
| prompt = ( | |
| "Please analyze the figure and create a python code that will reproduce the plot exactly, " | |
| "including colors, line types, point shapes, axis labels, axis ranges, etc. " | |
| f"Save the plot as a file \"{replot_path}\" only and do not show it. " | |
| "Use only matplotlib and numpy (plus the Python standard library). " | |
| "Respond with the code only so that it can be directly copied and executed, do not output the markdown code fence.\n\n" | |
| f"Use the following data on the plot: {data_csv}" | |
| ) | |
| if provider == "anthropic": | |
| convo = list(messages) | |
| convo.append({"role": "assistant", "content": data_csv}) | |
| convo.append({"role": "user", "content": prompt}) | |
| response = client.messages.create( | |
| model=model, | |
| max_tokens=2048, | |
| #temperature=0, | |
| messages=convo, | |
| ) | |
| return response.content[0].text, convo | |
| if provider == "gemini": | |
| convo = list(messages) | |
| convo.append({"role": "model", "parts": [{"text": data_csv}]}) | |
| convo.append({"role": "user", "parts": [{"text": prompt}]}) | |
| response = client.generate_content(model, convo, max_output_tokens=2048) | |
| return extract_gemini_text(response), convo | |
| convo = list(messages) | |
| convo.append({ | |
| "role": "user", | |
| "content": [ | |
| { | |
| "type": "input_text", | |
| "text": f"Use the following extracted CSV data:\n{data_csv}\n\n{prompt}", | |
| } | |
| ], | |
| }) | |
| response = client.responses.create( | |
| model=model, | |
| input=convo, | |
| max_output_tokens=2048, | |
| ) | |
| return extract_openai_text(response), convo | |
| def exec_replot_code(client, provider, model, code, replot_path, convo=None): | |
| error_output = None | |
| def strip_code_fences(text): | |
| if not text: | |
| return text | |
| cleaned = text.strip() | |
| if cleaned.startswith("```"): | |
| lines = cleaned.splitlines() | |
| lines = lines[1:] | |
| if lines and lines[-1].strip().startswith("```"): | |
| lines = lines[:-1] | |
| cleaned = "\n".join(lines).strip() | |
| return cleaned | |
| def run_code(script): | |
| exec_globals = {"__name__": "__main__"} | |
| try: | |
| import matplotlib | |
| matplotlib.use("Agg") | |
| import matplotlib.pyplot as plt | |
| plt.show = lambda *args, **kwargs: None | |
| exec_globals["plt"] = plt | |
| except Exception: | |
| pass | |
| exec(script, exec_globals) | |
| try: | |
| code = strip_code_fences(code) | |
| run_code(code) | |
| except Exception: | |
| error_output = traceback.format_exc() | |
| attempts = 0 | |
| while error_output and attempts < 2: | |
| attempts += 1 | |
| fix_prompt = error_output + "\n\n" + PROMPTS["code_fix"] | |
| if provider == "anthropic": | |
| fix_messages = [ | |
| {"role": "assistant", "content": code}, | |
| {"role": "user", "content": fix_prompt}, | |
| ] | |
| if convo: | |
| fix_messages = list(convo) + fix_messages | |
| response = client.messages.create( | |
| model=model, | |
| max_tokens=2048, | |
| #temperature=0, | |
| messages=fix_messages, | |
| ) | |
| code = strip_code_fences(response.content[0].text) | |
| elif provider == "gemini": | |
| fix_messages = [] | |
| if convo: | |
| fix_messages = list(convo) | |
| fix_messages.extend([ | |
| {"role": "model", "parts": [{"text": code}]}, | |
| {"role": "user", "parts": [{"text": fix_prompt}]}, | |
| ]) | |
| response = client.generate_content(model, fix_messages, max_output_tokens=2048) | |
| code = strip_code_fences(extract_gemini_text(response)) | |
| else: | |
| fix_messages = [] | |
| if convo: | |
| fix_messages = list(convo) + fix_messages | |
| fix_messages.append( | |
| { | |
| "role": "user", | |
| "content": [ | |
| { | |
| "type": "input_text", | |
| "text": f"Current code:\n{code}\n\nExecution error:\n{fix_prompt}", | |
| } | |
| ], | |
| } | |
| ) | |
| response = client.responses.create( | |
| model=model, | |
| input=fix_messages, | |
| max_output_tokens=2048, | |
| ) | |
| code = strip_code_fences(extract_openai_text(response)) | |
| try: | |
| run_code(code) | |
| error_output = None | |
| except Exception: | |
| error_output = traceback.format_exc() | |
| return code, error_output | |
| def run_comparisons(client, provider, model, stacked_path): | |
| base64_image = encode_image(stacked_path) | |
| media_type = media_type_for_path(stacked_path) | |
| results = {} | |
| wrong = False | |
| wrong_why = [] | |
| for key, label in [ | |
| ("compare_x", "X"), | |
| ("compare_y", "Y"), | |
| ("compare_number", "N"), | |
| ("compare_trend", "T"), | |
| ]: | |
| if provider == "anthropic": | |
| response = client.messages.create( | |
| model=model, | |
| max_tokens=32, | |
| #temperature=0, | |
| messages=[ | |
| { | |
| "role": "user", | |
| "content": [ | |
| { | |
| "type": "image", | |
| "source": { | |
| "type": "base64", | |
| "media_type": media_type, | |
| "data": base64_image, | |
| }, | |
| }, | |
| {"type": "text", "text": PROMPTS[key]}, | |
| ], | |
| } | |
| ], | |
| ) | |
| answer = response.content[0].text.strip().lower() | |
| elif provider == "gemini": | |
| response = client.generate_content( | |
| model, | |
| [ | |
| { | |
| "role": "user", | |
| "parts": [ | |
| { | |
| "inline_data": { | |
| "mime_type": media_type, | |
| "data": base64_image, | |
| }, | |
| }, | |
| {"text": PROMPTS[key]}, | |
| ], | |
| } | |
| ], | |
| max_output_tokens=32, | |
| ) | |
| answer = extract_gemini_text(response).strip().lower() | |
| else: | |
| image_data_url = f"data:{media_type};base64,{base64_image}" | |
| response = client.responses.create( | |
| model=model, | |
| input=[ | |
| { | |
| "role": "user", | |
| "content": [ | |
| {"type": "input_image", "image_url": image_data_url}, | |
| {"type": "input_text", "text": PROMPTS[key]}, | |
| ], | |
| } | |
| ], | |
| max_output_tokens=32, | |
| ) | |
| answer = extract_openai_text(response).strip().lower() | |
| results[key] = answer | |
| if answer.startswith("no"): | |
| wrong = True | |
| wrong_why.append(label) | |
| return { | |
| "answers": results, | |
| "wrong": wrong, | |
| "wrong_why": "; ".join(wrong_why) if wrong_why else "", | |
| } | |
| def stack_images_vertically(image1_path, image2_path, output_path, border_color=None, border_size=30): | |
| img1 = Image.open(image1_path).convert("RGB") | |
| img2 = Image.open(image2_path).convert("RGB") | |
| width = max(img1.width, img2.width) | |
| img1_resized = img1.resize((width, int(img1.height * width / img1.width)), Image.LANCZOS) | |
| img2_resized = img2.resize((width, int(img2.height * width / img2.width)), Image.LANCZOS) | |
| combined_height = img1_resized.height + img2_resized.height | |
| combined = Image.new("RGB", (width, combined_height), (0, 0, 0)) | |
| combined.paste(img1_resized, (0, 0)) | |
| combined.paste(img2_resized, (0, img1_resized.height)) | |
| if border_color: | |
| if "yes" in border_color.lower(): | |
| color = (16, 185, 129) | |
| elif "no" in border_color.lower(): | |
| color = (239, 68, 68) | |
| else: | |
| color = (99, 102, 241) | |
| combined = ImageOps.expand(combined, border=border_size, fill=color) | |
| combined.save(output_path) | |
| return output_path | |
| if __name__ == "__main__": | |
| os.makedirs(OUTPUT_DIR, exist_ok=True) | |
| app.run(debug=True, port=5173) | |