unstructured-io--unstructured/unstructured/staging/base.py

from __future__ import annotations

import base64
import csv
import io
import json
import re
import zlib
from copy import deepcopy
from datetime import datetime
from typing import Any, Iterable, Optional, Sequence, cast

from unstructured.documents.coordinates import PixelSpace
from unstructured.documents.elements import (
    TYPE_TO_TEXT_ELEMENT_MAP,
    CheckBox,
    Element,
    ElementMetadata,
    Formula,
    Image,
    Table,
    Title,
)
from unstructured.errors import DecompressedSizeExceededError
from unstructured.file_utils.ndjson import dumps as ndjson_dumps
from unstructured.partition.common.common import exactly_one
from unstructured.utils import Point, dependency_exists, requires_dependencies

if dependency_exists("pandas"):
    import pandas as pd

# ================================================================================================
# SERIALIZATION/DESERIALIZATION (SERDE) RELATED FUNCTIONS
# ================================================================================================
# These serde functions will likely relocate to `unstructured.documents.elements` since they are
# so closely related to elements and this staging "brick" is deprecated.
# ================================================================================================

# == DESERIALIZERS ===============================

MAX_DECOMPRESSED_SIZE = 200 * 1024 * 1024  # 200MB

FORMULA_MARKDOWN_AUTO = "auto"
FORMULA_MARKDOWN_DISPLAY_MATH = "display_math"
FORMULA_MARKDOWN_PLAIN = "plain"
_FORMULA_MARKDOWN_STYLES = frozenset(
    {FORMULA_MARKDOWN_AUTO, FORMULA_MARKDOWN_DISPLAY_MATH, FORMULA_MARKDOWN_PLAIN},
)

# Long OCR-heavy captions often masquerade as Formula; require strong LaTeX-like signals to wrap.
_FORMULA_PROSE_HINT = re.compile(
    r"\b(was|were|using|calculated|where|respectively|determined|following)\b",
    re.IGNORECASE,
)


def _normalize_formula_for_markdown(text: str) -> str:
    """Normalize common Unicode math glyphs to LaTeX-friendly tokens.

    This is intentionally conservative and only handles symbols that are very likely to be
    interpreted as math operators/relations. The Unicode square root sign (``√``) is not
    rewritten: mapping it to ``\\sqrt{}`` would require reparsing the radicand and could
    corrupt expressions like ``√2`` or ``√(x+1)``.
    """
    # Use `{}` after each LaTeX command so the next character cannot fuse into the name
    # (e.g. x∈S -> x\in{}S, not x\inS).
    substitutions = {
        "−": "-",  # Unicode minus -> ASCII hyphen-minus
        "×": r"\times{}",
        "÷": r"\div{}",
        "∞": r"\infty{}",
        "∈": r"\in{}",
        "∉": r"\notin{}",
        "≤": r"\leq{}",
        "≥": r"\geq{}",
        "≈": r"\approx{}",
        "≠": r"\neq{}",
    }
    normalized = text
    for source, target in substitutions.items():
        normalized = normalized.replace(source, target)
    return normalized


def _formula_has_unsafe_markdown_delimiters(text: str) -> bool:
    """True if wrapping in ``$$`` could break Markdown structure or confuse math renderers."""
    return "$" in text


def _formula_math_signal_score(text: str) -> int:
    """Rough score of how much the string looks like notation (not prose OCR)."""
    score = 0
    if re.search(r"\\[a-zA-Z]+", text):
        score += 3
    if "^" in text:
        score += 1
    if re.search(r"_(\{|[0-9A-Za-z])", text):
        score += 1
    rel_sym_count = len(re.findall(r"[∈∉≤≥≠≈×÷∞∑∫√∂∇]", text))
    score += min(rel_sym_count * 2, 6)
    if re.search(r"[¼½¾]", text):
        score += 1
    equals_like = len(
        re.findall(
            r"(?<=[A-Za-z0-9\)\]])\s*=\s*(?=[A-Za-z0-9\(\\])",
            text,
        ),
    )
    score += min(equals_like, 2)
    # e.g. FFN(x) = max(...); require "(" to be function-like, not "(0) =" from OCR tables
    if re.search(r"(?<=[A-Za-z])\([^)]*\)\s*=\s*", text):
        score += 2
    return score


