File size: 2,846 Bytes
a4d7d23 | 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 | """FastAPI service for DiagraMine.
POST /extract multipart image upload -> structured JSON + annotated PNG
(base64) + relationship CSV (text). Always returns
schema-valid JSON (the reliability guarantee).
GET /health liveness probe.
Run: uvicorn src.api:app --port 8000
"""
from __future__ import annotations
import base64
import csv
import io
import os
import tempfile
import cv2
from fastapi import FastAPI, File, Query, UploadFile
from fastapi.responses import JSONResponse, RedirectResponse
from pydantic import BaseModel
from src import pipeline
from src.schemas import ExtractionResult, PipelineConfig
app = FastAPI(
title="DiagraMine",
description="Diagram understanding & structure extraction (data-driven, no hardcoded connections).",
version="4.0",
)
class ExtractResponse(BaseModel):
structure: ExtractionResult
annotated_png_base64: str
relationship_csv: str
def _csv_string(result: ExtractionResult) -> str:
buf = io.StringIO()
w = csv.writer(buf)
w.writerow(["source", "target", "line_style", "direction", "relationship", "bidirectional"])
for r in result.relationships:
w.writerow([r.source, r.target, r.line_style, r.direction, r.relationship, r.bidirectional])
return buf.getvalue()
@app.get("/", include_in_schema=False)
async def root() -> RedirectResponse:
"""The service has no page of its own; send visitors to the interactive docs."""
return RedirectResponse(url="/docs")
@app.get("/health")
async def health() -> dict:
return {"status": "ok", "service": "diagramine", "version": "4.0"}
@app.post("/extract", response_model=ExtractResponse)
async def extract(
file: UploadFile = File(...),
text: str = Query("paddleocr"),
box: str = Query("canny_contours"),
arrow: str = Query("directed_lines"),
icon: str = Query("template_matching"),
outside_box_gate: bool = Query(False),
) -> JSONResponse:
config = PipelineConfig(
text_detector=text, box_detector=box, arrow_detector=arrow,
icon_detector=icon, outside_box_gate=outside_box_gate,
)
suffix = os.path.splitext(file.filename or "upload.png")[1] or ".png"
data = await file.read()
tmp = tempfile.NamedTemporaryFile(delete=False, suffix=suffix)
try:
tmp.write(data)
tmp.close()
result = pipeline.extract(tmp.name, config)
annotated = pipeline.annotate(tmp.name, result)
ok, png = cv2.imencode(".png", annotated)
png_b64 = base64.b64encode(png.tobytes()).decode("ascii") if ok else ""
finally:
os.unlink(tmp.name)
payload = ExtractResponse(
structure=result,
annotated_png_base64=png_b64,
relationship_csv=_csv_string(result),
)
return JSONResponse(content=payload.model_dump())
|