Spaces:

ngohel58
/

Drawing_Extraction

Runtime error

Upload 25 files

e425536 verified 3 months ago

10.9 kB

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 (
        '<div class="status-summary">'
        f'<span class="status-count">{extracted} extracted</span>'
        '<span class="status-separator">·</span>'
        f'<span class="status-review">{needs_review} need review</span>'
        '<span class="status-separator">·</span>'
        f'<span class="status-verified">{verified} approved by reviewer</span>'
        f'<span class="status-warning">{warning_text}</span>'
        "</div>"
    )


def metadata_markdown(result: ExtractionResult) -> str:
    lines = [
        f"- **Source:** {html.escape(result.source_name)}",
        f"- **Model:** {html.escape(result.model_id)}",
        f"- **Pages:** {result.page_count}",
    ]
    for key, value in result.drawing_metadata.items():
        title = str(key).replace("_", " ").title()
        lines.append(f"- **{html.escape(title)}:** {html.escape(str(value))}")
    if result.warnings:
        lines.append("\n**Extraction warnings**")
        lines.extend(f"- {html.escape(warning)}" for warning in result.warnings)
    return "\n".join(lines)


def notes_markdown(result: ExtractionResult) -> str:
    if not result.general_notes:
        return "No general notes were extracted."
    return "\n".join(
        f"{index}. {html.escape(note)}"
        for index, note in enumerate(result.general_notes, start=1)
    )