def _formula_looks_like_prose_sentence(text: str) -> bool:
    return len(text) >= 80 and _FORMULA_PROSE_HINT.search(text) is not None


def _formula_auto_use_display_math(text: str) -> bool:
    if _formula_looks_like_prose_sentence(text):
        return _formula_math_signal_score(text) >= 3
    return _formula_math_signal_score(text) >= 2


def _emit_formula_markdown(
    raw_text: str,
    *,
    normalize_formula: bool,
    formula_markdown_style: str,
) -> str:
    """Serialize Formula text for Markdown.

    Heuristic scoring for ``auto`` runs on **raw** stripped text so Unicode symbols
    are not replaced by ``\\command`` before the score is computed. Normalization
    applies only to text emitted inside ``$$`` blocks. ``plain`` never normalizes.
    """
    raw = raw_text.strip()
    if not raw:
        return raw

    style = formula_markdown_style.strip().lower()
    if style not in _FORMULA_MARKDOWN_STYLES:
        raise ValueError(
            "formula_markdown_style must be one of "
            f"{sorted(_FORMULA_MARKDOWN_STYLES)!r}, got {formula_markdown_style!r}",
        )
    if style == FORMULA_MARKDOWN_PLAIN:
        return raw

    if _formula_has_unsafe_markdown_delimiters(raw):
        return raw

    use_display_math = False
    if style == FORMULA_MARKDOWN_DISPLAY_MATH:
        use_display_math = True
    elif style == FORMULA_MARKDOWN_AUTO:
        use_display_math = _formula_auto_use_display_math(raw)

    if not use_display_math:
        return raw

    body = _normalize_formula_for_markdown(raw) if normalize_formula else raw
    return f"$$\n{body}\n$$"


def elements_from_base64_gzipped_json(b64_encoded_elements: str) -> list[Element]:
    """Restore Base64-encoded gzipped JSON elements to element objects.

    This is used to when deserializing `ElementMetadata.orig_elements` from its compressed form in
    JSON and dict forms and perhaps for other purposes.
    """
    # -- Base64 str -> gzip-encoded (JSON) bytes --
    decoded_b64_bytes = base64.b64decode(b64_encoded_elements)
    # -- undo gzip compression --
    dobj = zlib.decompressobj()
    elements_json_bytes = dobj.decompress(decoded_b64_bytes, max_length=MAX_DECOMPRESSED_SIZE)
    # -- Check if decompression completed successfully --
    if not dobj.eof:
        # Check if we hit the size limit or if data is actually incomplete
        if len(elements_json_bytes) >= MAX_DECOMPRESSED_SIZE:
            raise DecompressedSizeExceededError(
                max_size=MAX_DECOMPRESSED_SIZE,
            )
        else:
            raise zlib.error("Incomplete or corrupted compressed data")
    # -- JSON (bytes) to JSON (str) --
    elements_json_str = elements_json_bytes.decode("utf-8")
    # -- JSON (str) -> dicts --
    element_dicts = json.loads(elements_json_str)
    # -- dicts -> elements --
    return elements_from_dicts(element_dicts)


def elements_from_dicts(element_dicts: Iterable[dict[str, Any]]) -> list[Element]:
    """Convert a list of element-dicts to a list of elements."""
    elements: list[Element] = []

    for item in element_dicts:
        element_id: str = item.get("element_id", None)
        metadata = (
            ElementMetadata()
            if item.get("metadata") is None
            else ElementMetadata.from_dict(item["metadata"])
        )

        if item.get("type") in TYPE_TO_TEXT_ELEMENT_MAP:
            ElementCls = TYPE_TO_TEXT_ELEMENT_MAP[item["type"]]
            elements.append(ElementCls(text=item["text"], element_id=element_id, metadata=metadata))
        elif item.get("type") == "CheckBox":
            elements.append(
                CheckBox(checked=item["checked"], element_id=element_id, metadata=metadata)
            )

    return elements


# -- legacy aliases for elements_from_dicts() --
isd_to_elements = elements_from_dicts
dict_to_elements = elements_from_dicts


def elements_from_json(
    filename: str = "", text: str = "", encoding: str = "utf-8"
) -> list[Element]:
    """Loads a list of elements from a JSON file or a string."""
    exactly_one(filename=filename, text=text)

    if filename:
        with open(filename, encoding=encoding) as f:
            element_dicts = json.load(f)
    else:
        element_dicts = json.loads(text)

    return elements_from_dicts(element_dicts)


