import io import os import tempfile from pathlib import Path from fastapi import FastAPI, UploadFile, HTTPException, Query from fastapi.responses import StreamingResponse from pydantic import BaseModel from docling.document_converter import DocumentConverter from pdf2image import convert_from_bytes from PIL import Image from bbox import to_topleft_list app = FastAPI(title="Docling Studio - Document Parser") converter = DocumentConverter() MAX_FILE_SIZE = 50 * 1024 * 1024 # 50 MB # --- Response models --- class PageElement(BaseModel): type: str bbox: list[float] content: str class PageDetail(BaseModel): page_number: int width: float height: float elements: list[PageElement] class ParseResponse(BaseModel): filename: str page_count: int content_markdown: str content_html: str pages: list[PageDetail] # --- Helpers --- def extract_pages_detail(doc_result) -> list[PageDetail]: """Extract per-page element details with bounding boxes from Docling result.""" pages: dict[int, PageDetail] = {} document = doc_result.document # Get page dimensions from document pages if hasattr(document, 'pages') and document.pages: for page_key, page_obj in document.pages.items(): page_no = int(page_key) if isinstance(page_key, str) else page_key width = page_obj.size.width if hasattr(page_obj, 'size') and page_obj.size else 612.0 height = page_obj.size.height if hasattr(page_obj, 'size') and page_obj.size else 792.0 pages[page_no] = PageDetail( page_number=page_no, width=width, height=height, elements=[] ) # Iterate document items to extract elements with bounding boxes if hasattr(document, 'body') and hasattr(document.body, 'children'): _extract_elements_recursive(document, document.body.children, pages) # Fallback: if we have content items if hasattr(document, 'texts'): for text_item in document.texts: _process_content_item(text_item, pages) # Sort by page number return sorted(pages.values(), key=lambda p: p.page_number) def _process_content_item(item, pages: dict[int, PageDetail]): """Process a single content item and add it to the appropriate page.""" if not hasattr(item, 'prov') or not item.prov: return for prov in item.prov: page_no = prov.page_no if hasattr(prov, 'page_no') else 1 if page_no not in pages: pages[page_no] = PageDetail( page_number=page_no, width=612.0, height=792.0, elements=[] ) page_height = pages[page_no].height bbox = [0, 0, 0, 0] if hasattr(prov, 'bbox') and prov.bbox: b = prov.bbox if hasattr(b, 'l'): bbox = to_topleft_list(b, page_height) elif isinstance(b, (list, tuple)) and len(b) >= 4: bbox = list(b[:4]) element_type = _get_element_type(item) content = "" if hasattr(item, 'text'): content = item.text or "" elif hasattr(item, 'export_to_markdown'): content = item.export_to_markdown() pages[page_no].elements.append(PageElement( type=element_type, bbox=bbox, content=content[:500] # Truncate long content )) def _extract_elements_recursive(document, children, pages: dict[int, PageDetail]): """Recursively extract elements from document tree.""" if not children: return for child_ref in children: # Resolve reference if needed item = child_ref if hasattr(child_ref, '__ref__'): item = document.resolve_ref(child_ref) if item: _process_content_item(item, pages) if hasattr(item, 'children') and item.children: _extract_elements_recursive(document, item.children, pages) def _get_element_type(item) -> str: """Determine the element type from a Docling document item.""" type_name = type(item).__name__.lower() if 'table' in type_name: return 'table' if 'picture' in type_name or 'image' in type_name or 'figure' in type_name: return 'picture' if 'section' in type_name or 'heading' in type_name: return 'section_header' if 'list' in type_name: return 'list' if 'formula' in type_name or 'equation' in type_name: return 'formula' if 'caption' in type_name: return 'caption' if hasattr(item, 'label'): label = str(item.label).lower() if 'table' in label: return 'table' if 'picture' in label or 'figure' in label: return 'picture' if 'section' in label or 'head' in label: return 'section_header' if 'list' in label: return 'list' if 'formula' in label: return 'formula' return 'text' # --- Endpoints --- @app.post("/parse", response_model=ParseResponse) async def parse(file: UploadFile): if not file.filename: raise HTTPException(status_code=400, detail="No filename provided") content = await file.read() if len(content) > MAX_FILE_SIZE: raise HTTPException(status_code=413, detail="File too large (max 50MB)") suffix = Path(file.filename).suffix tmp_path = None try: with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp: tmp.write(content) tmp_path = tmp.name result = converter.convert(tmp_path) doc = result.document content_markdown = doc.export_to_markdown() content_html = doc.export_to_html() if hasattr(doc, 'export_to_html') else "" page_count = 0 if hasattr(doc, 'pages') and doc.pages: page_count = len(doc.pages) pages_detail = extract_pages_detail(result) if not pages_detail and page_count > 0: pages_detail = [ PageDetail(page_number=i + 1, width=612.0, height=792.0, elements=[]) for i in range(page_count) ] return ParseResponse( filename=file.filename, page_count=page_count or len(pages_detail) or 1, content_markdown=content_markdown, content_html=content_html, pages=pages_detail, ) except Exception as e: raise HTTPException(status_code=422, detail=f"Failed to parse document: {str(e)}") finally: if tmp_path and os.path.exists(tmp_path): os.unlink(tmp_path) @app.post("/preview") async def preview( file: UploadFile, page: int = Query(1, ge=1), dpi: int = Query(150, ge=72, le=300), ): """Generate a PNG preview of a specific page.""" if not file.filename: raise HTTPException(status_code=400, detail="No filename provided") content = await file.read() if len(content) > MAX_FILE_SIZE: raise HTTPException(status_code=413, detail="File too large (max 50MB)") try: images = convert_from_bytes(content, first_page=page, last_page=page, dpi=dpi) if not images: raise HTTPException(status_code=404, detail=f"Page {page} not found") buf = io.BytesIO() images[0].save(buf, format="PNG") buf.seek(0) return StreamingResponse(buf, media_type="image/png") except Exception as e: raise HTTPException(status_code=422, detail=f"Failed to generate preview: {str(e)}") @app.get("/health") def health(): return {"status": "ok"}