APP_CSS = r"""
:root {
  --di-accent: #b5121b;
  --di-accent-dark: #8f0d14;
  --di-ink: #171d26;
  --di-muted: #677283;
  --di-line: #dce1e8;
  --di-soft: #f6f8fa;
  --di-review: #b86a00;
  --di-good: #138a57;
}
.gradio-container {
  max-width: 1680px !important;
  margin: 0 auto !important;
  padding: 18px 22px 24px !important;
  background: #ffffff !important;
  color: var(--di-ink) !important;
}
.app-header { padding: 2px 2px 13px; }
.app-header h1 { margin: 0; font-size: 34px; line-height: 1.14; letter-spacing: -0.035em; color: var(--di-ink); }
.app-header p { margin: 5px 0 0; color: var(--di-muted); font-size: 16px; }
.control-strip, .workspace-shell, .bottom-strip {
  border: 1px solid var(--di-line) !important;
  border-radius: 9px !important;
  background: #fff !important;
  box-shadow: none !important;
}
.control-strip { padding: 13px 14px 5px !important; margin-bottom: 13px; }
.workspace-shell { overflow: hidden; }
.workspace-column { min-width: 0 !important; }
#drawing-viewer, #characteristics-table {
  border: 0 !important;
  border-radius: 0 !important;
  box-shadow: none !important;
}
#drawing-viewer { background: #fff; }
#drawing-viewer .image-container { min-height: 650px !important; background: #fff !important; }
#drawing-viewer img { object-fit: contain !important; }
#characteristics-table { border-left: 1px solid var(--di-line) !important; }
#characteristics-table table { font-size: 12px !important; line-height: 1.35 !important; }
#characteristics-table th { background: #f7f8fa !important; color: #303947 !important; font-weight: 650 !important; }
#characteristics-table td, #characteristics-table th { border-color: #e1e5ea !important; }
#characteristics-table .table-wrap { max-height: 685px !important; overflow: auto !important; }
.primary-action button, button.primary {
  background: var(--di-accent) !important;
  border-color: var(--di-accent) !important;
  color: #fff !important;
  min-height: 46px !important;
  font-weight: 650 !important;
}
.primary-action button:hover, button.primary:hover { background: var(--di-accent-dark) !important; }
button { box-shadow: none !important; font-weight: 600 !important; }
.bottom-strip { margin-top: 13px; padding: 9px 12px 2px !important; align-items: center; }
.status-summary { display: flex; flex-wrap: wrap; align-items: center; gap: 9px; min-height: 40px; font-size: 14px; }
.status-count { font-weight: 700; color: var(--di-ink); }
.status-separator { color: #a0a8b3; }
.status-review { color: var(--di-review); font-weight: 650; }
.status-verified { color: var(--di-good); font-weight: 650; }
.status-warning { color: var(--di-muted); }
.review-heading { padding: 10px 14px 6px; font-size: 18px; font-weight: 720; border-left: 1px solid var(--di-line); }
.assurance-note {
  margin-top: 10px; padding: 10px 12px; border-left: 3px solid #d29a3a;
  background: #fff8e9; color: #65491c; font-size: 13px; line-height: 1.45;
}
footer { display: none !important; }
@media (max-width: 980px) {
  .gradio-container { padding: 12px !important; }
  .app-header h1 { font-size: 28px; }
  #characteristics-table { border-left: 0 !important; border-top: 1px solid var(--di-line) !important; }
  .review-heading { border-left: 0; border-top: 1px solid var(--di-line); }
  #drawing-viewer .image-container { min-height: 440px !important; }
}
"""

	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 (
	'<div class="status-summary">'
	f'<span class="status-count">{extracted} extracted</span>'
	'<span class="status-separator">·</span>'
	f'<span class="status-review">{needs_review} need review</span>'
	'<span class="status-separator">·</span>'
	f'<span class="status-verified">{verified} approved by reviewer</span>'
	f'<span class="status-warning">{warning_text}</span>'
	"</div>"
	)


	def metadata_markdown(result: ExtractionResult) -> str:
	lines = [
	f"- Source: {html.escape(result.source_name)}",
	f"- Model: {html.escape(result.model_id)}",
	f"- Pages: {result.page_count}",
	]
	for key, value in result.drawing_metadata.items():
	title = str(key).replace("_", " ").title()
	lines.append(f"- {html.escape(title)}: {html.escape(str(value))}")
	if result.warnings:
	lines.append("\nExtraction warnings")
	lines.extend(f"- {html.escape(warning)}" for warning in result.warnings)
	return "\n".join(lines)


	def notes_markdown(result: ExtractionResult) -> str:
	if not result.general_notes:
	return "No general notes were extracted."
	return "\n".join(
	f"{index}. {html.escape(note)}"
	for index, note in enumerate(result.general_notes, start=1)
	)


	APP_CSS = r"""
	:root {
	--di-accent: #b5121b;
	--di-accent-dark: #8f0d14;
	--di-ink: #171d26;
	--di-muted: #677283;
	--di-line: #dce1e8;
	--di-soft: #f6f8fa;
	--di-review: #b86a00;
	--di-good: #138a57;
	}
	.gradio-container {
	max-width: 1680px !important;
	margin: 0 auto !important;
	padding: 18px 22px 24px !important;
	background: #ffffff !important;
	color: var(--di-ink) !important;
	}
	.app-header { padding: 2px 2px 13px; }
	.app-header h1 { margin: 0; font-size: 34px; line-height: 1.14; letter-spacing: -0.035em; color: var(--di-ink); }
	.app-header p { margin: 5px 0 0; color: var(--di-muted); font-size: 16px; }
	.control-strip, .workspace-shell, .bottom-strip {
	border: 1px solid var(--di-line) !important;
	border-radius: 9px !important;
	background: #fff !important;
	box-shadow: none !important;
	}
	.control-strip { padding: 13px 14px 5px !important; margin-bottom: 13px; }
	.workspace-shell { overflow: hidden; }
	.workspace-column { min-width: 0 !important; }
	#drawing-viewer, #characteristics-table {
	border: 0 !important;
	border-radius: 0 !important;
	box-shadow: none !important;
	}
	#drawing-viewer { background: #fff; }
	#drawing-viewer .image-container { min-height: 650px !important; background: #fff !important; }
	#drawing-viewer img { object-fit: contain !important; }
	#characteristics-table { border-left: 1px solid var(--di-line) !important; }
	#characteristics-table table { font-size: 12px !important; line-height: 1.35 !important; }
	#characteristics-table th { background: #f7f8fa !important; color: #303947 !important; font-weight: 650 !important; }
	#characteristics-table td, #characteristics-table th { border-color: #e1e5ea !important; }
	#characteristics-table .table-wrap { max-height: 685px !important; overflow: auto !important; }
	.primary-action button, button.primary {
	background: var(--di-accent) !important;
	border-color: var(--di-accent) !important;
	color: #fff !important;
	min-height: 46px !important;
	font-weight: 650 !important;
	}
	.primary-action button:hover, button.primary:hover { background: var(--di-accent-dark) !important; }
	button { box-shadow: none !important; font-weight: 600 !important; }
	.bottom-strip { margin-top: 13px; padding: 9px 12px 2px !important; align-items: center; }
	.status-summary { display: flex; flex-wrap: wrap; align-items: center; gap: 9px; min-height: 40px; font-size: 14px; }
	.status-count { font-weight: 700; color: var(--di-ink); }
	.status-separator { color: #a0a8b3; }
	.status-review { color: var(--di-review); font-weight: 650; }
	.status-verified { color: var(--di-good); font-weight: 650; }
	.status-warning { color: var(--di-muted); }
	.review-heading { padding: 10px 14px 6px; font-size: 18px; font-weight: 720; border-left: 1px solid var(--di-line); }
	.assurance-note {
	margin-top: 10px; padding: 10px 12px; border-left: 3px solid #d29a3a;
	background: #fff8e9; color: #65491c; font-size: 13px; line-height: 1.45;
	}
	footer { display: none !important; }
	@media (max-width: 980px) {
	.gradio-container { padding: 12px !important; }
	.app-header h1 { font-size: 28px; }
	#characteristics-table { border-left: 0 !important; border-top: 1px solid var(--di-line) !important; }
	.review-heading { border-left: 0; border-top: 1px solid var(--di-line); }
	#drawing-viewer .image-container { min-height: 440px !important; }
	}
	"""