# == SERIALIZERS =================================


def elements_to_base64_gzipped_json(elements: Iterable[Element]) -> str:
    """Convert `elements` to Base64-encoded gzipped JSON.

    This is used to when serializing `ElementMetadata.orig_elements` to make it as compact as
    possible when transported as JSON, for example in an HTTP response. This compressed form is also
    present when elements are in dict form ("element_dicts"). This function is not coupled to that
    purpose however and could have other uses.
    """
    # -- adjust floating-point precision of coordinates down for a more compact str value --
    precision_adjusted_elements = _fix_metadata_field_precision(elements)
    # -- serialize elements as dicts --
    element_dicts = elements_to_dicts(precision_adjusted_elements)
    # -- serialize the dicts to JSON (bytes) --
    json_bytes = json.dumps(element_dicts, sort_keys=True).encode("utf-8")
    # -- compress the JSON bytes with gzip compression --
    deflated_bytes = zlib.compress(json_bytes)
    # -- base64-encode those bytes so they can be serialized as a JSON string value --
    b64_deflated_bytes = base64.b64encode(deflated_bytes)
    # -- convert to a string suitable for serializing in JSON --
    return b64_deflated_bytes.decode("utf-8")


def elements_to_dicts(elements: Iterable[Element]) -> list[dict[str, Any]]:
    """Convert document elements to element-dicts."""
    return [e.to_dict() for e in elements]


# -- legacy aliases for elements_to_dicts() --
convert_to_isd = elements_to_dicts
convert_to_dict = elements_to_dicts


def element_to_md(
    element: Element,
    exclude_binary_image_data: bool = False,
    normalize_formula: bool = True,
    *,
    formula_markdown_style: str = FORMULA_MARKDOWN_AUTO,
) -> str:
    match element:
        case Title(text=text):
            return f"# {text}"
        case Formula(text=text):
            return _emit_formula_markdown(
                text,
                normalize_formula=normalize_formula,
                formula_markdown_style=formula_markdown_style,
            )
        case Table(metadata=metadata, text=text) if metadata.text_as_html is not None:
            return metadata.text_as_html
        case Image(metadata=metadata, text=text) if (
            metadata.image_base64 is not None
            and metadata.image_mime_type is None
            and not exclude_binary_image_data
        ):
            return f"![{text}](data:image/*;base64,{metadata.image_base64})"
        case Image(metadata=metadata, text=text) if (
            metadata.image_base64 is not None and not exclude_binary_image_data
        ):
            return f"![{text}](data:{metadata.image_mime_type};base64,{metadata.image_base64})"
        case Image(metadata=metadata, text=text) if metadata.image_url is not None:
            return f"![{text}]({metadata.image_url})"
        case _:
            return element.text


def elements_to_md(
    elements: Iterable[Element],
    filename: Optional[str] = None,
    exclude_binary_image_data: bool = False,
    encoding: str = "utf-8",
    normalize_formula: bool = True,
    *,
    formula_markdown_style: str = FORMULA_MARKDOWN_AUTO,
) -> str:
    """Convert elements to markdown format.

    Args:
        elements: Iterable of elements to convert
        filename: Optional file path to write the markdown to
        exclude_binary_image_data: If True, exclude base64 image data from output
        encoding: File encoding when writing to file
        normalize_formula: If True, map common Unicode math symbols to LaTeX-like tokens
            for `Formula` elements before wrapping with `$$ ... $$`. Placed after ``encoding``
            so legacy positional calls ``(..., filename, exclude_binary, encoding)`` remain valid.
        formula_markdown_style: How to serialize ``Formula`` elements: ``"auto"`` (default)
            uses display math only when content looks like notation and has no ``$`` delimiters;
            ``"display_math"`` always uses ``$$`` when safe; ``"plain"`` emits text only.
            Keyword-only so positional encoding arguments stay backward compatible.

    Returns:
        The markdown content as a string
    """
    markdown_content = "\n".join(
        [
            element_to_md(
                el,
                exclude_binary_image_data=exclude_binary_image_data,
                normalize_formula=normalize_formula,
                formula_markdown_style=formula_markdown_style,
            )
            for el in elements
        ]
    )

    if filename is not None:
        with open(filename, "w", encoding=encoding) as f:
            f.write(markdown_content)

    return markdown_content


