Spaces:
Runtime error
Runtime error
| from __future__ import annotations | |
| import json | |
| import math | |
| import tempfile | |
| from pathlib import Path | |
| from typing import Any | |
| import pandas as pd | |
| from openpyxl import load_workbook | |
| from openpyxl.formatting.rule import FormulaRule | |
| from openpyxl.styles import Alignment, Border, Font, PatternFill, Side | |
| from openpyxl.worksheet.datavalidation import DataValidation | |
| from openpyxl.worksheet.table import Table, TableStyleInfo | |
| from .schema import Characteristic, ExtractionResult | |
| TABLE_COLUMNS: list[str] = [ | |
| "Balloon", | |
| "Page", | |
| "Zone", | |
| "Category", | |
| "Feature", | |
| "Requirement", | |
| "Nominal", | |
| "+Tolerance", | |
| "-Tolerance", | |
| "Upper limit", | |
| "Lower limit", | |
| "Quantity", | |
| "Unit", | |
| "Modifier", | |
| "GD&T symbol", | |
| "GD&T tolerance", | |
| "Primary datum", | |
| "Secondary datum", | |
| "Tertiary datum", | |
| "Thread", | |
| "Surface finish", | |
| "Source text", | |
| "Confidence", | |
| "Status", | |
| "Validation issues", | |
| "Record ID", | |
| ] | |
| def records_to_dataframe(records: list[Characteristic]) -> pd.DataFrame: | |
| rows: list[dict[str, Any]] = [] | |
| for item in records: | |
| rows.append( | |
| { | |
| "Balloon": item.balloon_id, | |
| "Page": item.page, | |
| "Zone": item.zone, | |
| "Category": item.category.value, | |
| "Feature": item.feature, | |
| "Requirement": item.requirement, | |
| "Nominal": item.nominal, | |
| "+Tolerance": item.upper_tolerance, | |
| "-Tolerance": item.lower_tolerance, | |
| "Upper limit": item.upper_limit, | |
| "Lower limit": item.lower_limit, | |
| "Quantity": item.quantity if item.quantity is not None else "", | |
| "Unit": item.unit, | |
| "Modifier": item.modifier, | |
| "GD&T symbol": item.gdandt_symbol, | |
| "GD&T tolerance": item.gdandt_tolerance, | |
| "Primary datum": item.datum_primary, | |
| "Secondary datum": item.datum_secondary, | |
| "Tertiary datum": item.datum_tertiary, | |
| "Thread": item.thread, | |
| "Surface finish": item.surface_finish, | |
| "Source text": item.source_text, | |
| "Confidence": round(item.confidence, 3), | |
| "Status": item.status.value, | |
| "Validation issues": " | ".join(item.validation_issues), | |
| "Record ID": item.record_id, | |
| } | |
| ) | |
| return pd.DataFrame(rows, columns=TABLE_COLUMNS) | |
| def dataframe_to_characteristics(data: pd.DataFrame | list[list[Any]]) -> list[dict[str, Any]]: | |
| dataframe = ensure_dataframe(data) | |
| records: list[dict[str, Any]] = [] | |
| for row in dataframe.to_dict(orient="records"): | |
| records.append( | |
| { | |
| "balloon_id": _clean(row.get("Balloon")), | |
| "page": _integer(row.get("Page"), 1), | |
| "zone": _clean(row.get("Zone")), | |
| "category": _clean(row.get("Category")), | |
| "feature": _clean(row.get("Feature")), | |
| "requirement": _clean(row.get("Requirement")), | |
| "nominal": _clean(row.get("Nominal")), | |
| "upper_tolerance": _clean(row.get("+Tolerance")), | |
| "lower_tolerance": _clean(row.get("-Tolerance")), | |
| "upper_limit": _clean(row.get("Upper limit")), | |
| "lower_limit": _clean(row.get("Lower limit")), | |
| "quantity": _integer_or_blank(row.get("Quantity")), | |
| "unit": _clean(row.get("Unit")), | |
| "modifier": _clean(row.get("Modifier")), | |
| "gdandt_symbol": _clean(row.get("GD&T symbol")), | |
| "gdandt_tolerance": _clean(row.get("GD&T tolerance")), | |
| "datum_primary": _clean(row.get("Primary datum")), | |
| "datum_secondary": _clean(row.get("Secondary datum")), | |
| "datum_tertiary": _clean(row.get("Tertiary datum")), | |
| "thread": _clean(row.get("Thread")), | |
| "surface_finish": _clean(row.get("Surface finish")), | |
| "source_text": _clean(row.get("Source text")), | |
| "confidence": _float(row.get("Confidence")), | |
| "status": _clean(row.get("Status")), | |
| "validation_issues": [ | |
| part.strip() | |
| for part in _clean(row.get("Validation issues")).split("|") | |
| if part.strip() | |
| ], | |
| "record_id": _clean(row.get("Record ID")), | |
| } | |
| ) | |
| return records | |
| def ensure_dataframe(data: pd.DataFrame | list[list[Any]] | None) -> pd.DataFrame: | |
| if isinstance(data, pd.DataFrame): | |
| dataframe = data.copy() | |
| elif data is None: | |
| dataframe = pd.DataFrame(columns=TABLE_COLUMNS) | |
| else: | |
| dataframe = pd.DataFrame(data, columns=TABLE_COLUMNS[: len(data[0])] if data else TABLE_COLUMNS) | |
| for column in TABLE_COLUMNS: | |
| if column not in dataframe.columns: | |
| dataframe[column] = "" | |
| return dataframe[TABLE_COLUMNS] | |
| def create_export_files( | |
| table: pd.DataFrame | list[list[Any]], | |
| result: ExtractionResult, | |
| output_dir: str | Path | None = None, | |
| ) -> tuple[str, str]: | |
| directory = Path(output_dir or tempfile.mkdtemp(prefix="drawing-intelligence-export-")) | |
| directory.mkdir(parents=True, exist_ok=True) | |
| stem = _safe_stem(Path(result.source_name).stem or "drawing") | |
| excel_path = directory / f"{stem}_characteristics.xlsx" | |
| json_path = directory / f"{stem}_characteristics.json" | |
| dataframe = ensure_dataframe(table) | |
| metadata = pd.DataFrame( | |
| [(key, _clean(value)) for key, value in result.drawing_metadata.items()], | |
| columns=["Field", "Value"], | |
| ) | |
| notes = pd.DataFrame( | |
| [(index + 1, note) for index, note in enumerate(result.general_notes)], | |
| columns=["No.", "General note"], | |
| ) | |
| log = pd.DataFrame( | |
| [ | |
| { | |
| "Source": result.source_name, | |
| "Model": result.model_id, | |
| "Pages": result.page_count, | |
| "Warnings": " | ".join(result.warnings), | |
| "Approval": "AI output requires drawing-authorized human verification", | |
| } | |
| ] | |
| ) | |
| with pd.ExcelWriter(excel_path, engine="openpyxl") as writer: | |
| _excel_safe_dataframe(dataframe).to_excel( | |
| writer, sheet_name="Characteristics", index=False | |
| ) | |
| _excel_safe_dataframe(metadata).to_excel( | |
| writer, sheet_name="Drawing metadata", index=False | |
| ) | |
| _excel_safe_dataframe(notes).to_excel( | |
| writer, sheet_name="General notes", index=False | |
| ) | |
| _excel_safe_dataframe(log).to_excel( | |
| writer, sheet_name="Extraction log", index=False | |
| ) | |
| _style_workbook(excel_path) | |
| payload = { | |
| "source_name": result.source_name, | |
| "model_id": result.model_id, | |
| "page_count": result.page_count, | |
| "drawing_metadata": result.drawing_metadata, | |
| "general_notes": result.general_notes, | |
| "warnings": result.warnings, | |
| "approval_notice": "AI output requires drawing-authorized human verification.", | |
| "characteristics": dataframe_to_characteristics(dataframe), | |
| "raw_model_outputs": result.raw_outputs, | |
| } | |
| json_path.write_text( | |
| json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8" | |
| ) | |
| return str(excel_path), str(json_path) | |
| def _style_workbook(path: Path) -> None: | |
| workbook = load_workbook(path) | |
| header_fill = PatternFill("solid", fgColor="202832") | |
| header_font = Font(color="FFFFFF", bold=True) | |
| thin = Side(style="thin", color="D8DDE5") | |
| for worksheet in workbook.worksheets: | |
| worksheet.freeze_panes = "A2" | |
| worksheet.auto_filter.ref = worksheet.dimensions | |
| for cell in worksheet[1]: | |
| cell.fill = header_fill | |
| cell.font = header_font | |
| cell.alignment = Alignment(vertical="center") | |
| worksheet.row_dimensions[1].height = 28 | |
| for row in worksheet.iter_rows(): | |
| for cell in row: | |
| cell.alignment = Alignment(vertical="top", wrap_text=True) | |
| cell.border = Border(bottom=thin) | |
| for column_cells in worksheet.columns: | |
| letter = column_cells[0].column_letter | |
| maximum = max(len(str(cell.value or "")) for cell in column_cells[:100]) | |
| worksheet.column_dimensions[letter].width = min(44, max(10, maximum + 2)) | |
| sheet = workbook["Characteristics"] | |
| if sheet.max_row >= 2: | |
| table = Table(displayName="CharacteristicsTable", ref=sheet.dimensions) | |
| table.tableStyleInfo = TableStyleInfo( | |
| name="TableStyleMedium2", | |
| showFirstColumn=False, | |
| showLastColumn=False, | |
| showRowStripes=True, | |
| showColumnStripes=False, | |
| ) | |
| sheet.add_table(table) | |
| status_col = TABLE_COLUMNS.index("Status") + 1 | |
| status_letter = sheet.cell(1, status_col).column_letter | |
| validation = DataValidation( | |
| type="list", | |
| formula1='"Needs review,Verified,Corrected,Rejected,Needs mapping"', | |
| allow_blank=False, | |
| ) | |
| sheet.add_data_validation(validation) | |
| validation.add(f"{status_letter}2:{status_letter}{sheet.max_row}") | |
| fills = { | |
| "Verified": "DDF3E7", | |
| "Corrected": "E3EEF9", | |
| "Needs review": "FFF0D1", | |
| "Needs mapping": "FDE3DF", | |
| "Rejected": "F8D7DA", | |
| } | |
| for status, color in fills.items(): | |
| sheet.conditional_formatting.add( | |
| f"{status_letter}2:{status_letter}{sheet.max_row}", | |
| FormulaRule( | |
| formula=[f'${status_letter}2="{status}"'], | |
| fill=PatternFill("solid", fgColor=color), | |
| ), | |
| ) | |
| confidence_col = TABLE_COLUMNS.index("Confidence") + 1 | |
| for row in range(2, sheet.max_row + 1): | |
| sheet.cell(row, confidence_col).number_format = "0.0%" | |
| workbook.save(path) | |
| def _excel_safe_dataframe(dataframe: pd.DataFrame) -> pd.DataFrame: | |
| safe = dataframe.copy() | |
| for column in safe.columns: | |
| safe[column] = safe[column].map(_excel_safe_value) | |
| return safe | |
| def _excel_safe_value(value: Any) -> Any: | |
| if isinstance(value, str) and value.startswith(("=", "+", "-", "@")): | |
| return "'" + value | |
| return value | |
| def _safe_stem(value: str) -> str: | |
| safe = "".join(character if character.isalnum() or character in "-_" else "_" for character in value) | |
| return safe.strip("_")[:80] or "drawing" | |
| def _clean(value: Any) -> str: | |
| if value is None or (isinstance(value, float) and math.isnan(value)): | |
| return "" | |
| if isinstance(value, (list, dict)): | |
| return json.dumps(value, ensure_ascii=False) | |
| return str(value).strip() | |
| def _integer(value: Any, default: int) -> int: | |
| try: | |
| return int(float(value)) | |
| except (TypeError, ValueError): | |
| return default | |
| def _integer_or_blank(value: Any) -> int | None: | |
| if _clean(value) == "": | |
| return None | |
| return _integer(value, 0) | |
| def _float(value: Any) -> float: | |
| try: | |
| return float(value) | |
| except (TypeError, ValueError): | |
| return 0.0 | |