# -*- 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 base64 import zipfile import urllib.request import ssl import gradio as gr from PIL import Image 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"}) # Accept both PIL Image and base64 string if isinstance(image, str): # Strip data URI prefix if present (e.g. "data:image/png;base64,...") if image.startswith("data:"): image = image.split(",", 1)[1] img_bytes = base64.b64decode(image) print(f" Base64 decoded: {len(img_bytes)} bytes, first 4 bytes: {img_bytes[:4]}") pil_img = Image.open(io.BytesIO(img_bytes)).convert("RGB") img_rgb = np.array(pil_img) else: 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 def create_fastapi_app(): """Create FastAPI app with custom routes and mount Gradio.""" from fastapi import FastAPI, Request from fastapi.responses import JSONResponse from fastapi.middleware.cors import CORSMiddleware app = FastAPI() # Enable CORS for cross-origin requests from StarryDigitizer app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"], ) @app.post("/custom_api/digitize_base64") async def api_digitize_base64(request: Request): body = await request.json() data = body.get("data", []) if len(data) < 7: return JSONResponse({"error": "Expected 7 parameters in data array"}, status_code=400) image_base64 = data[0] print(f" Received image string: type={type(image_base64)}, len={len(image_base64) if isinstance(image_base64, str) else 'N/A'}, first 50 chars: {str(image_base64)[:50]}") auto_axis = bool(data[1]) downsample_mode = str(data[2]) max_points = int(data[3]) fixed_step = int(data[4]) sort_mode = str(data[5]) output_format = str(data[6]) result = digitize_chart( image_base64, auto_axis, downsample_mode, max_points, fixed_step, sort_mode, output_format ) return JSONResponse({"data": [result]}) return app if __name__ == "__main__": import uvicorn print("Loading models...") load_models() print("All models loaded. Starting server...") demo = create_app() demo.queue() fastapi_app = create_fastapi_app() # Mount Gradio app onto FastAPI gr.mount_gradio_app(fastapi_app, demo, path="/") uvicorn.run(fastapi_app, host="0.0.0.0", port=7860)