def create_file_from_elements(
    elements: Iterable[Element],
    output_format: str = "markdown",
    filename: Optional[str] = None,
    encoding: str = "utf-8",
    exclude_binary_image_data: bool = False,
    no_group_by_page: bool = True,
    normalize_formula: bool = True,
    *,
    formula_markdown_style: str = FORMULA_MARKDOWN_AUTO,
) -> str:
    """Re-create a document file from a list of elements (reverse of partition).

    Use this after partitioning a document, optionally modifying elements (e.g. replacing
    Image elements with NarrativeText using alt text), then writing back to a file.

    Supported formats: "markdown", "html", "text".

    Args:
        elements: Iterable of elements to convert (e.g. from partition_* or after editing).
        output_format: Output format: "markdown", "html", or "text".
        filename: Optional path to write the document to.
        encoding: File encoding when writing to file (all formats).
        exclude_binary_image_data: If True, omit base64 image data. Applies only to
            **markdown** and **html**; ignored for text.
        no_group_by_page: If True (default), include all elements in output. If False,
            group **html** by page (elements without metadata.page_number are skipped).
            Applies only to **html**; ignored for markdown and text. Placed before
            ``normalize_formula`` so legacy positional calls through ``exclude_binary_image_data``
            and ``no_group_by_page`` remain valid.
        normalize_formula: If True, map common Unicode math symbols to LaTeX-like tokens
            for `Formula` elements in **markdown** output. Ignored for html/text.
        formula_markdown_style: Passed to ``elements_to_md`` for **markdown** only.
            Keyword-only; see ``elements_to_md``.

    Returns:
        The document content as a string.

    Example:
        >>> from unstructured.partition.md import partition_md
        >>> from unstructured.staging.base import create_file_from_elements
        >>> elements = partition_md("README.md")
        >>> # ... modify elements (e.g. replace Image with NarrativeText) ...
        >>> create_file_from_elements(elements, output_format="markdown", filename="out.md")
    """
    format_lower = output_format.strip().lower()
    if format_lower not in ("markdown", "html", "text"):
        raise ValueError(
            f"Unsupported format: {output_format!r}. Supported formats: 'markdown', 'html', 'text'."
        )

    if format_lower == "markdown":
        content = elements_to_md(
            elements,
            filename=filename,
            exclude_binary_image_data=exclude_binary_image_data,
            normalize_formula=normalize_formula,
            encoding=encoding,
            formula_markdown_style=formula_markdown_style,
        )
        return content
    elif format_lower == "html":
        from unstructured.partition.html.convert import elements_to_html

        content = elements_to_html(
            list(elements),
            exclude_binary_image_data=exclude_binary_image_data,
            no_group_by_page=no_group_by_page,
        )
        if filename is not None:
            with open(filename, "w", encoding=encoding) as f:
                f.write(content)
        return content
    else:
        # text: delegate write to elements_to_text when filename is set
        content = convert_to_text(elements)
        if filename is not None:
            elements_to_text(elements, filename=filename, encoding=encoding)
        return content


def elements_to_json(
    elements: Iterable[Element],
    filename: Optional[str] = None,
    indent: int = 4,
    encoding: str = "utf-8",
) -> str:
    """Serialize `elements` to a JSON array.

    Also writes the JSON to `filename` if it is provided, encoded using `encoding`.

    The JSON is returned as a string.
    """
    # -- serialize `elements` as a JSON array (str) --
    precision_adjusted_elements = _fix_metadata_field_precision(elements)
    element_dicts = elements_to_dicts(precision_adjusted_elements)
    json_str = json.dumps(element_dicts, indent=indent, sort_keys=True)

    if filename is not None:
        with open(filename, "w", encoding=encoding) as f:
            f.write(json_str)

    return json_str


