from __future__ import annotations import html import json from dataclasses import asdict from typing import Any import gradio as gr import pandas as pd from PIL import Image from .config import MODEL_CHOICES, SETTINGS from .detector import BalloonCandidate, annotate_preview from .document import create_sample_drawing, load_document from .exporting import ( create_export_files, ensure_dataframe, records_to_dataframe, ) from .pipeline import run_extraction from .schema import ExtractionResult GRADIO_PROGRESS = gr.Progress(track_tqdm=False) def load_sample_ui() -> tuple[str, Image.Image, str]: path = create_sample_drawing() image = Image.open(path).convert("RGB") return path, image, "Sample drawing loaded. Select **Demo** or an AI model, then extract." def preview_file_ui(file_path: str | None) -> tuple[Image.Image | None, str]: if not file_path: return None, "Upload a ballooned PDF or image to begin." try: document = load_document(file_path, dpi=150, max_pages=SETTINGS.max_pages) return document.pages[0], f"Ready · {len(document.pages)} page(s)" except Exception as exc: return None, f"Could not preview: {html.escape(str(exc))}" def extract_ui( file_path: str | None, model_display: str, dpi: int, max_pages: int, tile_size: int, overlap_percent: int, progress: gr.Progress = GRADIO_PROGRESS, ) -> tuple[ Image.Image, pd.DataFrame, str, str, str, str, dict[str, Any], gr.Slider, ]: if not file_path: raise gr.Error("Upload a ballooned drawing first.", duration=8) model_key = MODEL_CHOICES.get(model_display) if model_key is None: raise gr.Error("Select a supported model mode.", duration=8) def notify(fraction: float, message: str) -> None: progress(fraction, desc=message) try: output = run_extraction( str(file_path), model_key=model_key, dpi=int(dpi), max_pages=int(max_pages), tile_size=int(tile_size), overlap=float(overlap_percent) / 100, progress=notify, ) except Exception as exc: raise gr.Error( f"Extraction stopped: {exc}", duration=14, print_exception=False, ) from exc result = output.result table = records_to_dataframe(result.characteristics) preview = annotate_preview( output.pages[0], result.characteristics, output.balloon_candidates[0], page_number=1, ) state = { "result": result.model_dump(mode="json"), "pages": output.pages, "candidates": [ [asdict(candidate) for candidate in page_candidates] for page_candidates in output.balloon_candidates ], } return ( preview, table, summary_html(table, result.warnings), metadata_markdown(result), notes_markdown(result), json.dumps(result.raw_outputs, ensure_ascii=False, indent=2), state, gr.Slider( minimum=1, maximum=max(1, result.page_count), value=1, step=1, visible=result.page_count > 1, label="Page", ), ) def page_ui(page: int, state: dict[str, Any] | None) -> Image.Image | None: if not state: return None page_number = max(1, min(int(page), len(state["pages"]))) result = ExtractionResult.model_validate(state["result"]) candidates = [ BalloonCandidate(**item) for item in state["candidates"][page_number - 1] ] return annotate_preview( state["pages"][page_number - 1], result.characteristics, candidates, page_number=page_number, ) def select_row_ui( state: dict[str, Any] | None, table: pd.DataFrame | list[list[Any]], current_page: int, event: gr.SelectData, ) -> tuple[Image.Image | None, gr.Slider]: if not state or event.index is None: return page_ui(current_page, state), gr.Slider() row_index = event.index[0] if isinstance(event.index, (tuple, list)) else int(event.index) dataframe = ensure_dataframe(table) if row_index >= len(dataframe): return page_ui(current_page, state), gr.Slider() row = dataframe.iloc[row_index] page_number = int(float(row.get("Page", current_page) or current_page)) balloon = str(row.get("Balloon", "")) result = ExtractionResult.model_validate(state["result"]) candidates = [ BalloonCandidate(**item) for item in state["candidates"][page_number - 1] ] preview = annotate_preview( state["pages"][page_number - 1], result.characteristics, candidates, selected_balloon=balloon, page_number=page_number, ) return preview, gr.Slider(value=page_number) def summary_ui( table: pd.DataFrame | list[list[Any]], state: dict[str, Any] | None ) -> str: warnings: list[str] = [] if state: warnings = ExtractionResult.model_validate(state["result"]).warnings return summary_html(ensure_dataframe(table), warnings) def export_excel_ui( table: pd.DataFrame | list[list[Any]], state: dict[str, Any] | None ) -> str: if not state: raise gr.Error("Run an extraction before exporting.") result = ExtractionResult.model_validate(state["result"]) excel, _ = create_export_files(table, result) return excel def export_json_ui( table: pd.DataFrame | list[list[Any]], state: dict[str, Any] | None ) -> str: if not state: raise gr.Error("Run an extraction before exporting.") result = ExtractionResult.model_validate(state["result"]) _, json_file = create_export_files(table, result) return json_file def summary_html(dataframe: pd.DataFrame, warnings: list[str]) -> str: extracted = len(dataframe) statuses = dataframe.get("Status", pd.Series(dtype=str)).astype(str) needs_review = int(statuses.isin(["Needs review", "Needs mapping", "nan", ""]).sum()) verified = int(statuses.isin(["Verified", "Corrected"]).sum()) warning_text = f" · {len(warnings)} warning(s)" if warnings else "" return ( '