# -*- coding: utf-8 -*- """ AutoLineDigitizer API Server for Hugging Face Spaces. Exposes chart line extraction as a Gradio API. """ import sys import os # Setup paths (same as desktop_app.py) SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__)) CHARTDETE_DIR = os.path.join(SCRIPT_DIR, "submodules", "chartdete") LINEFORMER_DIR = os.path.join(SCRIPT_DIR, "submodules", "lineformer") MMDET_DIR = os.path.join(LINEFORMER_DIR, "mmdetection") SRC_DIR = os.path.join(SCRIPT_DIR, "src") sys.path.insert(0, SCRIPT_DIR) sys.path.insert(0, SRC_DIR) sys.path.insert(0, MMDET_DIR) sys.path.insert(0, LINEFORMER_DIR) sys.path.insert(0, CHARTDETE_DIR) # Register ChartDete custom models CHARTDETE_AVAILABLE = False try: import mmdet # noqa: F401 from mmdet.models.roi_heads.cascade_roi_head_LGF import CascadeRoIHead_LGF # noqa: F401 CHARTDETE_AVAILABLE = True except Exception as e: print(f"ChartDete custom models not available: {e}") import cv2 import numpy as np import json import io import zipfile import urllib.request import ssl import gradio as gr from datetime import datetime, timezone # ============================================================ # Model download & loading # ============================================================ GITHUB_REPO = "t29mato/AutoLineDigitizer" GITHUB_RELEASE_TAG = "models" # Model files: filename -> download source MODEL_FILES = { "iter_3000.pth": { "source": "github", "url": f"https://github.com/{GITHUB_REPO}/releases/download/{GITHUB_RELEASE_TAG}/iter_3000.pth", }, "checkpoint.pth": { "source": "github", "url": f"https://github.com/{GITHUB_REPO}/releases/download/{GITHUB_RELEASE_TAG}/checkpoint.pth", }, } _infer_module = None _chartdete_module = None def download_file(url, dest_path): """Download a file with progress logging and SSL fallback.""" print(f" Downloading {os.path.basename(dest_path)} from {url} ...") tmp_path = dest_path + ".tmp" req = urllib.request.Request(url) try: response_ctx = urllib.request.urlopen(req) except Exception: ctx = ssl._create_unverified_context() response_ctx = urllib.request.urlopen(req, context=ctx) with response_ctx as response: total_size = int(response.headers.get("Content-Length", 0)) downloaded = 0 block_size = 1024 * 1024 # 1MB with open(tmp_path, "wb") as f: while True: chunk = response.read(block_size) if not chunk: break f.write(chunk) downloaded += len(chunk) if total_size > 0: pct = downloaded * 100 // total_size print(f" {os.path.basename(dest_path)}: {downloaded // (1024*1024)}MB / {total_size // (1024*1024)}MB ({pct}%)") os.replace(tmp_path, dest_path) print(f" {os.path.basename(dest_path)} downloaded successfully.") def ensure_models(): """Download model files if they don't exist.""" models_dir = os.path.join(SCRIPT_DIR, "models") os.makedirs(models_dir, exist_ok=True) for filename, info in MODEL_FILES.items(): dest = os.path.join(models_dir, filename) if not os.path.exists(dest): download_file(info["url"], dest) else: print(f" {filename} already exists, skipping download.") def load_models(): """Download (if needed) and load LineFormer and ChartDete models.""" global _infer_module, _chartdete_module print("Checking model files...") ensure_models() models_dir = os.path.join(SCRIPT_DIR, "models") # Load LineFormer import infer config_path = os.path.join(LINEFORMER_DIR, "lineformer_swin_t_config.py") ckpt_path = os.path.join(models_dir, "iter_3000.pth") infer.load_model(config_path, ckpt_path, "cpu") _infer_module = infer print("LineFormer model loaded.") # Load ChartDete if CHARTDETE_AVAILABLE: import chartdete_infer chartdete_config = os.path.join(SCRIPT_DIR, "config", "chartdete_config.py") chartdete_ckpt = os.path.join(models_dir, "checkpoint.pth") chartdete_infer.load_chartdete_model( config_path=chartdete_config, checkpoint_path=chartdete_ckpt, device="cpu", ) _chartdete_module = chartdete_infer print("ChartDete model loaded.") else: print("ChartDete not available, axis detection disabled.") def extract_lines(img): """Run LineFormer inference and return raw centerline points.""" line_dataseries = _infer_module.get_dataseries(img, to_clean=False) raw_lines = [] for line in line_dataseries: if len(line) == 0: continue raw_lines.append([[int(pt["x"]), int(pt["y"])] for pt in line]) return raw_lines def arc_length_resample(points, n_points): """Resample at equidistant intervals along pixel-space arc length.""" pts = np.array(points, dtype=float) diffs = np.diff(pts, axis=0) seg_lengths = np.sqrt((diffs ** 2).sum(axis=1)) cum_arc = np.zeros(len(pts)) cum_arc[1:] = np.cumsum(seg_lengths) total_length = cum_arc[-1] if total_length == 0: return [points[0]] target_distances = np.linspace(0, total_length, n_points) result = [] seg_idx = 0 for d in target_distances: while seg_idx < len(seg_lengths) - 1 and cum_arc[seg_idx + 1] < d: seg_idx += 1 seg_span = cum_arc[seg_idx + 1] - cum_arc[seg_idx] t = 0.0 if seg_span == 0 else (d - cum_arc[seg_idx]) / seg_span x = pts[seg_idx, 0] + t * (pts[seg_idx + 1, 0] - pts[seg_idx, 0]) y = pts[seg_idx, 1] + t * (pts[seg_idx + 1, 1] - pts[seg_idx, 1]) result.append([int(round(x)), int(round(y))]) return result def downsample_points(points, mode="max_points", max_points=20, fixed_step=10): """Downsample points based on mode.""" if len(points) <= 1: return points if mode == "none": return points elif mode == "fixed": return points[::fixed_step] elif mode == "max_points": if len(points) <= max_points: return points step = max(1, len(points) // max_points) return points[::step] elif mode == "arc_length": if len(points) <= max_points: return points return arc_length_resample(points, max_points) return points def detect_axis_calibration(img): """Detect axis calibration using ChartDete + OCR.""" if _chartdete_module is None: return None detections = _chartdete_module.detect_chart_elements(img, score_thr=0.3) axis_info = _chartdete_module.get_axis_info(detections, img=img, with_ocr=True) calibration = axis_info.get("calibration") if calibration is None: return None has_x = "x1_pixel" in calibration and "x2_pixel" in calibration has_y = "y1_pixel" in calibration and "y2_pixel" in calibration if not (has_x and has_y): return None plot_area = axis_info.get("plot_area") if plot_area: x_calib_y = plot_area[3] y_calib_x = plot_area[0] else: x_calib_y = img.shape[0] * 0.9 y_calib_x = img.shape[1] * 0.1 return { "x1_px": calibration["x1_pixel"], "x1_py": x_calib_y, "x1_val": calibration["x1_value"], "x2_px": calibration["x2_pixel"], "x2_py": x_calib_y, "x2_val": calibration["x2_value"], "y1_px": y_calib_x, "y1_py": calibration["y2_pixel"], "y1_val": calibration["y2_value"], "y2_px": y_calib_x, "y2_py": calibration["y1_pixel"], "y2_val": calibration["y1_value"], "xIsLogScale": False, "yIsLogScale": False, } def convert_to_starry_digitizer_format(data_series, img_shape, axis_config=None): """Convert extracted data to StarryDigitizer project.json format.""" timestamp = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%S.%f")[:-3] + "Z" if axis_config is None: axis_set = { "id": 1, "name": "XY Axes 1", "x1": {"name": "x1", "value": 0, "coord": {"xPx": 0, "yPx": float(img_shape[0])}}, "x2": {"name": "x2", "value": 100, "coord": {"xPx": float(img_shape[1]), "yPx": float(img_shape[0])}}, "y1": {"name": "y1", "value": 0, "coord": {"xPx": 0, "yPx": float(img_shape[0])}}, "y2": {"name": "y2", "value": 100, "coord": {"xPx": 0, "yPx": 0}}, "xIsLogScale": False, "yIsLogScale": False, "considerGraphTilt": False, "pointMode": 0, "isVisible": True, } else: axis_set = { "id": 1, "name": "XY Axes 1", "x1": {"name": "x1", "value": axis_config["x1_val"], "coord": {"xPx": axis_config["x1_px"], "yPx": axis_config["x1_py"]}}, "x2": {"name": "x2", "value": axis_config["x2_val"], "coord": {"xPx": axis_config["x2_px"], "yPx": axis_config["x2_py"]}}, "y1": {"name": "y1", "value": axis_config["y1_val"], "coord": {"xPx": axis_config["y1_px"], "yPx": axis_config["y1_py"]}}, "y2": {"name": "y2", "value": axis_config["y2_val"], "coord": {"xPx": axis_config["y2_px"], "yPx": axis_config["y2_py"]}}, "xIsLogScale": axis_config.get("xIsLogScale", False), "yIsLogScale": axis_config.get("yIsLogScale", False), "considerGraphTilt": False, "pointMode": 0, "isVisible": True, } datasets = [{ "id": 1, "name": "dataset 1", "axisSetId": 1, "points": [], "visiblePointIds": [], "manuallyAddedPointIds": [], }] for idx, series in enumerate(data_series): points = [] visible_ids = [] for pt_idx, pt in enumerate(series["points"]): pt_id = pt_idx + 1 points.append({"id": pt_id, "xPx": float(pt[0]), "yPx": float(pt[1])}) visible_ids.append(pt_id) datasets.append({ "id": idx + 2, "name": f"Line {idx + 1}", "axisSetId": 1, "points": points, "visiblePointIds": visible_ids, "manuallyAddedPointIds": [], }) return { "version": "1.11.2", "timestamp": timestamp, "axisSets": [axis_set], "activeAxisSetId": 1, "datasets": datasets, "activeDatasetId": len(datasets), "canvasHandler": {"scale": 1.0, "manualMode": 0}, } # ============================================================ # Gradio API endpoint # ============================================================ def digitize_chart( image, auto_axis_detection: bool = True, downsample_mode: str = "arc_length", max_points: int = 20, fixed_step: int = 10, sort_mode: str = "mean_y_desc", output_format: str = "starry_digitizer_json", ): """ Extract line data from a chart image. Args: image: Input chart image (PIL Image from Gradio) auto_axis_detection: Enable ChartDete + OCR axis detection downsample_mode: "none", "max_points", "fixed", or "arc_length" max_points: Max points per line (for max_points/arc_length modes) fixed_step: Step size (for fixed mode) sort_mode: "original", "mean_y_desc", or "mean_y_asc" output_format: "starry_digitizer_json", "starry_digitizer_zip", or "json" Returns: For starry_digitizer_json/json: JSON string For starry_digitizer_zip: ZIP file path """ if image is None: return json.dumps({"error": "No image provided"}) # Convert PIL Image to BGR numpy array img_rgb = np.array(image) img_bgr = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2BGR) # Extract lines raw_lines = extract_lines(img_bgr) # Downsample data_series = [] for all_points in raw_lines: points = downsample_points(all_points, downsample_mode, max_points, fixed_step) data_series.append({"points": points}) # Sort if sort_mode == "mean_y_desc" and len(data_series) > 0: data_series = sorted(data_series, key=lambda s: np.mean([pt[1] for pt in s["points"]])) elif sort_mode == "mean_y_asc" and len(data_series) > 0: data_series = sorted(data_series, key=lambda s: np.mean([pt[1] for pt in s["points"]]), reverse=True) # Axis detection axis_config = None if auto_axis_detection and _chartdete_module is not None: axis_config = detect_axis_calibration(img_bgr) # Build output if output_format == "starry_digitizer_zip": project_json = convert_to_starry_digitizer_format(data_series, img_bgr.shape, axis_config) zip_buffer = io.BytesIO() with zipfile.ZipFile(zip_buffer, "w", zipfile.ZIP_DEFLATED) as zf: _, img_encoded = cv2.imencode(".png", img_bgr) zf.writestr("image.png", img_encoded.tobytes()) zf.writestr("project.json", json.dumps(project_json, indent=2, ensure_ascii=False)) zip_buffer.seek(0) import tempfile tmp = tempfile.NamedTemporaryFile(suffix=".zip", delete=False) tmp.write(zip_buffer.read()) tmp.close() return tmp.name elif output_format == "starry_digitizer_json": project_json = convert_to_starry_digitizer_format(data_series, img_bgr.shape, axis_config) return json.dumps(project_json, ensure_ascii=False) else: # "json" - raw extraction result result = { "num_lines": len(data_series), "lines": [ {"line_index": i, "num_points": len(s["points"]), "points": s["points"]} for i, s in enumerate(data_series) ], "axis_config": axis_config, "image_shape": {"height": img_bgr.shape[0], "width": img_bgr.shape[1]}, } return json.dumps(result, ensure_ascii=False) # ============================================================ # Gradio Interface # ============================================================ def create_app(): """Create Gradio app.""" with gr.Blocks(title="AutoLineDigitizer API") as demo: gr.Markdown("# AutoLineDigitizer API") gr.Markdown( "Upload a chart image to automatically extract line data. " "Results can be imported into [StarryDigitizer](https://digitizer.starrydata.org/)." ) with gr.Row(): with gr.Column(): image_input = gr.Image(type="pil", label="Chart Image") auto_axis = gr.Checkbox(value=True, label="Auto Axis Detection (ChartDete + OCR)") downsample = gr.Dropdown( choices=["none", "max_points", "fixed", "arc_length"], value="arc_length", label="Downsample Mode", ) max_pts = gr.Slider(minimum=5, maximum=200, value=20, step=1, label="Max Points per Line") fixed_stp = gr.Slider(minimum=1, maximum=50, value=10, step=1, label="Fixed Step") sort = gr.Dropdown( choices=["original", "mean_y_desc", "mean_y_asc"], value="mean_y_desc", label="Sort Mode", ) out_fmt = gr.Dropdown( choices=["starry_digitizer_json", "starry_digitizer_zip", "json"], value="starry_digitizer_json", label="Output Format", ) run_btn = gr.Button("Extract Lines", variant="primary") with gr.Column(): output = gr.Textbox(label="Result (JSON)", lines=20, max_lines=50) file_output = gr.File(label="Download ZIP", visible=False) def on_run(image, auto_axis, downsample, max_pts, fixed_stp, sort, out_fmt): result = digitize_chart(image, auto_axis, downsample, max_pts, fixed_stp, sort, out_fmt) if out_fmt == "starry_digitizer_zip": return gr.update(value="ZIP file generated. Download below."), gr.update(value=result, visible=True) else: return gr.update(value=result), gr.update(visible=False) run_btn.click( fn=on_run, inputs=[image_input, auto_axis, downsample, max_pts, fixed_stp, sort, out_fmt], outputs=[output, file_output], ) return demo if __name__ == "__main__": print("Loading models...") load_models() print("All models loaded. Starting Gradio server...") demo = create_app() demo.queue() demo.launch(server_name="0.0.0.0", server_port=7860)