def elements_to_ndjson(
    elements: Iterable[Element],
    filename: Optional[str] = None,
    encoding: str = "utf-8",
) -> str:
    """Serialize `elements` to a JSON array.

    Also writes the JSON to `filename` if it is provided, encoded using `encoding`.

    The JSON is returned as a string.
    """
    # -- serialize `elements` as a JSON array (str) --
    precision_adjusted_elements = _fix_metadata_field_precision(elements)
    element_dicts = elements_to_dicts(precision_adjusted_elements)
    ndjson_str = ndjson_dumps(element_dicts, sort_keys=True)

    if filename is not None:
        with open(filename, "w", encoding=encoding) as f:
            f.write(ndjson_str)

    return ndjson_str


def _fix_metadata_field_precision(elements: Iterable[Element]) -> list[Element]:
    out_elements: list[Element] = []
    for element in elements:
        el = deepcopy(element)
        if el.metadata.coordinates:
            precision = 1 if isinstance(el.metadata.coordinates.system, PixelSpace) else 2
            points = el.metadata.coordinates.points
            assert points is not None
            rounded_points: list[Point] = []
            for point in points:
                x, y = point
                rounded_point = (round(x, precision), round(y, precision))
                rounded_points.append(rounded_point)
            el.metadata.coordinates.points = tuple(rounded_points)

        if el.metadata.detection_class_prob:
            el.metadata.detection_class_prob = round(el.metadata.detection_class_prob, 5)

        out_elements.append(el)

    return out_elements


# ================================================================================================


def _get_metadata_table_fieldnames() -> list[str]:
    metadata_fields = list(ElementMetadata.__annotations__.keys())
    metadata_fields.remove("coordinates")
    metadata_fields.extend(
        [
            "sender",
            "coordinates_points",
            "coordinates_system",
            "coordinates_layout_width",
            "coordinates_layout_height",
        ],
    )
    return metadata_fields


TABLE_FIELDNAMES: list[str] = [
    "type",
    "text",
    "element_id",
] + _get_metadata_table_fieldnames()


def convert_to_text(elements: Iterable[Element]) -> str:
    """Convert elements into clean, concatenated text."""
    return "\n".join([e.text for e in elements if hasattr(e, "text") and e.text])


def elements_to_text(
    elements: Iterable[Element], filename: Optional[str] = None, encoding: str = "utf-8"
) -> Optional[str]:
    """Convert text from each of `elements` into clean, concatenated text.

    Saves to a txt file if filename is specified. Otherwise, return the text of the elements as a
    string.
    """
    element_cct = convert_to_text(elements)
    if filename is not None:
        with open(filename, "w", encoding=encoding) as f:
            f.write(element_cct)
            return None
    else:
        return element_cct


def flatten_dict(
    dictionary: dict[str, Any],
    parent_key: str = "",
    separator: str = "_",
    flatten_lists: bool = False,
    remove_none: bool = False,
    keys_to_omit: Optional[Sequence[str]] = None,
) -> dict[str, Any]:
    """Flattens a nested dictionary into a single level dictionary.

    keys_to_omit is a list of keys that don't get flattened. If omitting a nested key, format as
    {parent_key}{separator}{key}. If flatten_lists is True, then lists and tuples are flattened as
    well. If remove_none is True, then None keys/values are removed from the flattened
    dictionary.
    """
    keys_to_omit = keys_to_omit if keys_to_omit else []
    flattened_dict: dict[str, Any] = {}
    for key, value in dictionary.items():
        new_key = f"{parent_key}{separator}{key}" if parent_key else key
        if new_key in keys_to_omit:
            flattened_dict[new_key] = value
        elif value is None and remove_none:
            continue
        elif isinstance(value, dict):
            value = cast("dict[str, Any]", value)
            flattened_dict.update(
                flatten_dict(
                    value, new_key, separator, flatten_lists, remove_none, keys_to_omit=keys_to_omit
                ),
            )
        elif isinstance(value, (list, tuple)) and flatten_lists:
            value = cast("list[Any] | tuple[Any]", value)
            for index, item in enumerate(value):
                flattened_dict.update(
                    flatten_dict(
                        {f"{new_key}{separator}{index}": item},
                        "",
                        separator,
                        flatten_lists,
                        remove_none,
                        keys_to_omit=keys_to_omit,
                    )
                )
        else:
            flattened_dict[new_key] = value

    return flattened_dict


def _get_table_fieldnames(rows: list[dict[str, Any]]):
    return list(TABLE_FIELDNAMES)


def convert_to_csv(elements: Iterable[Element]) -> str:
    """Convert `elements` to CSV format."""
    rows: list[dict[str, Any]] = elements_to_dicts(elements)
    table_fieldnames = _get_table_fieldnames(rows)
    # NOTE(robinson) - flatten metadata and add it to the table
    for row in rows:
        metadata = row.pop("metadata")
        for key, value in flatten_dict(metadata).items():
            if key in table_fieldnames:
                row[key] = value

        if row.get("sent_from"):
            row["sender"] = row.get("sent_from")
            if isinstance(row["sender"], list):
                row["sender"] = row["sender"][0]

    with io.StringIO() as buffer:
        csv_writer = csv.DictWriter(buffer, fieldnames=table_fieldnames)
        csv_writer.writeheader()
        csv_writer.writerows(rows)
        return buffer.getvalue()


# -- legacy alias for convert_to_csv --
convert_to_isd_csv = convert_to_csv


@requires_dependencies(["pandas"])
def get_default_pandas_dtypes() -> dict[str, Any]:
    return {
        "text": pd.StringDtype(),  # type: ignore
        "type": pd.StringDtype(),  # type: ignore
        "element_id": pd.StringDtype(),  # type: ignore
        "filename": pd.StringDtype(),  # Optional[str]  # type: ignore
        "filetype": pd.StringDtype(),  # Optional[str]  # type: ignore
        "file_directory": pd.StringDtype(),  # Optional[str]  # type: ignore
        "last_modified": pd.StringDtype(),  # Optional[str]  # type: ignore
        "attached_to_filename": pd.StringDtype(),  # Optional[str]  # type: ignore
        "parent_id": pd.StringDtype(),  # Optional[str],  # type: ignore
        "category_depth": "Int64",  # Optional[int]
        "image_path": pd.StringDtype(),  # Optional[str]  # type: ignore
        "languages": object,  # Optional[list[str]]
        "page_number": "Int64",  # Optional[int]
        "page_name": pd.StringDtype(),  # Optional[str]  # type: ignore
        "url": pd.StringDtype(),  # Optional[str]  # type: ignore
        "link_urls": pd.StringDtype(),  # Optional[str]  # type: ignore
        "link_texts": object,  # Optional[list[str]]
        "links": object,
        "sent_from": object,  # Optional[list[str]],
        "sent_to": object,  # Optional[list[str]]
        "subject": pd.StringDtype(),  # Optional[str]  # type: ignore
        "section": pd.StringDtype(),  # Optional[str]  # type: ignore
        "header_footer_type": pd.StringDtype(),  # Optional[str]  # type: ignore
        "emphasized_text_contents": object,  # Optional[list[str]]
        "emphasized_text_tags": object,  # Optional[list[str]]
        "text_as_html": pd.StringDtype(),  # Optional[str]  # type: ignore
        "max_characters": "Int64",  # Optional[int]
        "is_continuation": "boolean",  # Optional[bool]
        "num_carried_over_header_rows": "Int64",  # Optional[int]
        "detection_class_prob": float,  # Optional[float],
        "sender": pd.StringDtype(),  # type: ignore
        "coordinates_points": object,
        "coordinates_system": pd.StringDtype(),  # type: ignore
        "coordinates_layout_width": float,
        "coordinates_layout_height": float,
        "data_source_url": pd.StringDtype(),  # Optional[str]  # type: ignore
        "data_source_version": pd.StringDtype(),  # Optional[str]  # type: ignore
        "data_source_record_locator": object,
        "data_source_date_created": pd.StringDtype(),  # Optional[str]  # type: ignore
        "data_source_date_modified": pd.StringDtype(),  # Optional[str]  # type: ignore
        "data_source_date_processed": pd.StringDtype(),  # Optional[str]  # type: ignore
        "data_source_permissions_data": object,
        "embeddings": object,
    }


@requires_dependencies(["pandas"])
def convert_to_dataframe(
    elements: Iterable[Element], drop_empty_cols: bool = True, set_dtypes: bool = False
) -> "pd.DataFrame":
    """Convert `elements` to a pandas DataFrame.

    The dataframe contains the following columns:
        text: the element text
        type: the text type (NarrativeText, Title, etc)

    Output is pd.DataFrame
    """
    element_dicts = elements_to_dicts(elements)
    for d in element_dicts:
        if metadata := d.pop("metadata", None):
            d.update(flatten_dict(metadata, keys_to_omit=["data_source_record_locator"]))
    df = pd.DataFrame.from_dict(element_dicts)  # type: ignore
    if set_dtypes:
        dt = {k: v for k, v in get_default_pandas_dtypes().items() if k in df.columns}
        df = df.astype(dt)  # type: ignore
    if drop_empty_cols:
        df.dropna(axis=1, how="all", inplace=True)  # type: ignore
    return df


def filter_element_types(
    elements: Iterable[Element],
    include_element_types: Optional[Sequence[type[Element]]] = None,
    exclude_element_types: Optional[Sequence[type[Element]]] = None,
) -> list[Element]:
    """Filters document elements by element type"""
    exactly_one(
        include_element_types=include_element_types,
        exclude_element_types=exclude_element_types,
    )

    filtered_elements: list[Element] = []
    if include_element_types:
        for element in elements:
            if type(element) in include_element_types:
                filtered_elements.append(element)

        return filtered_elements

    elif exclude_element_types:
        for element in elements:
            if type(element) not in exclude_element_types:
                filtered_elements.append(element)

        return filtered_elements

    return list(elements)


def convert_to_coco(
    elements: Iterable[Element],
    dataset_description: Optional[str] = None,
    dataset_version: str = "1.0",
    contributors: tuple[str] = ("Unstructured Developers",),
) -> dict[str, Any]:
    from unstructured.documents.elements import TYPE_TO_TEXT_ELEMENT_MAP

    coco_dataset: dict[str, Any] = {}
    # Handle Info
    coco_dataset["info"] = {
        "description": (
            dataset_description
            if dataset_description
            else f"Unstructured COCO Dataset {datetime.now().strftime('%Y-%m-%d')}"
        ),
        "version": dataset_version,
        "year": datetime.now().year,
        "contributors": ",".join(contributors),
        "date_created": datetime.now().date().isoformat(),
    }
    element_dicts = elements_to_dicts(elements)
    # Handle Images
    images = [
        {
            "width": (
                el["metadata"]["coordinates"]["layout_width"]
                if el["metadata"].get("coordinates")
                else None
            ),
            "height": (
                el["metadata"]["coordinates"]["layout_height"]
                if el["metadata"].get("coordinates")
                else None
            ),
            "file_directory": el["metadata"].get("file_directory", ""),
            "file_name": el["metadata"].get("filename", ""),
            "page_number": el["metadata"].get("page_number", ""),
        }
        for el in element_dicts
    ]
    images = list({tuple(sorted(d.items())): d for d in images}.values())
    for index, d in enumerate(images):
        d["id"] = index + 1
    coco_dataset["images"] = images
    # Handle Categories
    categories = sorted(set(TYPE_TO_TEXT_ELEMENT_MAP.keys()))
    categories = [{"id": i + 1, "name": cat} for i, cat in enumerate(categories)]
    coco_dataset["categories"] = categories
    # Handle Annotations
    annotations = [
        {
            "id": el["element_id"],
            "category_id": [x["id"] for x in categories if x["name"] == el["type"]][0],
            "bbox": (
                [
                    float(el["metadata"].get("coordinates")["points"][0][0]),
                    float(el["metadata"].get("coordinates")["points"][0][1]),
                    float(
                        abs(
                            el["metadata"].get("coordinates")["points"][0][0]
                            - el["metadata"].get("coordinates")["points"][2][0]
                        )
                    ),
                    float(
                        abs(
                            el["metadata"].get("coordinates")["points"][0][1]
                            - el["metadata"].get("coordinates")["points"][1][1]
                        )
                    ),
                ]
                if el["metadata"].get("coordinates")
                else []
            ),
            "area": (
                (
                    float(
                        abs(
                            el["metadata"].get("coordinates")["points"][0][0]
                            - el["metadata"].get("coordinates")["points"][2][0]
                        )
                    )
                    * float(
                        abs(
                            el["metadata"].get("coordinates")["points"][0][1]
                            - el["metadata"].get("coordinates")["points"][1][1]
                        )
                    )
                )
                if el["metadata"].get("coordinates")
                else None
            ),
        }
        for el in element_dicts
    ]
    coco_dataset["annotations"] = annotations
    return coco_dataset