diff --git a/.gitattributes b/.gitattributes index a6344aac8c09253b3b630fb776ae94478aa0275b..54496fa24d48b490e10b78194529503941a84932 100644 --- a/.gitattributes +++ b/.gitattributes @@ -33,3 +33,20 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text *.zip filter=lfs diff=lfs merge=lfs -text *.zst filter=lfs diff=lfs merge=lfs -text *tfevents* filter=lfs diff=lfs merge=lfs -text +examples/scheme1.png filter=lfs diff=lfs merge=lfs -text +src/editable_image/fonts/Lato-Bold.ttf filter=lfs diff=lfs merge=lfs -text +src/editable_image/fonts/Lato-BoldItalic.ttf filter=lfs diff=lfs merge=lfs -text +src/editable_image/fonts/Lato-Italic.ttf filter=lfs diff=lfs merge=lfs -text +src/editable_image/fonts/Lato-Regular.ttf filter=lfs diff=lfs merge=lfs -text +src/editable_image/fonts/LiberationMono-Bold.ttf filter=lfs diff=lfs merge=lfs -text +src/editable_image/fonts/LiberationMono-BoldItalic.ttf filter=lfs diff=lfs merge=lfs -text +src/editable_image/fonts/LiberationMono-Italic.ttf filter=lfs diff=lfs merge=lfs -text +src/editable_image/fonts/LiberationMono-Regular.ttf filter=lfs diff=lfs merge=lfs -text +src/editable_image/fonts/LiberationSans-Bold.ttf filter=lfs diff=lfs merge=lfs -text +src/editable_image/fonts/LiberationSans-BoldItalic.ttf filter=lfs diff=lfs merge=lfs -text +src/editable_image/fonts/LiberationSans-Italic.ttf filter=lfs diff=lfs merge=lfs -text +src/editable_image/fonts/LiberationSans-Regular.ttf filter=lfs diff=lfs merge=lfs -text +src/editable_image/fonts/LiberationSerif-Bold.ttf filter=lfs diff=lfs merge=lfs -text +src/editable_image/fonts/LiberationSerif-BoldItalic.ttf filter=lfs diff=lfs merge=lfs -text +src/editable_image/fonts/LiberationSerif-Italic.ttf filter=lfs diff=lfs merge=lfs -text +src/editable_image/fonts/LiberationSerif-Regular.ttf filter=lfs diff=lfs merge=lfs -text diff --git a/README.md b/README.md index ad86e0a48bd9968e6c99307b48f12bb18020b799..5c80b54c292967a77fbb473cc047708bde8a2974 100644 --- a/README.md +++ b/README.md @@ -1,14 +1,105 @@ --- -title: Make Editable Image -emoji: 🐠 -colorFrom: green -colorTo: purple +title: Editable Image +emoji: āœļø +colorFrom: blue +colorTo: indigo sdk: gradio -sdk_version: 6.26.0 -python_version: '3.12' app_file: app.py +python_version: "3.11" pinned: false -short_description: Make editable --- -Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference +# Editable Image + +Editable Image is a local Python script that removes printed text from bitmap images and reconstructs it as a portable, editable SVG text layer. It uses compact PP-OCRv6 models through RapidOCR. + +## Install + +Python 3.11 or newer and `uv` are recommended. + +```bash +uv sync --extra cpu --extra test +uv run python make_editable_image.py models install ocr +``` + +For NVIDIA OCR acceleration, use `--extra cuda` instead of `--extra cpu`. Add `--extra lama` to install optional neural inpainting support. LaMa's approximately 200 MB model is only downloaded by an explicit command: + +```bash +uv run python make_editable_image.py models install lama +``` + +## Use + +```bash +uv run python make_editable_image.py image.png --output output +uv run python make_editable_image.py images/ --recursive --output output +``` + +Each input produces `//-background.png` and `-overlay.svg`. The SVG links the sibling background PNG, so keep both files together when moving the output. Run `uv run python make_editable_image.py --help` for all options. The installed `editable-image convert ...` command is equivalent. + +### Gradio / Hugging Face Spaces + +The repository includes a Gradio interface in `app.py` with controls for every +conversion option. To run it locally: + +```bash +uv run --with-requirements requirements.txt python app.py +``` + +For Hugging Face Spaces, create a Gradio Space and upload the repository. The +README metadata selects `app.py`, and `requirements.txt` installs the package +with CPU ONNX Runtime. The default OpenCV cleanup and bitmap extraction work +without additional model downloads. LaMa and SAM 2 remain available in the UI +when their optional dependencies and model files are added to the Space. + +Add `--vectorize-shapes` to promote confident diagram primitives—rounded rectangles, circles, ellipses, simple polygons, and connector-like paths—into editable SVG elements: + +```bash +uv run python make_editable_image.py image.png --output output --vectorize-shapes +``` + +Detailed illustrations are preserved as transparent PNG crops under `-assets/` and linked from the SVG. Move the complete per-image output directory together. Shape recognition is conservative: uncertain elements remain in the background bitmap. + +Bitmap masks use OpenCV by default. SAM 2 can optionally refine difficult masks: + +```bash +uv sync --extra cpu --extra sam2 +uv run python make_editable_image.py models install sam2 +uv run python make_editable_image.py image.png --output output \ + --vectorize-shapes --bitmap-extractor sam2 +``` + +The SAM 2 checkpoint is downloaded only by the explicit install command. + +Detected text orientation snaps to `0`, `45`, `90`, `-45`, or `-90` degrees when it is within 6 degrees of a target. Override this with `--snap-angles 0,30,60,90,-30,-60,-90` and `--snap-tolerance 4`. + +Add `--snap-font-sizes` to estimate the font-size distribution across each image and snap nearby estimates to its supported peaks. Isolated sizes, such as a single large heading, remain unchanged. + +SVGs reference font family names by default to keep files small. Add `--embed-fonts` to embed the used WOFF2 font faces for portable rendering. + +## Example output + +The repository includes a shape-vectorization result for [`examples/scheme1.png`](examples/scheme1.png) under [`examples/generated-output/scheme1/`](examples/generated-output/scheme1/). The directory contains the cleared background, editable SVG overlay, linked transparent bitmap assets, and a rendered PNG preview. + +## Notes + +- The default OpenCV cleanup is fast and works best on flat or slowly varying backgrounds. +- LaMa can improve larger removals over textured backgrounds but adds PyTorch and a much larger model. +- Font recovery is approximate. The app bundles OFL-licensed Lato and Liberation faces; embedding them is optional. +- The script accepts PNG, JPEG, and WebP images up to 50 MB and 40 megapixels. +- Shape vectorization targets clean diagram geometry, not photographs, logos, or detailed illustrations. + +# Method + +The conversion is a hybrid OCR and image-processing pipeline. It does not use a large generative model for the normal workflow. + +1. **Decode the bitmap.** Pillow reads PNG, JPEG, or WebP input, applies EXIF orientation, preserves the ICC profile, and converts the pixels to an RGBA NumPy array. +2. **Detect and recognize text.** RapidOCR runs the compact PP-OCRv6-small detection and recognition models through ONNX Runtime. ONNX Runtime provides either CPU execution or optional NVIDIA CUDA acceleration. The OCR result supplies the text, confidence, and four-corner polygon for each line. +3. **Estimate text appearance.** OpenCV measures the line angle and foreground color from pixels inside each OCR polygon. Pillow renders candidates from the bundled Lato faces, and OpenCV compares their masks with the source to estimate regular or bold weight, italic style, and font size. With `--snap-font-sizes`, NumPy evaluates a Gaussian kernel density over all estimated sizes, keeps peaks supported by multiple lines, and snaps nearby sizes to those peaks. This is an approximation rather than recovery of the original font file. +4. **Normalize orientation.** The measured line angle is snapped to a configured target when it falls within the snap tolerance. The default targets are `0`, `45`, `90`, `-45`, and `-90` degrees. +5. **Protect recognized text.** The OCR polygons are converted into a protection mask before graphical analysis. This prevents letters, bullets, and text strokes from being promoted into shapes and later provides transparent text holes in extracted bitmap layers. +6. **Recognize optional diagram geometry.** With `--vectorize-shapes`, OpenCV analyzes the untouched source at two scales. Median-filtered color clustering proposes large, low-frequency panels; LAB color clusters, connected components, contours, and edge rectangles propose smaller geometry. A dedicated arrow pass groups neutral and overlapping HSV color bands, then scores concave silhouettes for a sharp extremal tip, a supported head, a tail, and approximate symmetry. Repeated-arrow context inside panels suppresses arrow-like illustration fragments, while connectors between panels are retained. The remaining fitter tests rounded rectangles, ellipses, simple polygons, and low-complexity paths. Nested candidates and textured regions are rejected conservatively so detailed illustrations remain raster. +7. **Fit and score SVG primitives.** Every proposal is rendered back into a mask and scored using geometric overlap, visible-edge support, color support, size, and containment. Accepted shapes retain both their editable geometry and original source contour. Their appearance is fitted as a solid fill or, when it materially lowers reconstruction error, a two-stop linear gradient with optional stroke styling. +8. **Extract bitmap foreground layers.** For an accepted container such as a colored panel, the fitted fill is compared with the source in LAB space. OpenCV uses the residual as foreground/background seeds for GrabCut, filters small components, and derives a feathered alpha channel. `--bitmap-extractor sam2` optionally refines that mask with locally installed SAM 2 masks. Original RGB pixels are never traced: each retained illustration is written as a transparent PNG crop. +9. **Build the cleanup mask.** Text masks and accepted source contours are combined with any manual mask strokes. Large vectorized containers are cleared by estimating the underlying low-frequency canvas color from their outer perimeter; smaller text and shape remnants use OpenCV Telea inpainting or optional LaMa cleanup. Filled arrows are exported as editable closed SVG paths with `data-role="arrow"`; centerline arrows use SVG stroke markers so their arrowheads rotate and scale with the path. +10. **Write ordered hybrid output.** Pillow saves the cleared bitmap and transparent bitmap assets as PNG. The SVG links the cleared background, emits editable primitives and bitmap crops in layer order, and places editable `` in the final foreground group. With `--embed-fonts`, FontTools converts the used bundled faces to WOFF2 and stores them inside the SVG; otherwise the SVG names the font families. diff --git a/app.py b/app.py new file mode 100644 index 0000000000000000000000000000000000000000..2fcbeb6b2e0c3058bce21fb622828c33f42cdb4e --- /dev/null +++ b/app.py @@ -0,0 +1,323 @@ +from __future__ import annotations + +import base64 +import html +import re +import tempfile +import zipfile +from pathlib import Path + +import gradio as gr + +from editable_image.models import ( + BitmapExtractor, + CleanupMode, + DEFAULT_SNAP_ANGLES, + Device, + ProcessOptions, +) +from editable_image.pipeline import ProcessResult, export_result, process_path +from editable_image.svg_export import build_svg + + +APP_TITLE = "Editable Image" +APP_DESCRIPTION = ( + "Remove printed text from a bitmap and rebuild it as an editable SVG text layer. " + "Shape vectorization can also promote confident diagram geometry to SVG elements." +) +ANGLE_DEFAULT = ", ".join(f"{angle:g}" for angle in DEFAULT_SNAP_ANGLES) +CSS = """ +.gradio-container { max-width: 1280px !important; } +.hero { margin: 0 auto 0.6rem; max-width: 900px; text-align: center; } +.hero h1 { font-size: 2.15rem; margin-bottom: 0.25rem; } +.hero p { color: var(--body-text-color-subdued); font-size: 1rem; } +.svg-preview { + align-items: center; + background-color: #f4f4f5; + background-image: linear-gradient(45deg, #e4e4e7 25%, transparent 25%), + linear-gradient(-45deg, #e4e4e7 25%, transparent 25%), + linear-gradient(45deg, transparent 75%, #e4e4e7 75%), + linear-gradient(-45deg, transparent 75%, #e4e4e7 75%); + background-position: 0 0, 0 8px, 8px -8px, -8px 0; + background-size: 16px 16px; + border: 1px solid var(--border-color-primary); + border-radius: 10px; + display: flex; + justify-content: center; + min-height: 420px; + overflow: auto; + padding: 1rem; +} +.svg-preview img { display: block; max-height: 72vh; max-width: 100%; object-fit: contain; } +""" +GRADIO_MAJOR = int(gr.__version__.split(".", maxsplit=1)[0]) + + +def parse_snap_angles(value: str) -> list[float]: + try: + angles = [float(item.strip()) for item in value.split(",") if item.strip()] + except (AttributeError, ValueError) as exc: + raise ValueError("Snap angles must be comma-separated numbers.") from exc + if not angles: + raise ValueError("Enter at least one snap angle.") + if any(angle < -180 or angle > 180 for angle in angles): + raise ValueError("Snap angles must be between -180 and 180 degrees.") + return angles + + +def safe_stem(path: Path) -> str: + stem = re.sub(r"[^A-Za-z0-9._-]+", "-", path.stem).strip("-._") + return stem[:100] or "image" + + +def file_data_uri(path: Path, media_type: str) -> str: + payload = base64.b64encode(path.read_bytes()).decode("ascii") + return f"data:{media_type};base64,{payload}" + + +def preview_html( + result: ProcessResult, + background_path: Path, + asset_paths: list[Path], + embed_fonts: bool, +) -> str: + background_uri = file_data_uri(background_path, "image/png") + asset_by_name = {path.name: path for path in asset_paths} + bitmap_hrefs: dict[str, tuple[str, str]] = {} + for region in result.document.bitmap_regions: + asset_path = asset_by_name.get(region.asset_name) + if asset_path is not None: + uri = file_data_uri(asset_path, "image/png") + bitmap_hrefs[region.id] = (uri, uri) + svg = build_svg( + result.document, + background_href=background_uri, + background_fallback_href=background_uri, + embed_fonts=embed_fonts, + bitmap_hrefs=bitmap_hrefs, + ) + encoded = base64.b64encode(svg.encode("utf-8")).decode("ascii") + alt = html.escape(f"Editable reconstruction of {result.document.source_name}", quote=True) + return ( + '
' + f'{alt}' + "
" + ) + + +def make_bundle(output_directory: Path, stem: str) -> Path: + archive = output_directory.parent / f"{stem}-editable.zip" + with zipfile.ZipFile(archive, "w", compression=zipfile.ZIP_DEFLATED) as bundle: + for path in sorted(output_directory.rglob("*")): + if path.is_file(): + bundle.write(path, arcname=path.relative_to(output_directory.parent)) + return archive + + +def convert_image( + image_path: str | None, + confidence: float, + cleanup: str, + device: str, + max_pixels: float, + snap_angles: str, + snap_tolerance: float, + snap_font_sizes: bool, + embed_fonts: bool, + vectorize_shapes: bool, + bitmap_extractor: str, +) -> tuple[str, object, object, dict[str, object], str, str]: + if not image_path: + raise gr.Error("Upload a PNG, JPEG, or WebP image first.") + + source = Path(image_path) + try: + options = ProcessOptions( + confidence=confidence, + cleanup=CleanupMode(cleanup), + device=Device(device), + max_pixels=int(max_pixels), + snap_angles=parse_snap_angles(snap_angles), + snap_tolerance=snap_tolerance, + snap_font_sizes=snap_font_sizes, + vectorize_shapes=vectorize_shapes, + bitmap_extractor=BitmapExtractor(bitmap_extractor), + ) + result = process_path(source, options) + + job_root = Path(tempfile.mkdtemp(prefix="editable-image-")) + stem = safe_stem(source) + output_directory = job_root / stem + background_path, overlay_path, asset_paths = export_result( + result, + output_directory, + stem, + embed_fonts=embed_fonts, + ) + document_path = output_directory / f"{stem}-document.json" + document_path.write_text( + result.document.model_dump_json(indent=2), + encoding="utf-8", + ) + bundle_path = make_bundle(output_directory, stem) + rendered = preview_html(result, background_path, asset_paths, embed_fonts) + except Exception as exc: + # Gradio displays this as a concise toast instead of exposing a server traceback. + raise gr.Error(str(exc) or type(exc).__name__) from exc + + metadata: dict[str, object] = { + "source": result.document.source_name, + "dimensions": f"{result.image.width} Ɨ {result.image.height}", + "recognized_text_regions": len(result.document.regions), + "vector_shapes": len(result.document.shapes), + "bitmap_assets": len(result.bitmap_assets), + "cleanup": options.cleanup.value, + "device": options.device.value, + "document_version": result.document.version, + } + return ( + rendered, + result.background, + result.mask, + metadata, + str(overlay_path), + str(bundle_path), + ) + + +def build_demo() -> gr.Blocks: + blocks_options: dict[str, object] = {"title": APP_TITLE} + if GRADIO_MAJOR < 6: + blocks_options.update(theme=gr.themes.Soft(), css=CSS) + with gr.Blocks(**blocks_options) as demo: + gr.HTML( + '
' + f"

{APP_TITLE}

{APP_DESCRIPTION}

" + "
" + ) + with gr.Row(equal_height=False): + with gr.Column(scale=5): + source = gr.Image( + label="Source image", + type="filepath", + image_mode="RGBA", + sources=["upload", "clipboard"], + height=430, + ) + convert = gr.Button("Make editable", variant="primary", size="lg") + gr.Markdown( + "PNG, JPEG, and WebP are supported. Keep the downloaded bundle intact: " + "its SVG links to the sibling background and any extracted bitmap assets." + ) + with gr.Column(scale=4): + with gr.Accordion("OCR and orientation", open=True): + confidence = gr.Slider( + 0, + 1, + value=0.5, + step=0.01, + label="Minimum OCR confidence", + ) + snap_angles = gr.Textbox( + value=ANGLE_DEFAULT, + label="Snap angles (degrees)", + info="Comma-separated targets between -180 and 180.", + ) + snap_tolerance = gr.Slider( + 0, + 180, + value=6, + step=0.5, + label="Angle snap tolerance (degrees)", + ) + snap_font_sizes = gr.Checkbox( + value=False, + label="Snap font sizes to detected distribution peaks", + ) + embed_fonts = gr.Checkbox( + value=False, + label="Embed used WOFF2 fonts in the SVG", + ) + with gr.Accordion("Cleanup and vectorization", open=True): + cleanup = gr.Dropdown( + choices=[item.value for item in CleanupMode], + value=CleanupMode.OPENCV.value, + label="Cleanup method", + info="LaMa requires the optional dependency and installed model.", + ) + vectorize_shapes = gr.Checkbox( + value=False, + label="Vectorize confident diagram shapes", + ) + bitmap_extractor = gr.Dropdown( + choices=[item.value for item in BitmapExtractor], + value=BitmapExtractor.OPENCV.value, + label="Bitmap layer extractor", + info="Used during shape vectorization. SAM 2 requires its optional model.", + ) + with gr.Accordion("Runtime and limits", open=False): + device = gr.Dropdown( + choices=[item.value for item in Device], + value=Device.AUTO.value, + label="Compute device", + ) + max_pixels = gr.Number( + value=40_000_000, + minimum=1, + maximum=200_000_000, + step=1, + precision=0, + label="Maximum decoded pixels", + ) + + gr.Markdown("## Result") + with gr.Tabs(): + with gr.Tab("Editable preview"): + rendered_preview = gr.HTML( + '
Upload an image and choose ā€œMake editableā€.
', + label="Reconstructed SVG", + ) + with gr.Tab("Cleared background"): + background = gr.Image(label="Background PNG", interactive=False) + with gr.Tab("Removal mask"): + mask = gr.Image(label="Applied mask", interactive=False) + with gr.Tab("Details"): + metadata = gr.JSON(label="Conversion details") + with gr.Row(): + overlay = gr.File(label="Editable SVG") + bundle = gr.File(label="Complete bundle (.zip)") + + controls = [ + source, + confidence, + cleanup, + device, + max_pixels, + snap_angles, + snap_tolerance, + snap_font_sizes, + embed_fonts, + vectorize_shapes, + bitmap_extractor, + ] + outputs = [rendered_preview, background, mask, metadata, overlay, bundle] + convert.click(convert_image, inputs=controls, outputs=outputs, api_name="convert") + + example = Path(__file__).resolve().parent / "examples" / "scheme1.png" + if example.exists(): + gr.Examples( + examples=[[str(example)]], + inputs=[source], + label="Example", + ) + return demo + + +demo = build_demo() + + +if __name__ == "__main__": + launch_options: dict[str, object] = {} + if GRADIO_MAJOR >= 6: + launch_options.update(theme=gr.themes.Soft(), css=CSS) + demo.queue(default_concurrency_limit=1).launch(**launch_options) diff --git a/examples/scheme1.png b/examples/scheme1.png new file mode 100644 index 0000000000000000000000000000000000000000..c1da77ed8ea978028e21b73f204d470134e1d26f --- /dev/null +++ b/examples/scheme1.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:33322a4fc618424afc3c79ec844b0ecd472d6590824f462ad99a601b18e8c6d1 +size 1432222 diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000000000000000000000000000000000000..059d09e335efe18ec8c2b243aabeff61031088ba --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,54 @@ +[build-system] +requires = ["hatchling>=1.27"] +build-backend = "hatchling.build" + +[project] +name = "editable-image" +version = "0.1.0" +description = "Remove bitmap text into an editable SVG overlay" +readme = "README.md" +requires-python = ">=3.11" +license = { text = "MIT" } +dependencies = [ + "fonttools[woff]>=4.59,<5", + "numpy>=2.0,<3", + "opencv-python>=4.10,<6", + "pillow>=11,<13", + "platformdirs>=4.3,<5", + "pydantic>=2.10,<3", + "rapidocr>=3.9,<4", +] + +[project.optional-dependencies] +cpu = ["onnxruntime>=1.20,<2"] +cuda = ["onnxruntime-gpu>=1.20,<2"] +lama = ["torch>=2.5,<3"] +sam2 = [ + "huggingface-hub>=0.34,<2", + "torch>=2.5,<3", + "torchvision>=0.20,<1", + "transformers>=4.57,<6", +] +test = [ + "pytest>=8.3,<9", + "pytest-cov>=6,<8", +] + +[project.scripts] +editable-image = "editable_image.cli:main" + +[tool.hatch.build.targets.wheel] +packages = ["src/editable_image"] + +[tool.hatch.build] +exclude = ["src/editable_image/static"] + +[tool.hatch.build.targets.wheel.force-include] +"src/editable_image/fonts" = "editable_image/fonts" + +[tool.pytest.ini_options] +addopts = "-q" +testpaths = ["tests"] + +[tool.uv] +conflicts = [[{ extra = "cpu" }, { extra = "cuda" }]] diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..30a0adab85c21397f7892f0524f790c6f30865f5 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,2 @@ +-e .[cpu] +gradio>=5,<7 diff --git a/src/editable_image/__init__.py b/src/editable_image/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..935493748b7f1a08a36110c1f5167c9583ca83a6 --- /dev/null +++ b/src/editable_image/__init__.py @@ -0,0 +1,4 @@ +"""Bitmap text extraction into editable SVG overlays.""" + +__version__ = "0.1.0" + diff --git a/src/editable_image/__pycache__/__init__.cpython-313.pyc b/src/editable_image/__pycache__/__init__.cpython-313.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5f576dab990d4cccf0842746661783514415ea85 Binary files /dev/null and b/src/editable_image/__pycache__/__init__.cpython-313.pyc differ diff --git a/src/editable_image/__pycache__/bitmap_extract.cpython-313.pyc b/src/editable_image/__pycache__/bitmap_extract.cpython-313.pyc new file mode 100644 index 0000000000000000000000000000000000000000..61a612d96cbf975f2d6eb49312f589e5ad29ea84 Binary files /dev/null and b/src/editable_image/__pycache__/bitmap_extract.cpython-313.pyc differ diff --git a/src/editable_image/__pycache__/cli.cpython-313.pyc b/src/editable_image/__pycache__/cli.cpython-313.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0a113c3d7a48084116f91c158a991ac1ba423f67 Binary files /dev/null and b/src/editable_image/__pycache__/cli.cpython-313.pyc differ diff --git a/src/editable_image/__pycache__/fonts.cpython-313.pyc b/src/editable_image/__pycache__/fonts.cpython-313.pyc new file mode 100644 index 0000000000000000000000000000000000000000..26e7bccd49e05fe55074429dec085ad7e3e59996 Binary files /dev/null and b/src/editable_image/__pycache__/fonts.cpython-313.pyc differ diff --git a/src/editable_image/__pycache__/image_io.cpython-313.pyc b/src/editable_image/__pycache__/image_io.cpython-313.pyc new file mode 100644 index 0000000000000000000000000000000000000000..475cdf89566f4dbc69a7c24602164f3cec03e507 Binary files /dev/null and b/src/editable_image/__pycache__/image_io.cpython-313.pyc differ diff --git a/src/editable_image/__pycache__/inpaint.cpython-313.pyc b/src/editable_image/__pycache__/inpaint.cpython-313.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c4ed2bf19d522dde3145d9c0b930c72b773c0392 Binary files /dev/null and b/src/editable_image/__pycache__/inpaint.cpython-313.pyc differ diff --git a/src/editable_image/__pycache__/masking.cpython-313.pyc b/src/editable_image/__pycache__/masking.cpython-313.pyc new file mode 100644 index 0000000000000000000000000000000000000000..890e2be8020781456b7bd9fac9be59e9f4ae76ef Binary files /dev/null and b/src/editable_image/__pycache__/masking.cpython-313.pyc differ diff --git a/src/editable_image/__pycache__/models.cpython-313.pyc b/src/editable_image/__pycache__/models.cpython-313.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0f1c57a458dcf105f0d11004b3dded05b5f828ac Binary files /dev/null and b/src/editable_image/__pycache__/models.cpython-313.pyc differ diff --git a/src/editable_image/__pycache__/ocr.cpython-313.pyc b/src/editable_image/__pycache__/ocr.cpython-313.pyc new file mode 100644 index 0000000000000000000000000000000000000000..402de40edaae58329d14a7800cd1d6fcbaafe985 Binary files /dev/null and b/src/editable_image/__pycache__/ocr.cpython-313.pyc differ diff --git a/src/editable_image/__pycache__/pipeline.cpython-313.pyc b/src/editable_image/__pycache__/pipeline.cpython-313.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c4984d0471d6fd91b5c92ab1b0fbeb1d7480b717 Binary files /dev/null and b/src/editable_image/__pycache__/pipeline.cpython-313.pyc differ diff --git a/src/editable_image/__pycache__/sam2.cpython-313.pyc b/src/editable_image/__pycache__/sam2.cpython-313.pyc new file mode 100644 index 0000000000000000000000000000000000000000..73fe882b68bfda5586e2e56235622eabad78441d Binary files /dev/null and b/src/editable_image/__pycache__/sam2.cpython-313.pyc differ diff --git a/src/editable_image/__pycache__/shape_recognition.cpython-313.pyc b/src/editable_image/__pycache__/shape_recognition.cpython-313.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8b60dad5e2dac6299954998a1b771c57561643ca Binary files /dev/null and b/src/editable_image/__pycache__/shape_recognition.cpython-313.pyc differ diff --git a/src/editable_image/__pycache__/style.cpython-313.pyc b/src/editable_image/__pycache__/style.cpython-313.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5bd2427abf99c086234db19637f1fab1fd1272a8 Binary files /dev/null and b/src/editable_image/__pycache__/style.cpython-313.pyc differ diff --git a/src/editable_image/__pycache__/svg_export.cpython-313.pyc b/src/editable_image/__pycache__/svg_export.cpython-313.pyc new file mode 100644 index 0000000000000000000000000000000000000000..09c57dc09f203f1d03f14782786532970cffc606 Binary files /dev/null and b/src/editable_image/__pycache__/svg_export.cpython-313.pyc differ diff --git a/src/editable_image/bitmap_extract.py b/src/editable_image/bitmap_extract.py new file mode 100644 index 0000000000000000000000000000000000000000..724bb7ddd4983c267bbffd2e7e08a7a5c8b027a5 --- /dev/null +++ b/src/editable_image/bitmap_extract.py @@ -0,0 +1,203 @@ +from __future__ import annotations + +import cv2 +import numpy as np + +from .models import ( + BitmapExtractor, + BitmapRegion, + Box, + DocumentSpec, + LinearGradientPaint, + ShapeRegion, + SolidPaint, +) +from .shape_recognition import geometry_mask, shape_source_mask + + +def extract_bitmap_layers( + rgba: np.ndarray, + document: DocumentSpec, + text_mask: np.ndarray, + extractor: BitmapExtractor, + device: str, +) -> tuple[list[BitmapRegion], dict[str, np.ndarray]]: + height, width = rgba.shape[:2] + rgb = rgba[:, :, :3] + text_exclusion = cv2.dilate(text_mask, np.ones((5, 5), np.uint8)) + for region in document.regions: + polygon = np.array( + [(point.x, point.y) for point in region.source_quad], + dtype=np.int32, + ) + cv2.fillPoly(text_exclusion, [polygon], 255) + vector_foreground = np.zeros((height, width), dtype=np.uint8) + for shape in document.shapes: + if not shape.is_container and shape.include_overlay: + vector_foreground = cv2.bitwise_or( + vector_foreground, + geometry_mask(shape.geometry, width, height), + ) + + regions: list[BitmapRegion] = [] + assets: dict[str, np.ndarray] = {} + for shape in document.shapes: + if not shape.is_container or not shape.include_overlay: + continue + container_mask = geometry_mask(shape.geometry, width, height) + predicted = render_fill(shape, width, height) + source_lab = cv2.cvtColor(rgb, cv2.COLOR_RGB2LAB).astype(np.float32) + predicted_lab = cv2.cvtColor(predicted, cv2.COLOR_RGB2LAB).astype(np.float32) + delta = np.linalg.norm(source_lab - predicted_lab, axis=2) + border_exclusion = cv2.subtract( + container_mask, + cv2.erode(container_mask, np.ones((7, 7), np.uint8)), + ) + eligible = ( + (container_mask > 0) + & (text_exclusion == 0) + & (vector_foreground == 0) + & (border_exclusion == 0) + ) + values = delta[eligible] + if values.size < 20: + continue + baseline = float(np.percentile(values, 35)) + threshold = max(9.0, baseline + 6.0) + initial = np.zeros((height, width), dtype=np.uint8) + initial[eligible & (delta >= threshold * 0.65)] = cv2.GC_PR_FGD + initial[eligible & (delta >= threshold * 1.35)] = cv2.GC_FGD + initial[(container_mask > 0) & (initial == 0)] = cv2.GC_BGD + if np.count_nonzero(initial == cv2.GC_FGD) < 8: + continue + mask = refine_opencv(rgb, initial, container_mask) + if extractor == BitmapExtractor.SAM2: + from .sam2 import refine_with_sam2 + + mask = refine_with_sam2(rgb, mask, container_mask, device) + mask[text_exclusion > 0] = 0 + mask[vector_foreground > 0] = 0 + mask = filter_components(mask, max(8, int(width * height * 0.00001))) + if cv2.countNonZero(mask) < 12: + continue + alpha = soft_alpha(delta, mask, threshold) + x, y, w, h = cv2.boundingRect((alpha > 0).astype(np.uint8)) + if w <= 0 or h <= 0: + continue + padding = 2 + x0, y0 = max(0, x - padding), max(0, y - padding) + x1, y1 = min(width, x + w + padding), min(height, y + h + padding) + crop = np.dstack([rgb[y0:y1, x0:x1], alpha[y0:y1, x0:x1]]) + asset_name = f"bitmap-{len(regions) + 1:04d}.png" + region = BitmapRegion( + id=f"bitmap-{len(regions) + 1:04d}", + asset_name=asset_name, + box=Box(x=x0, y=y0, width=x1 - x0, height=y1 - y0), + z_index=500, + ) + regions.append(region) + assets[region.id] = crop + return regions, assets + + +def refine_opencv(rgb: np.ndarray, initial: np.ndarray, container_mask: np.ndarray) -> np.ndarray: + mask = initial.copy() + try: + cv2.grabCut( + cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR), + mask, + None, + np.zeros((1, 65), np.float64), + np.zeros((1, 65), np.float64), + 2, + cv2.GC_INIT_WITH_MASK, + ) + output = np.isin(mask, [cv2.GC_FGD, cv2.GC_PR_FGD]).astype(np.uint8) * 255 + except cv2.error: + output = (initial >= cv2.GC_PR_FGD).astype(np.uint8) * 255 + output[container_mask == 0] = 0 + return output + + +def filter_components(mask: np.ndarray, minimum_area: int) -> np.ndarray: + count, labels, stats, _ = cv2.connectedComponentsWithStats((mask > 0).astype(np.uint8), 8) + output = np.zeros_like(mask, dtype=np.uint8) + for index in range(1, count): + if stats[index, cv2.CC_STAT_AREA] >= minimum_area: + output[labels == index] = 255 + return output + + +def soft_alpha(delta: np.ndarray, mask: np.ndarray, threshold: float) -> np.ndarray: + strength = np.clip((delta - threshold * 0.45) / max(threshold, 1), 0, 1) + alpha = np.rint(strength * 255).astype(np.uint8) + alpha[mask == 0] = 0 + alpha = cv2.GaussianBlur(alpha, (3, 3), 0.65) + alpha[cv2.dilate(mask, np.ones((3, 3), np.uint8)) == 0] = 0 + return alpha + + +def render_fill(shape: ShapeRegion, width: int, height: int) -> np.ndarray: + paint = shape.style.fill + if isinstance(paint, SolidPaint): + color = parse_hex(paint.color) + return np.broadcast_to(color, (height, width, 3)).copy() + assert isinstance(paint, LinearGradientPaint) + first, last = paint.stops[0], paint.stops[-1] + color0 = parse_hex(first.color).astype(np.float32) + color1 = parse_hex(last.color).astype(np.float32) + radians = np.deg2rad(paint.angle) + direction = np.array([np.cos(radians), np.sin(radians)], dtype=np.float32) + yy, xx = np.mgrid[0:height, 0:width] + projection = xx * direction[0] + yy * direction[1] + minimum, maximum = float(projection.min()), float(projection.max()) + t = (projection - minimum) / max(maximum - minimum, 1) + result = color0[None, None, :] * (1 - t[:, :, None]) + color1[None, None, :] * t[:, :, None] + return np.clip(np.rint(result), 0, 255).astype(np.uint8) + + +def reconstruct_background( + rgba: np.ndarray, + document: DocumentSpec, + combined_mask: np.ndarray, + inpaint_callback, +) -> np.ndarray: + height, width = rgba.shape[:2] + working = rgba.copy() + container_union = np.zeros((height, width), dtype=np.uint8) + for shape in sorted(document.shapes, key=lambda item: geometry_area(item), reverse=True): + if not shape.is_container or not shape.remove_source: + continue + mask = shape_source_mask(shape, width, height) + if not cv2.countNonZero(mask): + continue + ring = cv2.subtract( + cv2.dilate(mask, np.ones((15, 15), np.uint8)), + cv2.dilate(mask, np.ones((3, 3), np.uint8)), + ) + samples = working[:, :, :3][ring > 0] + fill = np.median(samples, axis=0) if samples.size else np.array([255, 255, 255]) + working[:, :, :3][mask > 0] = np.rint(fill).astype(np.uint8) + container_union = cv2.bitwise_or(container_union, mask) + remaining = combined_mask.copy() + remaining[container_union > 0] = 0 + if cv2.countNonZero(remaining): + working = inpaint_callback(working, remaining) + return working + + +def geometry_area(shape: ShapeRegion) -> float: + geometry = shape.geometry + if hasattr(geometry, "box"): + return geometry.box.width * geometry.box.height + if hasattr(geometry, "rx"): + return float(np.pi * geometry.rx * geometry.ry) + points = np.array([(point.x, point.y) for point in geometry.points], dtype=np.float32) + return abs(float(cv2.contourArea(points))) if len(points) >= 3 else 0 + + +def parse_hex(value: str) -> np.ndarray: + value = value.lstrip("#") + if len(value) == 3: + value = "".join(character * 2 for character in value) + return np.array([int(value[index:index + 2], 16) for index in (0, 2, 4)], dtype=np.uint8) diff --git a/src/editable_image/cli.py b/src/editable_image/cli.py new file mode 100644 index 0000000000000000000000000000000000000000..4e2490850fadaa39441447255ce15c1c7d36667d --- /dev/null +++ b/src/editable_image/cli.py @@ -0,0 +1,220 @@ +from __future__ import annotations + +import argparse +import importlib.util +import json +import sys +from pathlib import Path + +from .image_io import InvalidImage +from .inpaint import InpaintUnavailable, install_lama, lama_installed, resolve_torch_device +from .models import BitmapExtractor, CleanupMode, DEFAULT_SNAP_ANGLES, Device, ProcessOptions +from .ocr import OcrUnavailable, is_ocr_available +from .pipeline import export_result, process_path +from .sam2 import ( + ShapeExtractionUnavailable, + install_sam2, + sam2_dependencies_installed, + sam2_installed, +) + + +IMAGE_SUFFIXES = {".png", ".jpg", ".jpeg", ".webp"} + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(prog="editable-image") + subparsers = parser.add_subparsers(dest="command", required=True) + + convert = subparsers.add_parser("convert", help="convert bitmap images") + convert.add_argument("inputs", nargs="+", type=Path) + convert.add_argument("--output", "-o", type=Path, required=True) + convert.add_argument("--recursive", action="store_true") + convert.add_argument("--cleanup", choices=[item.value for item in CleanupMode], default="opencv") + convert.add_argument("--device", choices=[item.value for item in Device], default="auto") + convert.add_argument("--confidence", type=float, default=0.5) + convert.add_argument( + "--snap-angles", + type=parse_angles, + default=list(DEFAULT_SNAP_ANGLES), + metavar="DEGREES", + help="comma-separated target angles (default: 0,45,90,-45,-90)", + ) + convert.add_argument( + "--snap-tolerance", + type=float, + default=6.0, + metavar="DEGREES", + help="maximum distance from a target angle (default: 6)", + ) + convert.add_argument( + "--snap-font-sizes", + action="store_true", + help="snap estimated font sizes to peaks in the image's size distribution", + ) + convert.add_argument( + "--embed-fonts", + action="store_true", + help="embed used font faces in the SVG (default: link system fonts)", + ) + convert.add_argument( + "--vectorize-shapes", + action="store_true", + help="promote confident diagram primitives into editable SVG shapes", + ) + convert.add_argument( + "--bitmap-extractor", + choices=[item.value for item in BitmapExtractor], + default=BitmapExtractor.OPENCV.value, + help="bitmap-layer mask extractor used with --vectorize-shapes (default: opencv)", + ) + convert.add_argument("--overwrite", action="store_true") + + models = subparsers.add_parser("models", help="manage optional model files") + model_commands = models.add_subparsers(dest="model_command", required=True) + install = model_commands.add_parser("install") + install.add_argument("model", choices=["ocr", "lama", "sam2"]) + model_commands.add_parser("status") + return parser + + +def main(argv: list[str] | None = None) -> int: + args = build_parser().parse_args(argv) + if args.command == "models": + return model_command(args) + return convert_command(args) + + +def model_command(args: argparse.Namespace) -> int: + if args.model_command == "install": + if args.model == "lama": + path = install_lama() + print(f"Installed LaMa at {path}") + return 0 + if args.model == "sam2": + path = install_sam2() + print(f"Installed SAM 2 at {path}") + return 0 + if not is_ocr_available(): + print("Install the cpu or cuda project extra before installing OCR models", file=sys.stderr) + return 2 + from .ocr import RapidOcrEngine + + RapidOcrEngine(Device.CPU) + print("OCR models are ready") + return 0 + print( + json.dumps( + { + "ocr": is_ocr_available(), + "lama": lama_installed(), + "sam2": sam2_installed(), + "sam2_dependencies": sam2_dependencies_installed(), + "devices": available_devices(), + }, + indent=2, + ) + ) + return 0 + + +def convert_command(args: argparse.Namespace) -> int: + if args.bitmap_extractor != BitmapExtractor.OPENCV.value and not args.vectorize_shapes: + print("--bitmap-extractor requires --vectorize-shapes", file=sys.stderr) + return 2 + options = ProcessOptions( + confidence=args.confidence, + cleanup=CleanupMode(args.cleanup), + device=Device(args.device), + snap_angles=args.snap_angles, + snap_tolerance=args.snap_tolerance, + snap_font_sizes=args.snap_font_sizes, + vectorize_shapes=args.vectorize_shapes, + bitmap_extractor=BitmapExtractor(args.bitmap_extractor), + ) + inputs = collect_inputs(args.inputs, args.recursive) + if not inputs: + print("No supported images found", file=sys.stderr) + return 2 + failures = 0 + for source in inputs: + destination = args.output / source.stem + expected = [destination / f"{source.stem}-background.png", destination / f"{source.stem}-overlay.svg"] + if not args.overwrite and any(path.exists() for path in expected): + failures += 1 + print(f"skip {source}: output exists (use --overwrite)", file=sys.stderr) + continue + try: + result = process_path(source, options) + background, overlay, assets = export_result( + result, + destination, + source.stem, + embed_fonts=args.embed_fonts, + ) + print( + json.dumps( + { + "source": str(source), + "background": str(background), + "overlay": str(overlay), + "assets": [str(path) for path in assets], + } + ) + ) + except ( + InvalidImage, + OcrUnavailable, + InpaintUnavailable, + ShapeExtractionUnavailable, + OSError, + ValueError, + ) as exc: + failures += 1 + print(f"failed {source}: {exc}", file=sys.stderr) + return 1 if failures else 0 + + +def collect_inputs(paths: list[Path], recursive: bool) -> list[Path]: + output: list[Path] = [] + for path in paths: + if path.is_file() and path.suffix.lower() in IMAGE_SUFFIXES: + output.append(path) + elif path.is_dir(): + iterator = path.rglob("*") if recursive else path.glob("*") + output.extend(item for item in iterator if item.is_file() and item.suffix.lower() in IMAGE_SUFFIXES) + return sorted(set(output)) + + +def parse_angles(value: str) -> list[float]: + try: + angles = [float(item.strip()) for item in value.split(",") if item.strip()] + except ValueError as exc: + raise argparse.ArgumentTypeError("angles must be comma-separated numbers") from exc + if not angles: + raise argparse.ArgumentTypeError("at least one snap angle is required") + if any(angle < -180 or angle > 180 for angle in angles): + raise argparse.ArgumentTypeError("snap angles must be between -180 and 180") + return angles + + +def available_devices() -> list[str]: + devices = ["cpu"] + try: + import onnxruntime as ort + + if "CUDAExecutionProvider" in ort.get_available_providers(): + devices.append("cuda") + except ImportError: + pass + if importlib.util.find_spec("torch") is not None: + try: + if resolve_torch_device(Device.AUTO) == "mps": + devices.append("mps") + except InpaintUnavailable: + pass + return devices + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/src/editable_image/fonts.py b/src/editable_image/fonts.py new file mode 100644 index 0000000000000000000000000000000000000000..7abaa4a5c8ac7dc863e90a2ae46b8334b3775bf2 --- /dev/null +++ b/src/editable_image/fonts.py @@ -0,0 +1,69 @@ +from __future__ import annotations + +import base64 +from io import BytesIO +from dataclasses import dataclass +from functools import lru_cache +from importlib.resources import files +from pathlib import Path + + +@dataclass(frozen=True, slots=True) +class FontFace: + family: str + weight: int + italic: bool + filename: str + + @property + def path(self) -> Path: + return Path(str(files("editable_image").joinpath("fonts", self.filename))) + + +FACES = ( + FontFace("Lato", 400, False, "Lato-Regular.ttf"), + FontFace("Lato", 400, True, "Lato-Italic.ttf"), + FontFace("Lato", 700, False, "Lato-Bold.ttf"), + FontFace("Lato", 700, True, "Lato-BoldItalic.ttf"), + FontFace("Liberation Sans", 400, False, "LiberationSans-Regular.ttf"), + FontFace("Liberation Sans", 400, True, "LiberationSans-Italic.ttf"), + FontFace("Liberation Sans", 700, False, "LiberationSans-Bold.ttf"), + FontFace("Liberation Sans", 700, True, "LiberationSans-BoldItalic.ttf"), + FontFace("Liberation Serif", 400, False, "LiberationSerif-Regular.ttf"), + FontFace("Liberation Serif", 400, True, "LiberationSerif-Italic.ttf"), + FontFace("Liberation Serif", 700, False, "LiberationSerif-Bold.ttf"), + FontFace("Liberation Serif", 700, True, "LiberationSerif-BoldItalic.ttf"), + FontFace("Liberation Mono", 400, False, "LiberationMono-Regular.ttf"), + FontFace("Liberation Mono", 400, True, "LiberationMono-Italic.ttf"), + FontFace("Liberation Mono", 700, False, "LiberationMono-Bold.ttf"), + FontFace("Liberation Mono", 700, True, "LiberationMono-BoldItalic.ttf"), +) + + +def families() -> list[str]: + return list(dict.fromkeys(face.family for face in FACES if face.path.exists())) + + +def get_face(family: str, weight: int, italic: bool) -> FontFace: + candidates = [face for face in FACES if face.family == family and face.path.exists()] + if not candidates: + candidates = [face for face in FACES if face.path.exists()] + if not candidates: + raise FileNotFoundError("No bundled fonts are installed") + return min( + candidates, + key=lambda face: abs(face.weight - weight) + (1000 if face.italic != italic else 0), + ) + + +@lru_cache(maxsize=32) +def font_data_uri(filename: str) -> str: + path = Path(str(files("editable_image").joinpath("fonts", filename))) + from fontTools.ttLib import TTFont + + font = TTFont(path) + font.flavor = "woff2" + output = BytesIO() + font.save(output) + data = base64.b64encode(output.getvalue()).decode("ascii") + return f"data:font/woff2;base64,{data}" diff --git a/src/editable_image/fonts/LICENSES.txt b/src/editable_image/fonts/LICENSES.txt new file mode 100644 index 0000000000000000000000000000000000000000..ac2d86d356e61722d496accabf121173a29d6b44 --- /dev/null +++ b/src/editable_image/fonts/LICENSES.txt @@ -0,0 +1,10 @@ +Lato +Copyright (c) 2010-2014 by tyPoland Lukasz Dziedzic with Reserved Font Name "Lato". +Licensed under the SIL Open Font License, Version 1.1. +https://scripts.sil.org/OFL + +Liberation Fonts +Copyright (c) 2007 Red Hat, Inc. with Reserved Font Name "Liberation". +Licensed under the SIL Open Font License, Version 1.1. +https://scripts.sil.org/OFL + diff --git a/src/editable_image/fonts/Lato-Bold.ttf b/src/editable_image/fonts/Lato-Bold.ttf new file mode 100644 index 0000000000000000000000000000000000000000..07bed44d77b2f40b1af6de99d4f6a3facfb0b257 --- /dev/null +++ b/src/editable_image/fonts/Lato-Bold.ttf @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:bf1b8130069b44b9148eeece35e5423bedac49777ba746615b826b8276574a7b +size 657188 diff --git a/src/editable_image/fonts/Lato-BoldItalic.ttf b/src/editable_image/fonts/Lato-BoldItalic.ttf new file mode 100644 index 0000000000000000000000000000000000000000..50ad2cad2dc87e6a50b679b82d31ddd4328d59ba --- /dev/null +++ b/src/editable_image/fonts/Lato-BoldItalic.ttf @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:e2f9244fdcd20840f2480f91370cea885bf99230c49d153199fa64853ad74b69 +size 699008 diff --git a/src/editable_image/fonts/Lato-Italic.ttf b/src/editable_image/fonts/Lato-Italic.ttf new file mode 100644 index 0000000000000000000000000000000000000000..ec7b2fbb3d6dd92ac28d6a9ae3f2b5a1b0c7acaa --- /dev/null +++ b/src/editable_image/fonts/Lato-Italic.ttf @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:8c863f833b12bc8486671d26e3a58b13b0f9eaff1d6b10398e7c107a9b7ee05a +size 723544 diff --git a/src/editable_image/fonts/Lato-Regular.ttf b/src/editable_image/fonts/Lato-Regular.ttf new file mode 100644 index 0000000000000000000000000000000000000000..378e885d252fec567211cf48d84817f26a2dd1ba --- /dev/null +++ b/src/editable_image/fonts/Lato-Regular.ttf @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:6f6940be0835c3ddec9199e5fc42be4cbc61ebcfd58c623fdf719366253f1780 +size 657212 diff --git a/src/editable_image/fonts/LiberationMono-Bold.ttf b/src/editable_image/fonts/LiberationMono-Bold.ttf new file mode 100644 index 0000000000000000000000000000000000000000..37fbf423f224a1e032c2ff8a7f5551c4fbbe0967 --- /dev/null +++ b/src/editable_image/fonts/LiberationMono-Bold.ttf @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:626655e94dd82f3f42549daf995c921b0915fa8ab1f4b839559e8892ea41d240 +size 308068 diff --git a/src/editable_image/fonts/LiberationMono-BoldItalic.ttf b/src/editable_image/fonts/LiberationMono-BoldItalic.ttf new file mode 100644 index 0000000000000000000000000000000000000000..1ef2c44b57ec53640e99826b86f23a98d27a5724 --- /dev/null +++ b/src/editable_image/fonts/LiberationMono-BoldItalic.ttf @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:15eb161953e3ecc7fc05a3fec8a59e0f4e0a54a4e375736adf8e98d65981f813 +size 284184 diff --git a/src/editable_image/fonts/LiberationMono-Italic.ttf b/src/editable_image/fonts/LiberationMono-Italic.ttf new file mode 100644 index 0000000000000000000000000000000000000000..4459153820c1e9746af235e64ad9cdd8be2054aa --- /dev/null +++ b/src/editable_image/fonts/LiberationMono-Italic.ttf @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a71b2c25c89da05cf0e7c4dbba8d473fdead0b181ae56165217747d2c1f39215 +size 281608 diff --git a/src/editable_image/fonts/LiberationMono-Regular.ttf b/src/editable_image/fonts/LiberationMono-Regular.ttf new file mode 100644 index 0000000000000000000000000000000000000000..7cb636108eb06c956a08d6e58b699421a6dbcb1e --- /dev/null +++ b/src/editable_image/fonts/LiberationMono-Regular.ttf @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:395fa5ab8d40c8eba390ced528744ea75a7f69aabf3e68b6f925ca0e39a27370 +size 319624 diff --git a/src/editable_image/fonts/LiberationSans-Bold.ttf b/src/editable_image/fonts/LiberationSans-Bold.ttf new file mode 100644 index 0000000000000000000000000000000000000000..265356a957aff227753fec3cf4ad056821e711a5 --- /dev/null +++ b/src/editable_image/fonts/LiberationSans-Bold.ttf @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:3973aa5054fb467dd5627245d3dc82e37bf16fe075756156a570455871351582 +size 414568 diff --git a/src/editable_image/fonts/LiberationSans-BoldItalic.ttf b/src/editable_image/fonts/LiberationSans-BoldItalic.ttf new file mode 100644 index 0000000000000000000000000000000000000000..fe42459166a63bf3075c481e4f719e501fb09eb4 --- /dev/null +++ b/src/editable_image/fonts/LiberationSans-BoldItalic.ttf @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:c80fa7f2ffa0e01d4d8dcd6a6d1e43eda665222d1b4db597dde1174c456006cf +size 409108 diff --git a/src/editable_image/fonts/LiberationSans-Italic.ttf b/src/editable_image/fonts/LiberationSans-Italic.ttf new file mode 100644 index 0000000000000000000000000000000000000000..1e7f891f15f058728ce3ad6b9ace78f429f61593 --- /dev/null +++ b/src/editable_image/fonts/LiberationSans-Italic.ttf @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:830c5fa600505fb4c1a271b4271c53c44bae43f492b2a240d0e98a3a7a380121 +size 415920 diff --git a/src/editable_image/fonts/LiberationSans-Regular.ttf b/src/editable_image/fonts/LiberationSans-Regular.ttf new file mode 100644 index 0000000000000000000000000000000000000000..5ab11a333c8677ee94c1672159bb37eb1d594930 --- /dev/null +++ b/src/editable_image/fonts/LiberationSans-Regular.ttf @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:4659bc0c58c5028dd488ec928d41d9265db43d9b669fc14ca8b0832daca7b144 +size 410820 diff --git a/src/editable_image/fonts/LiberationSerif-Bold.ttf b/src/editable_image/fonts/LiberationSerif-Bold.ttf new file mode 100644 index 0000000000000000000000000000000000000000..2ad37722f99a2f77e7247ce68ab010c9466b20db --- /dev/null +++ b/src/editable_image/fonts/LiberationSerif-Bold.ttf @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:9e66c25e20868756bf12eb24a520a4552727c9cd108f577f7ae33e3c0110e39c +size 370196 diff --git a/src/editable_image/fonts/LiberationSerif-BoldItalic.ttf b/src/editable_image/fonts/LiberationSerif-BoldItalic.ttf new file mode 100644 index 0000000000000000000000000000000000000000..ad28df7fb761ff81cb6e42e45683cc25a7e0aaa2 --- /dev/null +++ b/src/editable_image/fonts/LiberationSerif-BoldItalic.ttf @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:379010e87421a883f7bbfa7936d23dfe0257a54ce850819b74556e9f2c615a8a +size 376892 diff --git a/src/editable_image/fonts/LiberationSerif-Italic.ttf b/src/editable_image/fonts/LiberationSerif-Italic.ttf new file mode 100644 index 0000000000000000000000000000000000000000..b26bb60f404fddb274f48749c3c8d8f6bd6f62ef --- /dev/null +++ b/src/editable_image/fonts/LiberationSerif-Italic.ttf @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:eaeaf7f2b12544ecde64290c1fa7ec403ff70fd58b313637b265400b20249242 +size 375760 diff --git a/src/editable_image/fonts/LiberationSerif-Regular.ttf b/src/editable_image/fonts/LiberationSerif-Regular.ttf new file mode 100644 index 0000000000000000000000000000000000000000..239a5cffacb38af0f2ec97f44522f88f439d4810 --- /dev/null +++ b/src/editable_image/fonts/LiberationSerif-Regular.ttf @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:705903ae1382f0150f115d6f10a30f5b2f3d6ade649c474afe6366c39174fb65 +size 393692 diff --git a/src/editable_image/image_io.py b/src/editable_image/image_io.py new file mode 100644 index 0000000000000000000000000000000000000000..94c7fbe6538ea3ea4fb33f20ba6b3a75a45fc4bf --- /dev/null +++ b/src/editable_image/image_io.py @@ -0,0 +1,73 @@ +from __future__ import annotations + +from dataclasses import dataclass +from io import BytesIO +from pathlib import Path + +import numpy as np +from PIL import Image, ImageOps, UnidentifiedImageError + + +ALLOWED_FORMATS = {"PNG", "JPEG", "WEBP"} +MAX_FILE_BYTES = 50 * 1024 * 1024 + + +class InvalidImage(ValueError): + pass + + +@dataclass(slots=True) +class LoadedImage: + rgba: np.ndarray + icc_profile: bytes | None + source_format: str + + @property + def width(self) -> int: + return int(self.rgba.shape[1]) + + @property + def height(self) -> int: + return int(self.rgba.shape[0]) + + +def load_image_bytes(data: bytes, max_pixels: int = 40_000_000) -> LoadedImage: + if not data: + raise InvalidImage("The image is empty") + if len(data) > MAX_FILE_BYTES: + raise InvalidImage("The image exceeds the 50 MB limit") + + try: + with Image.open(BytesIO(data)) as source: + if source.format not in ALLOWED_FORMATS: + raise InvalidImage("Only PNG, JPEG, and WebP images are supported") + if getattr(source, "is_animated", False): + raise InvalidImage("Animated images are not supported") + source.load() + image = ImageOps.exif_transpose(source) + if image.width * image.height > max_pixels: + raise InvalidImage(f"The decoded image exceeds {max_pixels:,} pixels") + rgba = np.asarray(image.convert("RGBA"), dtype=np.uint8).copy() + return LoadedImage( + rgba=rgba, + icc_profile=source.info.get("icc_profile"), + source_format=source.format, + ) + except (UnidentifiedImageError, OSError) as exc: + raise InvalidImage("The uploaded file is not a readable image") from exc + + +def load_image_path(path: Path, max_pixels: int = 40_000_000) -> LoadedImage: + try: + return load_image_bytes(path.read_bytes(), max_pixels=max_pixels) + except OSError as exc: + raise InvalidImage(f"Could not read {path}") from exc + + +def save_png(image: LoadedImage | np.ndarray, path: Path, icc_profile: bytes | None = None) -> None: + rgba = image.rgba if isinstance(image, LoadedImage) else image + if isinstance(image, LoadedImage): + icc_profile = image.icc_profile + target = Image.fromarray(rgba, "RGBA") + target.save(path, format="PNG", icc_profile=icc_profile, optimize=True) + diff --git a/src/editable_image/inpaint.py b/src/editable_image/inpaint.py new file mode 100644 index 0000000000000000000000000000000000000000..5df5c9b1e6392335a17abd80b7619e9f45981b3d --- /dev/null +++ b/src/editable_image/inpaint.py @@ -0,0 +1,103 @@ +from __future__ import annotations + +import hashlib +import importlib.util +import urllib.request +from pathlib import Path + +import cv2 +import numpy as np +from platformdirs import user_cache_path + +from .models import CleanupMode, Device + + +LAMA_URL = "https://github.com/Sanster/models/releases/download/add_big_lama/big-lama.pt" +LAMA_MD5 = "e3aa4aaa15225a33ec84f9f4bc47e500" + + +class InpaintUnavailable(RuntimeError): + pass + + +def lama_path() -> Path: + return user_cache_path("editable-image") / "models" / "big-lama.pt" + + +def lama_installed() -> bool: + path = lama_path() + return path.exists() and file_md5(path) == LAMA_MD5 + + +def file_md5(path: Path) -> str: + digest = hashlib.md5(usedforsecurity=False) + with path.open("rb") as stream: + for chunk in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def install_lama() -> Path: + target = lama_path() + target.parent.mkdir(parents=True, exist_ok=True) + temporary = target.with_suffix(".download") + urllib.request.urlretrieve(LAMA_URL, temporary) + if file_md5(temporary) != LAMA_MD5: + temporary.unlink(missing_ok=True) + raise InpaintUnavailable("The downloaded LaMa model failed checksum verification") + temporary.replace(target) + return target + + +def inpaint(rgba: np.ndarray, mask: np.ndarray, mode: CleanupMode, device: Device) -> np.ndarray: + if not np.any(mask): + return rgba.copy() + if mode == CleanupMode.LAMA: + return lama_inpaint(rgba, mask, device) + return opencv_inpaint(rgba, mask) + + +def opencv_inpaint(rgba: np.ndarray, mask: np.ndarray) -> np.ndarray: + rgb = rgba[:, :, :3] + alpha = rgba[:, :, 3] + radius = max(2, min(9, int(round(np.sqrt(np.count_nonzero(mask)) / 120)))) + restored_rgb = cv2.inpaint(cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR), mask, radius, cv2.INPAINT_TELEA) + restored_alpha = cv2.inpaint(alpha, mask, radius, cv2.INPAINT_TELEA) + return np.dstack([cv2.cvtColor(restored_rgb, cv2.COLOR_BGR2RGB), restored_alpha]) + + +def resolve_torch_device(requested: Device) -> str: + if importlib.util.find_spec("torch") is None: + raise InpaintUnavailable("Install the lama project extra to use neural cleanup") + import torch + + if requested == Device.CPU: + return "cpu" + if requested in {Device.AUTO, Device.CUDA} and torch.cuda.is_available(): + return "cuda" + if requested in {Device.AUTO, Device.MPS} and torch.backends.mps.is_available(): + return "mps" + if requested in {Device.CUDA, Device.MPS}: + raise InpaintUnavailable(f"The requested {requested.value} device is unavailable") + return "cpu" + + +def lama_inpaint(rgba: np.ndarray, mask: np.ndarray, device: Device) -> np.ndarray: + if not lama_installed(): + raise InpaintUnavailable("LaMa is not installed. Run 'editable-image models install lama'.") + import torch + + target_device = resolve_torch_device(device) + model = torch.jit.load(str(lama_path()), map_location=target_device).eval() + height, width = rgba.shape[:2] + pad_height = (8 - height % 8) % 8 + pad_width = (8 - width % 8) % 8 + rgb = np.pad(rgba[:, :, :3], ((0, pad_height), (0, pad_width), (0, 0)), mode="reflect") + padded_mask = np.pad(mask, ((0, pad_height), (0, pad_width)), mode="constant") + image_tensor = torch.from_numpy(rgb.transpose(2, 0, 1)).float().div(255).unsqueeze(0).to(target_device) + mask_tensor = torch.from_numpy((padded_mask > 0).astype(np.float32)).unsqueeze(0).unsqueeze(0).to(target_device) + with torch.inference_mode(): + restored = model(image_tensor, mask_tensor)[0].permute(1, 2, 0).clamp(0, 1) + result = (restored.cpu().numpy()[:height, :width] * 255).astype(np.uint8) + return np.dstack([result, rgba[:, :, 3]]) + diff --git a/src/editable_image/masking.py b/src/editable_image/masking.py new file mode 100644 index 0000000000000000000000000000000000000000..562ecfd9dacd994ed91d174f0e9ea53a165def14 --- /dev/null +++ b/src/editable_image/masking.py @@ -0,0 +1,94 @@ +from __future__ import annotations + +import cv2 +import numpy as np + +from .models import DocumentSpec, MaskStroke, Point, TextRegion +from .shape_recognition import shape_source_mask + + +def build_mask(rgba: np.ndarray, document: DocumentSpec) -> np.ndarray: + height, width = rgba.shape[:2] + combined = np.zeros((height, width), dtype=np.uint8) + for region in document.regions: + if region.remove_text: + combined = cv2.bitwise_or(combined, region_mask(rgba, region)) + for shape in document.shapes: + if shape.remove_source: + shape_mask = shape_source_mask(shape, width, height) + padding = max(0, int(round(shape.mask_padding))) + if padding: + kernel = cv2.getStructuringElement( + cv2.MORPH_ELLIPSE, + (padding * 2 + 1, padding * 2 + 1), + ) + shape_mask = cv2.dilate(shape_mask, kernel) + combined = cv2.bitwise_or(combined, shape_mask) + for stroke in document.mask_strokes: + apply_stroke(combined, stroke) + return combined + + +def region_mask(rgba: np.ndarray, region: TextRegion) -> np.ndarray: + height, width = rgba.shape[:2] + polygon = np.array([[point.x, point.y] for point in region.source_quad], dtype=np.int32) + bounds = cv2.boundingRect(polygon) + x, y, w, h = bounds + x0, y0 = max(0, x), max(0, y) + x1, y1 = min(width, x + w), min(height, y + h) + output = np.zeros((height, width), dtype=np.uint8) + if x1 <= x0 or y1 <= y0: + return output + + crop = rgba[y0:y1, x0:x1, :3] + lab = cv2.cvtColor(crop, cv2.COLOR_RGB2LAB).astype(np.float32) + local_polygon = polygon - np.array([x0, y0]) + polygon_mask = np.zeros(crop.shape[:2], dtype=np.uint8) + cv2.fillPoly(polygon_mask, [local_polygon], 255) + + perimeter = cv2.morphologyEx(polygon_mask, cv2.MORPH_GRADIENT, np.ones((3, 3), np.uint8)) + border_pixels = lab[perimeter > 0] + background = np.median(border_pixels, axis=0) if border_pixels.size else np.median(lab.reshape(-1, 3), axis=0) + delta = np.linalg.norm(lab - background, axis=2) + inside = delta[polygon_mask > 0] + if inside.size: + otsu_threshold, _ = cv2.threshold( + np.clip(inside, 0, 255).astype(np.uint8), + 0, + 255, + cv2.THRESH_BINARY + cv2.THRESH_OTSU, + ) + threshold = max(8.0, float(otsu_threshold) * 0.72) + else: + threshold = 8.0 + foreground = ((delta >= threshold) & (polygon_mask > 0)).astype(np.uint8) * 255 + + count, labels, stats, _ = cv2.connectedComponentsWithStats(foreground, 8) + filtered = np.zeros_like(foreground) + max_area = max(9, int(w * h * 0.22)) + min_area = 1 + for index in range(1, count): + area = stats[index, cv2.CC_STAT_AREA] + component_height = stats[index, cv2.CC_STAT_HEIGHT] + if min_area <= area <= max_area and component_height >= 1: + filtered[labels == index] = 255 + if cv2.countNonZero(filtered) < max(3, w * h * 0.005): + filtered = foreground + + padding = max(1, int(round(region.mask_padding))) + kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (padding * 2 + 1, padding * 2 + 1)) + filtered = cv2.dilate(filtered, kernel) + output[y0:y1, x0:x1] = filtered + return output + + +def apply_stroke(mask: np.ndarray, stroke: MaskStroke) -> None: + value = 255 if stroke.operation == "add" else 0 + radius = max(1, int(round(stroke.radius))) + points = [(int(round(point.x)), int(round(point.y))) for point in stroke.points] + if len(points) == 1: + cv2.circle(mask, points[0], radius, value, -1, lineType=cv2.LINE_AA) + return + for start, end in zip(points, points[1:]): + cv2.line(mask, start, end, value, radius * 2, lineType=cv2.LINE_AA) + cv2.circle(mask, end, radius, value, -1, lineType=cv2.LINE_AA) diff --git a/src/editable_image/models.py b/src/editable_image/models.py new file mode 100644 index 0000000000000000000000000000000000000000..a55ea874c032d8931b2af7b2f25e769fe8cda1cd --- /dev/null +++ b/src/editable_image/models.py @@ -0,0 +1,244 @@ +from __future__ import annotations + +from enum import StrEnum +from typing import Annotated, Literal + +from pydantic import BaseModel, ConfigDict, Field, field_validator + + +Pixel = Annotated[float, Field(ge=0)] +DEFAULT_SNAP_ANGLES = (0.0, 45.0, 90.0, -45.0, -90.0) + + +class CleanupMode(StrEnum): + OPENCV = "opencv" + LAMA = "lama" + + +class Device(StrEnum): + AUTO = "auto" + CPU = "cpu" + CUDA = "cuda" + MPS = "mps" + + +class BitmapExtractor(StrEnum): + OPENCV = "opencv" + SAM2 = "sam2" + + +class Point(BaseModel): + x: Pixel + y: Pixel + + +class Box(BaseModel): + x: Pixel + y: Pixel + width: Annotated[float, Field(gt=0)] + height: Annotated[float, Field(gt=0)] + + +class TextStyle(BaseModel): + font_family: str = "Lato" + font_size: Annotated[float, Field(ge=4, le=1000)] = 16 + font_weight: Literal[400, 700] = 400 + italic: bool = False + fill: str = "#111111" + opacity: Annotated[float, Field(ge=0, le=1)] = 1 + stroke: str = "none" + stroke_width: Annotated[float, Field(ge=0, le=50)] = 0 + letter_spacing: Annotated[float, Field(ge=-20, le=100)] = 0 + anchor: Literal["start", "middle", "end"] = "start" + + @field_validator("fill", "stroke") + @classmethod + def validate_color(cls, value: str) -> str: + if value == "none" or ( + len(value) in {4, 7, 9} and value.startswith("#") + ): + return value + raise ValueError("expected a CSS hexadecimal color or 'none'") + + +class TextRegion(BaseModel): + model_config = ConfigDict(extra="forbid") + + id: str + text: str + confidence: Annotated[float, Field(ge=0, le=1)] + source_quad: Annotated[list[Point], Field(min_length=4, max_length=4)] + overlay_box: Box + rotation: Annotated[float, Field(ge=-180, le=180)] = 0 + remove_text: bool = True + include_overlay: bool = True + mask_padding: Annotated[float, Field(ge=0, le=100)] = 3 + style: TextStyle = Field(default_factory=TextStyle) + z_index: int = 1000 + + +class SolidPaint(BaseModel): + model_config = ConfigDict(extra="forbid") + + type: Literal["solid"] = "solid" + color: str + + @field_validator("color") + @classmethod + def validate_color(cls, value: str) -> str: + if len(value) in {4, 7, 9} and value.startswith("#"): + return value + raise ValueError("expected a CSS hexadecimal color") + + +class GradientStop(BaseModel): + offset: Annotated[float, Field(ge=0, le=1)] + color: str + opacity: Annotated[float, Field(ge=0, le=1)] = 1 + + @field_validator("color") + @classmethod + def validate_color(cls, value: str) -> str: + if len(value) in {4, 7, 9} and value.startswith("#"): + return value + raise ValueError("expected a CSS hexadecimal color") + + +class LinearGradientPaint(BaseModel): + model_config = ConfigDict(extra="forbid") + + type: Literal["linear-gradient"] = "linear-gradient" + angle: Annotated[float, Field(ge=-180, le=180)] = 0 + stops: Annotated[list[GradientStop], Field(min_length=2, max_length=8)] + + +Paint = SolidPaint | LinearGradientPaint + + +class ShapeStyle(BaseModel): + fill: Paint = Field(default_factory=lambda: SolidPaint(color="#000000")) + fill_opacity: Annotated[float, Field(ge=0, le=1)] = 1 + stroke: str = "none" + stroke_width: Annotated[float, Field(ge=0, le=100)] = 0 + stroke_opacity: Annotated[float, Field(ge=0, le=1)] = 1 + linecap: Literal["butt", "round", "square"] = "round" + linejoin: Literal["miter", "round", "bevel"] = "round" + + @field_validator("stroke") + @classmethod + def validate_stroke(cls, value: str) -> str: + if value == "none" or (len(value) in {4, 7, 9} and value.startswith("#")): + return value + raise ValueError("expected a CSS hexadecimal color or 'none'") + + +class RectGeometry(BaseModel): + type: Literal["rect"] = "rect" + box: Box + rx: Annotated[float, Field(ge=0)] = 0 + ry: Annotated[float, Field(ge=0)] = 0 + rotation: Annotated[float, Field(ge=-180, le=180)] = 0 + + +class EllipseGeometry(BaseModel): + type: Literal["ellipse"] = "ellipse" + cx: Pixel + cy: Pixel + rx: Annotated[float, Field(gt=0)] + ry: Annotated[float, Field(gt=0)] + rotation: Annotated[float, Field(ge=-180, le=180)] = 0 + + +class PolygonGeometry(BaseModel): + type: Literal["polygon"] = "polygon" + points: Annotated[list[Point], Field(min_length=3)] + + +class PathGeometry(BaseModel): + type: Literal["path"] = "path" + points: Annotated[list[Point], Field(min_length=2)] + closed: bool = False + smooth: bool = False + arrow_start: bool = False + arrow_end: bool = False + + +ShapeGeometry = RectGeometry | EllipseGeometry | PolygonGeometry | PathGeometry + + +class ShapeRegion(BaseModel): + model_config = ConfigDict(extra="forbid") + + id: str + confidence: Annotated[float, Field(ge=0, le=1)] + geometry: ShapeGeometry + source_contours: list[list[Point]] = Field(default_factory=list) + style: ShapeStyle = Field(default_factory=ShapeStyle) + remove_source: bool = True + include_overlay: bool = True + mask_padding: Annotated[float, Field(ge=0, le=100)] = 1 + z_index: int = 100 + is_container: bool = False + semantic_role: Literal["shape", "arrow"] = "shape" + + +class BitmapRegion(BaseModel): + model_config = ConfigDict(extra="forbid") + + id: str + asset_name: str + box: Box + opacity: Annotated[float, Field(ge=0, le=1)] = 1 + z_index: int = 500 + + +class MaskStroke(BaseModel): + operation: Literal["add", "erase"] + radius: Annotated[float, Field(gt=0, le=500)] + points: Annotated[list[Point], Field(min_length=1)] + + +class Canvas(BaseModel): + width: Annotated[int, Field(gt=0)] + height: Annotated[int, Field(gt=0)] + + +class DocumentSpec(BaseModel): + model_config = ConfigDict(extra="forbid") + + version: Literal[1, 2] = 1 + source_name: str + canvas: Canvas + regions: list[TextRegion] + shapes: list[ShapeRegion] = Field(default_factory=list) + bitmap_regions: list[BitmapRegion] = Field(default_factory=list) + mask_strokes: list[MaskStroke] = Field(default_factory=list) + + +class ProcessOptions(BaseModel): + confidence: Annotated[float, Field(ge=0, le=1)] = 0.5 + device: Device = Device.AUTO + cleanup: CleanupMode = CleanupMode.OPENCV + max_pixels: Annotated[int, Field(gt=0)] = 40_000_000 + snap_angles: Annotated[ + list[Annotated[float, Field(ge=-180, le=180)]], + Field(min_length=1), + ] = Field(default_factory=lambda: list(DEFAULT_SNAP_ANGLES)) + snap_tolerance: Annotated[float, Field(ge=0, le=180)] = 6.0 + snap_font_sizes: bool = False + vectorize_shapes: bool = False + bitmap_extractor: BitmapExtractor = BitmapExtractor.OPENCV + + +class JobResponse(BaseModel): + id: str + document: DocumentSpec + original_url: str + background_url: str + + +class Capabilities(BaseModel): + ocr_available: bool + lama_installed: bool + devices: list[str] + fonts: list[str] diff --git a/src/editable_image/ocr.py b/src/editable_image/ocr.py new file mode 100644 index 0000000000000000000000000000000000000000..f87e653b43ce468216b45080b3cf29bca41e3961 --- /dev/null +++ b/src/editable_image/ocr.py @@ -0,0 +1,92 @@ +from __future__ import annotations + +import importlib.util +import math +from dataclasses import dataclass + +import numpy as np + +from .models import Device, Point + + +class OcrUnavailable(RuntimeError): + pass + + +@dataclass(slots=True) +class OcrLine: + text: str + confidence: float + quad: list[Point] + + +def is_ocr_available() -> bool: + return importlib.util.find_spec("rapidocr") is not None and ( + importlib.util.find_spec("onnxruntime") is not None + ) + + +class RapidOcrEngine: + def __init__(self, device: Device = Device.AUTO) -> None: + if not is_ocr_available(): + raise OcrUnavailable( + "ONNX Runtime is missing. Install the cpu or cuda project extra." + ) + from rapidocr import RapidOCR + + # RapidOCR 3.9 defaults to PP-OCRv6-small detection and recognition. + self._engine = RapidOCR(params={"Global.use_cls": False}) + + def recognize(self, rgba: np.ndarray, confidence: float) -> list[OcrLine]: + rgb = rgba[:, :, :3] + result = self._engine(rgb) + boxes = getattr(result, "boxes", None) + txts = getattr(result, "txts", None) + scores = getattr(result, "scores", None) + if boxes is None or txts is None or scores is None: + return [] + + lines: list[OcrLine] = [] + for box, text, score in zip(boxes, txts, scores, strict=True): + value = float(score) + clean = str(text).strip() + if value < confidence or not clean: + continue + quad = [Point(x=max(0, float(x)), y=max(0, float(y))) for x, y in box] + if len(quad) != 4 or polygon_area(quad) < 4: + continue + lines.append(OcrLine(text=clean, confidence=value, quad=quad)) + return sorted(lines, key=lambda line: (quad_box(line.quad)[1], quad_box(line.quad)[0])) + + +def polygon_area(points: list[Point]) -> float: + return abs( + sum( + point.x * points[(index + 1) % len(points)].y + - points[(index + 1) % len(points)].x * point.y + for index, point in enumerate(points) + ) + / 2 + ) + + +def quad_box(quad: list[Point]) -> tuple[float, float, float, float]: + xs = [point.x for point in quad] + ys = [point.y for point in quad] + return min(xs), min(ys), max(xs) - min(xs), max(ys) - min(ys) + + +def quad_rotation(quad: list[Point]) -> float: + first, second = quad[0], quad[1] + return math.degrees(math.atan2(second.y - first.y, second.x - first.x)) + + +def snap_rotation(angle: float, targets: list[float], tolerance: float) -> float: + nearest = min(targets, key=lambda target: angular_distance(angle, target)) + if angular_distance(angle, nearest) <= tolerance: + return float(nearest) + return angle + + +def angular_distance(first: float, second: float) -> float: + return abs((first - second + 180) % 360 - 180) diff --git a/src/editable_image/pipeline.py b/src/editable_image/pipeline.py new file mode 100644 index 0000000000000000000000000000000000000000..39cd7c71a4a17effdab6078923aabdc0366cb0c9 --- /dev/null +++ b/src/editable_image/pipeline.py @@ -0,0 +1,153 @@ +from __future__ import annotations + +from dataclasses import dataclass, field +from pathlib import Path +from uuid import uuid4 + +import numpy as np + +from .image_io import LoadedImage, load_image_path, save_png +from .bitmap_extract import extract_bitmap_layers, reconstruct_background +from .inpaint import inpaint +from .masking import build_mask +from .models import Canvas, DocumentSpec, ProcessOptions, TextRegion +from .ocr import RapidOcrEngine, snap_rotation +from .shape_recognition import detect_shapes +from .style import estimate_style, snap_sizes_to_distribution_peaks +from .svg_export import build_svg + + +@dataclass(slots=True) +class ProcessResult: + image: LoadedImage + document: DocumentSpec + mask: np.ndarray + background: np.ndarray + bitmap_assets: dict[str, np.ndarray] = field(default_factory=dict) + + +class EditableImagePipeline: + def __init__(self, options: ProcessOptions) -> None: + self.options = options + + def analyze(self, image: LoadedImage, source_name: str) -> ProcessResult: + engine = RapidOcrEngine(self.options.device) + lines = engine.recognize(image.rgba, self.options.confidence) + regions: list[TextRegion] = [] + for line in lines: + estimated = estimate_style(image.rgba, line.text, line.quad) + regions.append( + TextRegion( + id=f"text-{uuid4().hex[:12]}", + text=line.text, + confidence=line.confidence, + source_quad=line.quad, + overlay_box=estimated.box, + rotation=snap_rotation( + estimated.rotation, + self.options.snap_angles, + self.options.snap_tolerance, + ), + style=estimated.style, + ) + ) + if self.options.snap_font_sizes: + snapped_sizes = snap_sizes_to_distribution_peaks( + [region.style.font_size for region in regions] + ) + for region, font_size in zip(regions, snapped_sizes, strict=True): + region.style.font_size = font_size + document = DocumentSpec( + version=2 if self.options.vectorize_shapes else 1, + source_name=source_name, + canvas=Canvas(width=image.width, height=image.height), + regions=regions, + ) + if self.options.vectorize_shapes: + text_mask = build_mask(image.rgba, document) + document.shapes = detect_shapes(image.rgba, text_mask) + return self.rebuild(image, document) + + def rebuild(self, image: LoadedImage, document: DocumentSpec) -> ProcessResult: + if (document.canvas.width, document.canvas.height) != (image.width, image.height): + raise ValueError("Document dimensions do not match the source image") + text_only = document.model_copy(deep=True) + text_only.shapes = [] + text_only.bitmap_regions = [] + text_mask = build_mask(image.rgba, text_only) + bitmap_assets: dict[str, np.ndarray] = {} + if document.shapes: + bitmap_regions, bitmap_assets = extract_bitmap_layers( + image.rgba, + document, + text_mask, + self.options.bitmap_extractor, + self.options.device.value, + ) + document.bitmap_regions = bitmap_regions + mask = build_mask(image.rgba, document) + if document.shapes: + background = reconstruct_background( + image.rgba, + document, + mask, + lambda source, remaining: inpaint( + source, + remaining, + self.options.cleanup, + self.options.device, + ), + ) + else: + background = inpaint(image.rgba, mask, self.options.cleanup, self.options.device) + return ProcessResult( + image=image, + document=document, + mask=mask, + background=background, + bitmap_assets=bitmap_assets, + ) + + +def export_result( + result: ProcessResult, + directory: Path, + stem: str, + *, + embed_fonts: bool = False, +) -> tuple[Path, Path, list[Path]]: + directory.mkdir(parents=True, exist_ok=True) + background_path = directory / f"{stem}-background.png" + overlay_path = directory / f"{stem}-overlay.svg" + asset_directory = directory / f"{stem}-assets" + save_png(result.background, background_path, result.image.icc_profile) + asset_paths: list[Path] = [] + bitmap_hrefs: dict[str, tuple[str, str]] = {} + if result.bitmap_assets: + asset_directory.mkdir(parents=True, exist_ok=True) + regions = {region.id: region for region in result.document.bitmap_regions} + for region_id, pixels in result.bitmap_assets.items(): + region = regions[region_id] + asset_path = asset_directory / region.asset_name + save_png(pixels, asset_path) + asset_paths.append(asset_path) + bitmap_hrefs[region_id] = ( + f"{asset_directory.name}/{region.asset_name}", + asset_path.resolve().as_uri(), + ) + overlay_path.write_text( + build_svg( + result.document, + background_href=background_path.name, + background_fallback_href=background_path.name, + embed_fonts=embed_fonts, + bitmap_hrefs=bitmap_hrefs, + ), + encoding="utf-8", + ) + return background_path, overlay_path, asset_paths + + +def process_path(path: Path, options: ProcessOptions) -> ProcessResult: + image = load_image_path(path, max_pixels=options.max_pixels) + return EditableImagePipeline(options).analyze(image, path.name) diff --git a/src/editable_image/sam2.py b/src/editable_image/sam2.py new file mode 100644 index 0000000000000000000000000000000000000000..95e9373329fe1796bff0da72fa39036c1a05cbd5 --- /dev/null +++ b/src/editable_image/sam2.py @@ -0,0 +1,112 @@ +from __future__ import annotations + +import importlib.util +from functools import lru_cache +from pathlib import Path + +import cv2 +import numpy as np +from platformdirs import user_cache_path + +from .models import Device + + +SAM2_REPOSITORY = "facebook/sam2.1-hiera-tiny" + + +class ShapeExtractionUnavailable(RuntimeError): + pass + + +def sam2_path() -> Path: + return user_cache_path("editable-image") / "models" / "sam2.1-hiera-tiny" + + +def sam2_installed() -> bool: + target = sam2_path() + return (target / "config.json").exists() and (target / "model.safetensors").exists() + + +def sam2_dependencies_installed() -> bool: + return all(importlib.util.find_spec(name) is not None for name in ("torch", "transformers")) + + +def install_sam2() -> Path: + if importlib.util.find_spec("huggingface_hub") is None: + raise ShapeExtractionUnavailable( + "Install the sam2 project extra before downloading the SAM 2 model" + ) + from huggingface_hub import snapshot_download + + target = sam2_path() + target.mkdir(parents=True, exist_ok=True) + snapshot_download( + repo_id=SAM2_REPOSITORY, + local_dir=target, + allow_patterns=["*.json", "*.safetensors", "*.txt"], + ) + if not sam2_installed(): + raise ShapeExtractionUnavailable("The downloaded SAM 2 model is incomplete") + return target + + +def refine_with_sam2( + rgb: np.ndarray, + residual_mask: np.ndarray, + container_mask: np.ndarray, + device: str, +) -> np.ndarray: + if not sam2_dependencies_installed() or not sam2_installed(): + raise ShapeExtractionUnavailable( + "SAM 2 is unavailable. Install the sam2 extra and run " + "'editable-image models install sam2'." + ) + generator = mask_generator(resolve_device(Device(device))) + outputs = generator(rgb, points_per_batch=32) + masks = outputs.get("masks", []) + selected = np.zeros_like(residual_mask, dtype=np.uint8) + container_area = max(cv2.countNonZero(container_mask), 1) + for raw in masks: + candidate = np.asarray(raw, dtype=np.uint8) + if candidate.shape != residual_mask.shape: + candidate = cv2.resize( + candidate, + (residual_mask.shape[1], residual_mask.shape[0]), + interpolation=cv2.INTER_NEAREST, + ) + candidate = (candidate > 0).astype(np.uint8) * 255 + area = cv2.countNonZero(candidate) + if area == 0 or area > container_area * 0.8: + continue + inside = cv2.countNonZero(cv2.bitwise_and(candidate, container_mask)) / area + residual_overlap = cv2.countNonZero(cv2.bitwise_and(candidate, residual_mask)) / area + if inside >= 0.8 and residual_overlap >= 0.12: + selected = cv2.bitwise_or(selected, candidate) + selected[container_mask == 0] = 0 + return selected if cv2.countNonZero(selected) else residual_mask + + +def resolve_device(requested: Device) -> str: + import torch + + if requested == Device.CPU: + return "cpu" + if requested in {Device.AUTO, Device.CUDA} and torch.cuda.is_available(): + return "cuda" + if requested in {Device.AUTO, Device.MPS} and torch.backends.mps.is_available(): + return "mps" + if requested in {Device.CUDA, Device.MPS}: + raise ShapeExtractionUnavailable(f"The requested {requested.value} device is unavailable") + return "cpu" + + +@lru_cache(maxsize=3) +def mask_generator(device: str): + from transformers import pipeline + + pipeline_device: int | str = 0 if device == "cuda" else device + return pipeline( + "mask-generation", + model=str(sam2_path()), + device=pipeline_device, + ) diff --git a/src/editable_image/shape_recognition.py b/src/editable_image/shape_recognition.py new file mode 100644 index 0000000000000000000000000000000000000000..41c3afd746a53a1b63cb1faf6e3dbd0e356ffb17 --- /dev/null +++ b/src/editable_image/shape_recognition.py @@ -0,0 +1,769 @@ +from __future__ import annotations + +from dataclasses import dataclass +from math import atan2, degrees + +import cv2 +import numpy as np + +from .models import ( + Box, + EllipseGeometry, + GradientStop, + LinearGradientPaint, + PathGeometry, + Point, + PolygonGeometry, + RectGeometry, + ShapeRegion, + ShapeStyle, + SolidPaint, +) + + +@dataclass(slots=True) +class Candidate: + shape: ShapeRegion + mask: np.ndarray + score: float + + +def detect_shapes(rgba: np.ndarray, protected_mask: np.ndarray) -> list[ShapeRegion]: + """Recognize low-complexity diagram primitives without tracing illustrations.""" + rgb = rgba[:, :, :3] + height, width = rgb.shape[:2] + clusters = color_clusters(rgb, protected_mask) + candidates: list[Candidate] = smoothed_container_candidates(rgb, protected_mask) + candidates.extend( + arrow_candidates( + rgb, + protected_mask, + [*clusters, *arrow_color_masks(rgb, protected_mask)], + ) + ) + for cluster_index, cluster_mask in enumerate(clusters): + cleaned = cv2.morphologyEx(cluster_mask, cv2.MORPH_OPEN, np.ones((2, 2), np.uint8)) + cleaned = cv2.morphologyEx(cleaned, cv2.MORPH_CLOSE, np.ones((5, 5), np.uint8)) + contours, _ = cv2.findContours(cleaned, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) + for contour in contours: + candidate = fit_cluster_contour( + rgb, + cleaned, + contour, + protected_mask, + width, + height, + cluster_index, + ) + if candidate is not None: + candidates.append(candidate) + candidates.extend(edge_rectangle_candidates(rgb, protected_mask)) + accepted = deduplicate(candidates) + containers = [candidate for candidate in accepted if candidate.shape.is_container] + accepted = [ + candidate + for candidate in accepted + if not ( + isinstance(candidate.shape.geometry, EllipseGeometry) + and any( + cv2.countNonZero(cv2.bitwise_and(candidate.mask, container.mask)) + / max(cv2.countNonZero(candidate.mask), 1) + >= 0.9 + for container in containers + ) + and cv2.countNonZero(cv2.bitwise_and(candidate.mask, protected_mask)) < 3 + ) + ] + accepted = filter_arrow_context(accepted) + accepted.sort( + key=lambda item: (-item.shape.is_container, -int(item.mask.sum()), item.shape.id) + ) + for index, candidate in enumerate(accepted, 1): + candidate.shape.id = f"shape-{index:04d}" + candidate.shape.z_index = 100 if candidate.shape.is_container else 700 + return [candidate.shape for candidate in accepted] + + +def color_clusters(rgb: np.ndarray, protected_mask: np.ndarray, count: int = 10) -> list[np.ndarray]: + lab = cv2.cvtColor(rgb, cv2.COLOR_RGB2LAB).astype(np.float32) + pixels = lab.reshape(-1, 3) + allowed = protected_mask.reshape(-1) == 0 + indices = np.flatnonzero(allowed) + if indices.size == 0: + return [] + stride = max(1, indices.size // 50_000) + sample = pixels[indices[::stride]].copy() + count = min(count, max(2, sample.shape[0] // 100)) + cv2.setRNGSeed(1729) + _, _, centers = cv2.kmeans( + sample, + count, + None, + (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.4), + 3, + cv2.KMEANS_PP_CENTERS, + ) + distances = np.linalg.norm(pixels[:, None, :] - centers[None, :, :], axis=2) + labels = np.argmin(distances, axis=1).reshape(lab.shape[:2]) + output: list[np.ndarray] = [] + for index in range(count): + mask = ((labels == index) & (protected_mask == 0)).astype(np.uint8) * 255 + output.append(mask) + return output + + +def arrow_candidates( + rgb: np.ndarray, + protected_mask: np.ndarray, + clusters: list[np.ndarray], +) -> list[Candidate]: + """Find connected, concave silhouettes with a supported arrow tip and tail.""" + height, width = rgb.shape[:2] + image_area = width * height + output: list[Candidate] = [] + for cluster_index, cluster_mask in enumerate(clusters): + cleaned = cv2.morphologyEx(cluster_mask, cv2.MORPH_OPEN, np.ones((2, 2), np.uint8)) + cleaned = cv2.morphologyEx(cleaned, cv2.MORPH_CLOSE, np.ones((3, 3), np.uint8)) + contours, _ = cv2.findContours(cleaned, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) + for contour in contours: + area = float(cv2.contourArea(contour)) + x, y, w, h = cv2.boundingRect(contour) + aspect = max(w / max(h, 1), h / max(w, 1)) + if ( + area < max(45, image_area * 0.000025) + or w * h < 600 + or max(w, h) < 24 + # Compact arrowheads and short curved arrows can be almost square. + or aspect < 1.03 + ): + continue + perimeter = cv2.arcLength(contour, True) + if perimeter <= 0: + continue + approx = cv2.approxPolyDP(contour, max(1.2, perimeter * 0.008), True) + if not 5 <= len(approx) <= 48: + continue + hull = cv2.convexHull(contour) + hull_area = float(cv2.contourArea(hull)) + solidity = area / max(hull_area, 1) + fill_ratio = area / max(w * h, 1) + if ( + solidity >= 0.93 + or fill_ratio > 0.55 + or (len(approx) > 14 and fill_ratio >= 0.28) + or (area > image_area * 0.003 and aspect < 4) + ): + continue + tip = supported_arrow_tip(approx) + if tip is None: + continue + tip_index, tip_score = tip + points = [Point(x=float(item[0][0]), y=float(item[0][1])) for item in approx] + geometry = PathGeometry(points=points, closed=True) + fitted_mask = np.zeros((height, width), dtype=np.uint8) + cv2.drawContours(fitted_mask, [contour], -1, 255, -1) + sample_mask = cv2.bitwise_and(cleaned, fitted_mask) + sample_mask[protected_mask > 0] = 0 + style = estimate_style(rgb, sample_mask, geometry) + confidence = min(0.99, 0.72 + tip_score * 0.2 + (1 - solidity) * 0.1) + shape = ShapeRegion( + id=f"arrow-{cluster_index}-{tip_index}", + confidence=confidence, + geometry=geometry, + source_contours=[contour_points(contour)], + style=style, + semantic_role="arrow", + ) + output.append(Candidate(shape=shape, mask=fitted_mask, score=confidence)) + return output + + +def arrow_color_masks(rgb: np.ndarray, protected_mask: np.ndarray) -> list[np.ndarray]: + hsv = cv2.cvtColor(rgb, cv2.COLOR_RGB2HSV) + masks: list[np.ndarray] = [] + gray = ( + (hsv[:, :, 1] < 55) + & (hsv[:, :, 2] >= 45) + & (hsv[:, :, 2] <= 220) + & (protected_mask == 0) + ) + masks.append(gray.astype(np.uint8) * 255) + # Overlapping bins keep anti-aliased and shaded arrow pixels connected while + # still separating unrelated diagram colors. + for start in range(0, 165, 15): + colored = ( + (hsv[:, :, 0] >= start) + & (hsv[:, :, 0] < start + 30) + & (hsv[:, :, 1] >= 55) + & (hsv[:, :, 2] >= 35) + & (hsv[:, :, 2] <= 235) + & (protected_mask == 0) + ) + if np.count_nonzero(colored) >= 40: + masks.append(colored.astype(np.uint8) * 255) + red = ( + ((hsv[:, :, 0] < 15) | (hsv[:, :, 0] >= 165)) + & (hsv[:, :, 1] >= 55) + & (hsv[:, :, 2] >= 35) + & (hsv[:, :, 2] <= 235) + & (protected_mask == 0) + ) + if np.count_nonzero(red) >= 40: + masks.append(red.astype(np.uint8) * 255) + return masks + + +def filter_arrow_context(candidates: list[Candidate]) -> list[Candidate]: + containers = [candidate for candidate in candidates if candidate.shape.is_container] + arrow_groups: dict[int, list[Candidate]] = {index: [] for index in range(len(containers))} + outside: set[int] = set() + for candidate in candidates: + if candidate.shape.semantic_role != "arrow": + continue + area = max(cv2.countNonZero(candidate.mask), 1) + containing = [ + index + for index, container in enumerate(containers) + if cv2.countNonZero(cv2.bitwise_and(candidate.mask, container.mask)) / area >= 0.9 + ] + if not containing: + outside.add(id(candidate)) + else: + arrow_groups[containing[0]].append(candidate) + + accepted_inside: set[int] = set() + for container_index, group in arrow_groups.items(): + container = containers[container_index] + eligible = [ + candidate + for candidate in group + if color_distance(candidate.shape, container.shape) >= 25 + ] + peer_counts = { + id(candidate): sum( + color_distance(candidate.shape, other.shape) <= 75 + and similar_arrow_scale(candidate.mask, other.mask) + for other in eligible + ) + for candidate in eligible + } + stable = [candidate for candidate in eligible if peer_counts[id(candidate)] >= 4] + for candidate in stable: + stable_peers = sum( + color_distance(candidate.shape, other.shape) <= 75 + and similar_arrow_scale(candidate.mask, other.mask) + for other in stable + ) + # A coherent repeated set is strong evidence inside illustration + # panels; isolated arrow-like glyphs are usually raster detail. + if stable_peers >= 4: + accepted_inside.add(id(candidate)) + return [ + candidate + for candidate in candidates + if candidate.shape.semantic_role != "arrow" + or id(candidate) in outside + or id(candidate) in accepted_inside + ] + + +def color_distance(first: ShapeRegion, second: ShapeRegion) -> float: + if not isinstance(first.style.fill, SolidPaint) or not isinstance(second.style.fill, SolidPaint): + return float("inf") + return float(np.linalg.norm(hex_rgb(first.style.fill.color) - hex_rgb(second.style.fill.color))) + + +def similar_arrow_scale(first: np.ndarray, second: np.ndarray) -> bool: + first_area = max(cv2.countNonZero(first), 1) + second_area = max(cv2.countNonZero(second), 1) + ratio = max(first_area, second_area) / min(first_area, second_area) + return ratio <= 3.5 + + +def hex_rgb(value: str) -> np.ndarray: + value = value.lstrip("#") + if len(value) == 3: + value = "".join(character * 2 for character in value) + return np.array([int(value[index:index + 2], 16) for index in (0, 2, 4)], dtype=np.float32) + + +def supported_arrow_tip(approx: np.ndarray) -> tuple[int, float] | None: + points = approx[:, 0, :].astype(np.float32) + centroid = np.mean(points, axis=0) + best: tuple[int, float] | None = None + for index, tip in enumerate(points): + previous = points[(index - 1) % len(points)] + following = points[(index + 1) % len(points)] + side_a = previous - tip + side_b = following - tip + norm_a, norm_b = np.linalg.norm(side_a), np.linalg.norm(side_b) + if min(norm_a, norm_b) < 2: + continue + cosine = float(np.clip(np.dot(side_a, side_b) / (norm_a * norm_b), -1, 1)) + angle = float(np.degrees(np.arccos(cosine))) + if angle > 95: + continue + base_midpoint = (previous + following) / 2 + direction = tip - base_midpoint + direction_norm = np.linalg.norm(direction) + if direction_norm < 2: + continue + direction /= direction_norm + projections = (points - centroid) @ direction + tip_projection = float((tip - centroid) @ direction) + projection_span = float(projections.max() - projections.min()) + if projection_span < 8 or tip_projection < float(projections.max()) - projection_span * 0.08: + continue + head_width = float(np.linalg.norm(previous - following)) + if head_width < 4 or projection_span < head_width * 0.65: + continue + base_projection = float((base_midpoint - centroid) @ direction) + tail_points = int(np.count_nonzero(projections < base_projection - max(2, projection_span * 0.08))) + if tail_points < 2: + continue + symmetry = 1 - min(1.0, abs(norm_a - norm_b) / max(norm_a, norm_b)) + sharpness = max(0.0, 1 - angle / 100) + extremity = min(1.0, (tip_projection - float(np.median(projections))) / max(projection_span * 0.45, 1)) + score = sharpness * 0.4 + symmetry * 0.25 + extremity * 0.35 + if score >= 0.54 and (best is None or score > best[1]): + best = (index, score) + return best + + +def smoothed_container_candidates( + rgb: np.ndarray, + protected_mask: np.ndarray, +) -> list[Candidate]: + """Find large low-frequency panels after suppressing their foreground details.""" + height, width = rgb.shape[:2] + scale = min(1.0, 500 / max(width, height)) + small_width = max(32, int(round(width * scale))) + small_height = max(32, int(round(height * scale))) + small = cv2.resize(rgb, (small_width, small_height), interpolation=cv2.INTER_AREA) + kernel = min(21, min(small_width, small_height) // 2 * 2 - 1) + kernel = max(5, kernel) + smooth = cv2.medianBlur(small, kernel) + lab = cv2.cvtColor(smooth, cv2.COLOR_RGB2LAB).astype(np.float32) + pixels = lab.reshape(-1, 3) + cv2.setRNGSeed(2718) + _, labels, centers = cv2.kmeans( + pixels, + 12, + None, + (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.35), + 4, + cv2.KMEANS_PP_CENTERS, + ) + labels = labels.reshape(small_height, small_width) + output: list[Candidate] = [] + for index, center in enumerate(centers): + mask = (labels == index).astype(np.uint8) * 255 + mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, np.ones((5, 5), np.uint8)) + contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) + for contour in contours: + x, y, w, h = cv2.boundingRect(contour) + area = cv2.contourArea(contour) + if not ( + width * 0.12 <= w / scale <= width * 0.28 + and height * 0.45 <= h / scale <= height * 0.78 + and area / max(w * h, 1) >= 0.55 + ): + continue + box = Box( + x=x / scale, + y=y / scale, + width=min(width - x / scale, w / scale), + height=min(height - y / scale, h / scale), + ) + box = refine_container_box(rgb, box) + geometry = RectGeometry( + box=box, + rx=min(box.width, box.height) * 0.035, + ry=min(box.width, box.height) * 0.035, + ) + fitted = geometry_mask(geometry, width, height) + full_lab = cv2.cvtColor(rgb, cv2.COLOR_RGB2LAB).astype(np.float32) + close = np.linalg.norm(full_lab - center[None, None, :], axis=2) < 12 + sample = ((fitted > 0) & close & (protected_mask == 0)).astype(np.uint8) * 255 + style = estimate_style(rgb, sample, geometry) + corners = [ + Point(x=box.x, y=box.y), + Point(x=box.x + box.width, y=box.y), + Point(x=box.x + box.width, y=box.y + box.height), + Point(x=box.x, y=box.y + box.height), + ] + confidence = min(0.98, 0.82 + area / max(w * h, 1) * 0.16) + shape = ShapeRegion( + id=f"smooth-{index}", + confidence=confidence, + geometry=geometry, + source_contours=[corners], + style=style, + is_container=True, + ) + output.append(Candidate(shape=shape, mask=fitted, score=confidence)) + return output + + +def refine_container_box(rgb: np.ndarray, box: Box) -> Box: + height, width = rgb.shape[:2] + lab = cv2.cvtColor(rgb, cv2.COLOR_RGB2LAB).astype(np.float32) + border = np.concatenate( + [lab[:8].reshape(-1, 3), lab[-8:].reshape(-1, 3), lab[:, :8].reshape(-1, 3), lab[:, -8:].reshape(-1, 3)] + ) + delta = np.linalg.norm(lab - np.median(border, axis=0), axis=2) + x0, y0 = max(0, int(box.x)), max(0, int(box.y)) + x1 = min(width, int(np.ceil(box.x + box.width))) + y1 = min(height, int(np.ceil(box.y + box.height))) + row_signal = np.median(delta[y0:y1, x0:x1], axis=1) > 2.5 + row_run = longest_run(row_signal) + if row_run and row_run[1] - row_run[0] >= box.height * 0.5: + y0, y1 = y0 + row_run[0], y0 + row_run[1] + column_signal = np.median(delta[y0:y1, x0:x1], axis=0) > 2.5 + column_run = longest_run(column_signal) + if column_run and column_run[1] - column_run[0] >= box.width * 0.65: + x0, x1 = x0 + column_run[0], x0 + column_run[1] + return Box(x=x0, y=y0, width=max(1, x1 - x0), height=max(1, y1 - y0)) + + +def longest_run(values: np.ndarray) -> tuple[int, int] | None: + mask = (values.astype(np.uint8) * 255)[:, None] + mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, np.ones((7, 1), np.uint8))[:, 0] > 0 + padded = np.pad(mask.astype(np.int8), (1, 1)) + changes = np.diff(padded) + starts = np.flatnonzero(changes == 1) + ends = np.flatnonzero(changes == -1) + if not len(starts): + return None + index = int(np.argmax(ends - starts)) + return int(starts[index]), int(ends[index]) + + +def fit_cluster_contour( + rgb: np.ndarray, + cluster_mask: np.ndarray, + contour: np.ndarray, + protected_mask: np.ndarray, + width: int, + height: int, + cluster_index: int, +) -> Candidate | None: + area = float(cv2.contourArea(contour)) + image_area = width * height + if area < max(45, image_area * 0.00003): + return None + perimeter = cv2.arcLength(contour, True) + if perimeter <= 0: + return None + approx = cv2.approxPolyDP(contour, max(1.5, perimeter * 0.012), True) + x, y, w, h = cv2.boundingRect(contour) + if w < 5 or h < 5: + return None + contour_mask = np.zeros((height, width), dtype=np.uint8) + cv2.drawContours(contour_mask, [contour], -1, 255, -1) + circularity = 4 * np.pi * area / max(perimeter * perimeter, 1) + geometry: RectGeometry | EllipseGeometry | PolygonGeometry | PathGeometry + fitted_mask = np.zeros_like(contour_mask) + rectangularity = area / max(w * h, 1) + if len(approx) == 4 and rectangularity >= 0.58: + rotated = cv2.minAreaRect(contour) + (cx, cy), (rw, rh), angle = rotated + if rw < rh: + rw, rh = rh, rw + angle += 90 + angle = normalize_angle(angle) + if cx - rw / 2 < 0 or cy - rh / 2 < 0: + return None + radius = min(rw, rh) * (0.045 if rectangularity < 0.97 else 0) + geometry = RectGeometry( + box=Box(x=cx - rw / 2, y=cy - rh / 2, width=rw, height=rh), + rx=radius, + ry=radius, + rotation=angle, + ) + fitted_mask = geometry_mask(geometry, width, height) + elif circularity >= 0.72 and len(contour) >= 5: + (cx, cy), (ew, eh), angle = cv2.fitEllipse(contour) + if ew < 5 or eh < 5: + return None + geometry = EllipseGeometry(cx=cx, cy=cy, rx=ew / 2, ry=eh / 2, rotation=angle) + fitted_mask = geometry_mask(geometry, width, height) + elif 3 <= len(approx) <= 10: + points = [Point(x=float(item[0][0]), y=float(item[0][1])) for item in approx] + geometry = PolygonGeometry(points=points) + fitted_mask = geometry_mask(geometry, width, height) + elif len(approx) <= 20 and (w / h >= 3.5 or h / w >= 3.5): + points = [Point(x=float(item[0][0]), y=float(item[0][1])) for item in approx] + geometry = PathGeometry(points=points, closed=True) + fitted_mask = geometry_mask(geometry, width, height) + else: + return None + + large_container = area >= image_area * 0.012 and isinstance(geometry, RectGeometry) + if not large_container: + if isinstance(geometry, RectGeometry) and area < image_area * 0.00045: + return None + if isinstance(geometry, EllipseGeometry) and min(geometry.rx, geometry.ry) < 12: + return None + if isinstance(geometry, PolygonGeometry): + aspect = max(w / h, h / w) + if area < image_area * 0.00018 or not ( + len(geometry.points) == 3 or aspect >= 1.8 or area >= image_area * 0.0015 + ): + return None + + fitted_area = cv2.countNonZero(fitted_mask) + if fitted_area == 0: + return None + support = cv2.countNonZero(cv2.bitwise_and(cluster_mask, fitted_mask)) / fitted_area + overlap = cv2.countNonZero(cv2.bitwise_and(contour_mask, fitted_mask)) + union = cv2.countNonZero(cv2.bitwise_or(contour_mask, fitted_mask)) + iou = overlap / max(union, 1) + minimum_support = 0.38 if large_container else 0.64 + if support < minimum_support or iou < (0.58 if large_container else 0.68): + return None + + sample_mask = cv2.bitwise_and(cluster_mask, fitted_mask) + sample_mask[protected_mask > 0] = 0 + style = estimate_style(rgb, sample_mask, geometry) + source = contour_points(contour) + confidence = min(0.99, 0.45 * support + 0.45 * iou + 0.1) + shape = ShapeRegion( + id=f"cluster-{cluster_index}", + confidence=confidence, + geometry=geometry, + source_contours=[source], + style=style, + is_container=large_container, + ) + return Candidate(shape=shape, mask=fitted_mask, score=confidence) + + +def edge_rectangle_candidates(rgb: np.ndarray, protected_mask: np.ndarray) -> list[Candidate]: + height, width = rgb.shape[:2] + gray = cv2.cvtColor(rgb, cv2.COLOR_RGB2GRAY) + edges = cv2.Canny(gray, 40, 120) + edges[cv2.dilate(protected_mask, np.ones((3, 3), np.uint8)) > 0] = 0 + edges = cv2.morphologyEx(edges, cv2.MORPH_CLOSE, np.ones((5, 5), np.uint8)) + contours, _ = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) + output: list[Candidate] = [] + image_area = width * height + for contour in contours: + area = float(cv2.contourArea(contour)) + if area < image_area * 0.001: + continue + perimeter = cv2.arcLength(contour, True) + approx = cv2.approxPolyDP(contour, max(2, perimeter * 0.02), True) + x, y, w, h = cv2.boundingRect(contour) + if not 4 <= len(approx) <= 8 or w < 20 or h < 20: + continue + rectangularity = area / max(w * h, 1) + circularity = 4 * np.pi * area / max(perimeter * perimeter, 1) + if 0.8 <= w / h <= 1.2 and circularity >= 0.68: + continue + if rectangularity < 0.72: + continue + geometry = RectGeometry( + box=Box(x=x, y=y, width=w, height=h), + rx=min(w, h) * (0.045 if rectangularity < 0.97 else 0), + ry=min(w, h) * (0.045 if rectangularity < 0.97 else 0), + ) + mask = geometry_mask(geometry, width, height) + interior = cv2.erode(mask, np.ones((5, 5), np.uint8)) + interior[protected_mask > 0] = 0 + style = estimate_style(rgb, interior, geometry) + ring = cv2.subtract(mask, cv2.erode(mask, np.ones((5, 5), np.uint8))) + ring_pixels = rgb[ring > 0] + if ring_pixels.size: + style.stroke = rgb_hex(np.median(ring_pixels, axis=0)) + style.stroke_width = 2 + confidence = min(0.94, 0.55 + rectangularity * 0.4) + shape = ShapeRegion( + id="edge-rect", + confidence=confidence, + geometry=geometry, + source_contours=[contour_points(contour)], + style=style, + is_container=area >= image_area * 0.012, + ) + output.append(Candidate(shape=shape, mask=mask, score=confidence)) + return output + + +def estimate_style( + rgb: np.ndarray, + sample_mask: np.ndarray, + geometry: RectGeometry | EllipseGeometry | PolygonGeometry | PathGeometry, +) -> ShapeStyle: + ys, xs = np.nonzero(sample_mask) + if xs.size == 0: + return ShapeStyle(fill=SolidPaint(color="#000000")) + values = rgb[ys, xs].astype(np.float32) + median = np.median(values, axis=0) + solid_error = float(np.mean((values - median) ** 2)) + if xs.size >= 100 and solid_error > 5: + x_span = max(float(xs.max() - xs.min()), 1) + y_span = max(float(ys.max() - ys.min()), 1) + tx = (xs - xs.min()) / x_span + ty = (ys - ys.min()) / y_span + fit_x, error_x = linear_color_fit(values, tx) + fit_y, error_y = linear_color_fit(values, ty) + fit, error, angle = (fit_x, error_x, 0.0) if error_x <= error_y else (fit_y, error_y, 90.0) + if error < solid_error * 0.65 and np.linalg.norm(fit[1] - fit[0]) >= 5: + return ShapeStyle( + fill=LinearGradientPaint( + angle=angle, + stops=[ + GradientStop(offset=0, color=rgb_hex(fit[0])), + GradientStop(offset=1, color=rgb_hex(fit[1])), + ], + ) + ) + return ShapeStyle(fill=SolidPaint(color=rgb_hex(median))) + + +def linear_color_fit(values: np.ndarray, t: np.ndarray) -> tuple[np.ndarray, float]: + design = np.column_stack([np.ones_like(t), t]) + coefficients, _, _, _ = np.linalg.lstsq(design, values, rcond=None) + predicted = design @ coefficients + endpoints = np.clip(np.vstack([coefficients[0], coefficients[0] + coefficients[1]]), 0, 255) + return endpoints, float(np.mean((values - predicted) ** 2)) + + +def deduplicate(candidates: list[Candidate]) -> list[Candidate]: + accepted: list[Candidate] = [] + # Preserve semantic interpretations when their silhouette overlaps a more + # generic polygon/path fit. Containers still establish the layout first. + for candidate in sorted( + candidates, + key=lambda item: ( + item.shape.is_container, + item.shape.semantic_role == "arrow", + item.score, + ), + reverse=True, + ): + if ( + isinstance(candidate.shape.geometry, RectGeometry) + and not candidate.shape.is_container + and any( + existing.shape.is_container + and cv2.countNonZero(cv2.bitwise_and(candidate.mask, existing.mask)) + / max(cv2.countNonZero(candidate.mask), 1) + >= 0.9 + for existing in accepted + ) + ): + continue + duplicate = False + for existing in accepted: + intersection = cv2.countNonZero(cv2.bitwise_and(candidate.mask, existing.mask)) + union = cv2.countNonZero(cv2.bitwise_or(candidate.mask, existing.mask)) + smaller = min(cv2.countNonZero(candidate.mask), cv2.countNonZero(existing.mask)) + container_containment = ( + candidate.shape.is_container + and existing.shape.is_container + and smaller > 0 + and intersection / smaller >= 0.78 + ) + arrow_containment = ( + candidate.shape.semantic_role == "arrow" + and existing.shape.semantic_role == "arrow" + and smaller > 0 + and intersection / smaller >= 0.7 + ) + if union and ( + intersection / union >= 0.78 + or container_containment + or arrow_containment + ): + duplicate = True + break + if not duplicate: + accepted.append(candidate) + return accepted + + +def geometry_mask( + geometry: RectGeometry | EllipseGeometry | PolygonGeometry | PathGeometry, + width: int, + height: int, +) -> np.ndarray: + mask = np.zeros((height, width), dtype=np.uint8) + if isinstance(geometry, RectGeometry): + box = geometry.box + center = (box.x + box.width / 2, box.y + box.height / 2) + if abs(geometry.rotation) > 0.01: + points = cv2.boxPoints((center, (box.width, box.height), geometry.rotation)).astype(np.int32) + cv2.fillPoly(mask, [points], 255) + else: + p1 = (int(round(box.x)), int(round(box.y))) + p2 = (int(round(box.x + box.width)), int(round(box.y + box.height))) + radius = int(round(max(geometry.rx, geometry.ry))) + if radius > 0: + cv2.rectangle(mask, (p1[0] + radius, p1[1]), (p2[0] - radius, p2[1]), 255, -1) + cv2.rectangle(mask, (p1[0], p1[1] + radius), (p2[0], p2[1] - radius), 255, -1) + for cx, cy in ( + (p1[0] + radius, p1[1] + radius), + (p2[0] - radius, p1[1] + radius), + (p1[0] + radius, p2[1] - radius), + (p2[0] - radius, p2[1] - radius), + ): + cv2.circle(mask, (cx, cy), radius, 255, -1) + else: + cv2.rectangle(mask, p1, p2, 255, -1) + elif isinstance(geometry, EllipseGeometry): + cv2.ellipse( + mask, + (int(round(geometry.cx)), int(round(geometry.cy))), + (int(round(geometry.rx)), int(round(geometry.ry))), + geometry.rotation, + 0, + 360, + 255, + -1, + ) + else: + points = np.array([(point.x, point.y) for point in geometry.points], dtype=np.int32) + if isinstance(geometry, PolygonGeometry) or geometry.closed: + cv2.fillPoly(mask, [points], 255) + else: + cv2.polylines(mask, [points], False, 255, 2, cv2.LINE_AA) + return mask + + +def shape_source_mask(shape: ShapeRegion, width: int, height: int) -> np.ndarray: + if not shape.source_contours: + return geometry_mask(shape.geometry, width, height) + mask = np.zeros((height, width), dtype=np.uint8) + contours = [ + np.array([(point.x, point.y) for point in contour], dtype=np.int32) + for contour in shape.source_contours + if len(contour) >= 3 + ] + if contours: + cv2.fillPoly(mask, contours, 255) + return mask + + +def contour_points(contour: np.ndarray) -> list[Point]: + simplified = cv2.approxPolyDP(contour, max(1, cv2.arcLength(contour, True) * 0.004), True) + return [Point(x=float(item[0][0]), y=float(item[0][1])) for item in simplified] + + +def rgb_hex(value: np.ndarray) -> str: + red, green, blue = np.clip(np.rint(value), 0, 255).astype(np.uint8) + return f"#{red:02x}{green:02x}{blue:02x}" + + +def normalize_angle(value: float) -> float: + while value > 90: + value -= 180 + while value < -90: + value += 180 + return float(value) + + +def path_angle(points: list[Point]) -> float: + start, end = points[0], points[-1] + return degrees(atan2(end.y - start.y, end.x - start.x)) diff --git a/src/editable_image/style.py b/src/editable_image/style.py new file mode 100644 index 0000000000000000000000000000000000000000..d9832a3e51ecf4b4537d6029c29d50c0f7ade85e --- /dev/null +++ b/src/editable_image/style.py @@ -0,0 +1,201 @@ +from __future__ import annotations + +from dataclasses import dataclass + +import cv2 +import numpy as np +from PIL import Image, ImageDraw, ImageFont + +from .fonts import FACES +from .models import Box, Point, TextStyle +from .ocr import quad_box, quad_rotation + + +@dataclass(slots=True) +class StyleResult: + box: Box + rotation: float + style: TextStyle + + +def estimate_style(rgba: np.ndarray, text: str, quad: list[Point]) -> StyleResult: + x, y, width, height = quad_box(quad) + max_height, max_width = rgba.shape[:2] + left = max(0, int(x)) + top = max(0, int(y)) + right = min(max_width, int(np.ceil(x + width))) + bottom = min(max_height, int(np.ceil(y + height))) + crop = rgba[top:bottom, left:right, :3] + font_size = max(6.0, height * 0.92) + foreground = estimate_foreground_color(crop) + face, adjusted_size = match_font(crop, text, font_size, width) + style = TextStyle( + font_family=face.family, + font_size=adjusted_size, + font_weight=700 if face.weight >= 700 else 400, + italic=face.italic, + fill=foreground, + ) + return StyleResult( + box=Box(x=x, y=y, width=max(1, width), height=max(1, height)), + rotation=max(-180, min(180, quad_rotation(quad))), + style=style, + ) + + +def estimate_foreground_color(crop: np.ndarray) -> str: + if crop.size == 0: + return "#111111" + pixels = crop.reshape(-1, 3).astype(np.float32) + border = np.concatenate( + [crop[0, :, :3], crop[-1, :, :3], crop[:, 0, :3], crop[:, -1, :3]], axis=0 + ).astype(np.float32) + background = np.median(border, axis=0) + distances = np.linalg.norm(pixels - background, axis=1) + threshold = max(24.0, float(np.percentile(distances, 75))) + candidates = pixels[distances >= threshold] + if len(candidates) < 3: + candidates = pixels[np.argsort(distances)[-max(3, len(pixels) // 10) :]] + color = np.median(candidates, axis=0).clip(0, 255).astype(int) + return "#{:02x}{:02x}{:02x}".format(*color) + + +def match_font(crop: np.ndarray, text: str, size: float, target_width: float): + available = [face for face in FACES if face.family == "Lato" and face.path.exists()] + if not available: + raise FileNotFoundError("Bundled fonts are missing") + best_face = available[0] + best_size = max(6, int(round(size))) + best_score = float("inf") + observed = foreground_binary(crop) + likely_italic = estimate_italic(observed) + available = [face for face in available if face.italic == likely_italic] + for face in available: + candidate_size = max(6, int(round(size))) + font = ImageFont.truetype(str(face.path), candidate_size) + bounds = font.getbbox(text or " ") + rendered_width = max(1, bounds[2] - bounds[0]) + scaled_size = int(np.clip(candidate_size * target_width / rendered_width, 6, 1000)) + mask = render_text_mask(text, face.path, scaled_size, crop.shape[:2]) + score = mask_distance(observed, mask) + score += 0.002 if face.weight >= 700 else 0 + if score < best_score: + best_score = score + best_face = face + best_size = scaled_size + return best_face, float(best_size) + + +def foreground_binary(crop: np.ndarray) -> np.ndarray: + if crop.size == 0: + return np.zeros((1, 1), dtype=np.uint8) + gray = cv2.cvtColor(crop, cv2.COLOR_RGB2GRAY) + border = np.concatenate([gray[0], gray[-1], gray[:, 0], gray[:, -1]]) + background = float(np.median(border)) + delta = cv2.absdiff(gray, np.full_like(gray, int(background))) + _, mask = cv2.threshold(delta, 18, 255, cv2.THRESH_BINARY) + return mask + + +def render_text_mask(text: str, path, size: int, shape: tuple[int, int]) -> np.ndarray: + height, width = max(1, shape[0]), max(1, shape[1]) + image = Image.new("L", (width, height), 0) + draw = ImageDraw.Draw(image) + font = ImageFont.truetype(str(path), size) + bounds = draw.textbbox((0, 0), text, font=font) + y = max(-bounds[1], (height - (bounds[3] - bounds[1])) // 2 - bounds[1]) + draw.text((0, y), text, font=font, fill=255) + return np.asarray(image) + + +def mask_distance(observed: np.ndarray, candidate: np.ndarray) -> float: + if observed.shape != candidate.shape: + candidate = cv2.resize(candidate, (observed.shape[1], observed.shape[0])) + observed_pixels = observed > 0 + candidate_pixels = candidate > 0 + if not np.any(observed_pixels) or not np.any(candidate_pixels): + return 1.0 + distance_to_observed = cv2.distanceTransform( + (~observed_pixels).astype(np.uint8), cv2.DIST_L2, 3 + ) + distance_to_candidate = cv2.distanceTransform( + (~candidate_pixels).astype(np.uint8), cv2.DIST_L2, 3 + ) + chamfer = ( + float(distance_to_observed[candidate_pixels].mean()) + + float(distance_to_candidate[observed_pixels].mean()) + ) / (2 * max(1, observed.shape[0])) + density_delta = abs( + float(candidate_pixels.mean()) - float(observed_pixels.mean()) + ) + return chamfer + density_delta * 0.45 + + +def estimate_italic(mask: np.ndarray) -> bool: + if mask.size == 0 or np.count_nonzero(mask) < 12: + return False + height, width = mask.shape + shears = np.linspace(-0.3, 0.3, 13) + scores: list[float] = [] + for shear in shears: + transform = np.float32([[1, shear, -shear * height / 2], [0, 1, 0]]) + deskewed = cv2.warpAffine(mask, transform, (width, height)) > 0 + columns = deskewed.sum(axis=0).astype(np.float32) + scores.append(float(np.square(columns).sum() / max(1, columns.sum()))) + return float(shears[int(np.argmax(scores))]) >= 0.12 + + +def snap_sizes_to_distribution_peaks(sizes: list[float]) -> list[float]: + if len(sizes) < 2: + return list(sizes) + values = np.asarray(sizes, dtype=np.float64) + bandwidth = max(0.75, float(np.median(values)) * 0.06) + peaks = font_size_distribution_peaks(values, bandwidth) + if not peaks: + return list(sizes) + + snap_distance = bandwidth * 2 + snapped: list[float] = [] + for value in values: + nearest = min(peaks, key=lambda peak: abs(peak - value)) + snapped.append(nearest if abs(nearest - value) <= snap_distance else float(value)) + return snapped + + +def font_size_distribution_peaks(values: np.ndarray, bandwidth: float) -> list[float]: + lower = float(values.min()) + upper = float(values.max()) + if upper - lower < 1e-6: + return [lower] + + sample_count = min( + 4096, + max(64, int(np.ceil((upper - lower) / (bandwidth / 12))) + 1), + ) + grid = np.linspace(lower, upper, sample_count) + offsets = (grid[:, None] - values[None, :]) / bandwidth + density = np.exp(-0.5 * np.square(offsets)).sum(axis=1) + + candidates: list[int] = [] + if density[0] >= density[1]: + candidates.append(0) + candidates.extend( + index + for index in range(1, sample_count - 1) + if density[index] >= density[index - 1] + and density[index] >= density[index + 1] + ) + if density[-1] >= density[-2]: + candidates.append(sample_count - 1) + + peaks: list[float] = [] + for index in sorted(candidates, key=lambda item: density[item], reverse=True): + center = grid[index] + distances = np.abs(values - center) + supported = values[distances <= bandwidth * 2] + if len(supported) < 2: + continue + peak = round(float(np.mean(supported)) * 2) / 2 + if all(abs(peak - existing) > bandwidth * 2 for existing in peaks): + peaks.append(peak) + return sorted(peaks) diff --git a/src/editable_image/svg_export.py b/src/editable_image/svg_export.py new file mode 100644 index 0000000000000000000000000000000000000000..0c28b5f557b6de0bfb77170d818b469db48282fc --- /dev/null +++ b/src/editable_image/svg_export.py @@ -0,0 +1,308 @@ +from __future__ import annotations + +from urllib.parse import quote +from xml.etree import ElementTree as ET + +from .fonts import font_data_uri, get_face +from .models import ( + BitmapRegion, + DocumentSpec, + EllipseGeometry, + LinearGradientPaint, + PathGeometry, + PolygonGeometry, + RectGeometry, + ShapeRegion, + SolidPaint, + TextRegion, +) + + +SVG_NS = "http://www.w3.org/2000/svg" +XLINK_NS = "http://www.w3.org/1999/xlink" +ET.register_namespace("", SVG_NS) +ET.register_namespace("xlink", XLINK_NS) + + +def build_svg( + document: DocumentSpec, + background_href: str | None = None, + background_fallback_href: str | None = None, + embed_fonts: bool = False, + bitmap_hrefs: dict[str, tuple[str, str]] | None = None, +) -> str: + root = ET.Element( + f"{{{SVG_NS}}}svg", + { + "width": str(document.canvas.width), + "height": str(document.canvas.height), + "viewBox": f"0 0 {document.canvas.width} {document.canvas.height}", + "version": "1.1", + }, + ) + bitmap_hrefs = bitmap_hrefs or {} + needs_markers = any( + isinstance(shape.geometry, PathGeometry) + and (shape.geometry.arrow_start or shape.geometry.arrow_end) + for shape in document.shapes + ) + if embed_fonts or document.shapes or needs_markers: + definitions = ET.SubElement(root, f"{{{SVG_NS}}}defs") + else: + definitions = None + if embed_fonts: + used = { + (region.style.font_family, region.style.font_weight, region.style.italic) + for region in document.regions + if region.include_overlay and region.text + } + assert definitions is not None + style = ET.SubElement(definitions, f"{{{SVG_NS}}}style", {"type": "text/css"}) + style.text = "\n" + "\n".join(font_rule(*item) for item in sorted(used)) + "\n" + if definitions is not None: + for shape in document.shapes: + if isinstance(shape.style.fill, LinearGradientPaint): + definitions.append(gradient_element(shape)) + if needs_markers: + add_arrow_marker(definitions) + + if background_href: + relative_href = encoded_image_href(background_href) + fallback_href = encoded_image_href(background_fallback_href or background_href) + ET.SubElement( + root, + f"{{{SVG_NS}}}image", + { + "id": "cleared-background", + "href": relative_href, + f"{{{XLINK_NS}}}href": fallback_href, + "x": "0", + "y": "0", + "width": str(document.canvas.width), + "height": str(document.canvas.height), + "preserveAspectRatio": "none", + }, + ) + + layer = ET.SubElement(root, f"{{{SVG_NS}}}g", {"id": "editable-content"}) + content: list[tuple[int, str, ET.Element]] = [] + for shape in document.shapes: + if shape.include_overlay: + content.append((shape.z_index, shape.id, shape_element(shape))) + for bitmap in document.bitmap_regions: + hrefs = bitmap_hrefs.get(bitmap.id) + if hrefs: + content.append((bitmap.z_index, bitmap.id, bitmap_element(bitmap, *hrefs))) + for _, _, element in sorted(content, key=lambda item: (item[0], item[1])): + layer.append(element) + text_layer = ET.SubElement(root, f"{{{SVG_NS}}}g", {"id": "editable-text"}) + for region in sorted(document.regions, key=lambda item: (item.z_index, item.id)): + if region.include_overlay and region.text: + text_layer.append(text_element(region)) + return ET.tostring(root, encoding="unicode", xml_declaration=True) + + +def gradient_element(shape: ShapeRegion) -> ET.Element: + paint = shape.style.fill + assert isinstance(paint, LinearGradientPaint) + import math + + radians = math.radians(paint.angle) + dx, dy = math.cos(radians) * 50, math.sin(radians) * 50 + gradient = ET.Element( + f"{{{SVG_NS}}}linearGradient", + { + "id": f"gradient-{shape.id}", + "x1": f"{compact(50 - dx)}%", + "y1": f"{compact(50 - dy)}%", + "x2": f"{compact(50 + dx)}%", + "y2": f"{compact(50 + dy)}%", + }, + ) + for stop in paint.stops: + ET.SubElement( + gradient, + f"{{{SVG_NS}}}stop", + { + "offset": f"{compact(stop.offset * 100)}%", + "stop-color": stop.color, + "stop-opacity": compact(stop.opacity), + }, + ) + return gradient + + +def add_arrow_marker(definitions: ET.Element) -> None: + marker = ET.SubElement( + definitions, + f"{{{SVG_NS}}}marker", + { + "id": "editable-arrow", + "viewBox": "0 0 10 10", + "refX": "9", + "refY": "5", + "markerWidth": "6", + "markerHeight": "6", + "orient": "auto-start-reverse", + }, + ) + ET.SubElement(marker, f"{{{SVG_NS}}}path", {"d": "M 0 0 L 10 5 L 0 10 z", "fill": "context-stroke"}) + + +def shape_element(shape: ShapeRegion) -> ET.Element: + geometry = shape.geometry + style = shape.style + fill = ( + style.fill.color + if isinstance(style.fill, SolidPaint) + else f"url(#gradient-{shape.id})" + ) + attributes = { + "id": shape.id, + "data-role": shape.semantic_role, + "fill": fill, + "fill-opacity": compact(style.fill_opacity), + "stroke": style.stroke, + "stroke-width": compact(style.stroke_width), + "stroke-opacity": compact(style.stroke_opacity), + "stroke-linecap": style.linecap, + "stroke-linejoin": style.linejoin, + "data-confidence": compact(shape.confidence), + "data-source-contours": " | ".join( + " ".join(f"{compact(point.x)},{compact(point.y)}" for point in contour) + for contour in shape.source_contours + ), + } + if isinstance(geometry, RectGeometry): + box = geometry.box + attributes.update( + { + "x": compact(box.x), + "y": compact(box.y), + "width": compact(box.width), + "height": compact(box.height), + "rx": compact(geometry.rx), + "ry": compact(geometry.ry), + } + ) + if abs(geometry.rotation) > 0.01: + cx, cy = box.x + box.width / 2, box.y + box.height / 2 + attributes["transform"] = f"rotate({compact(geometry.rotation)} {compact(cx)} {compact(cy)})" + return ET.Element(f"{{{SVG_NS}}}rect", attributes) + if isinstance(geometry, EllipseGeometry): + attributes.update( + { + "cx": compact(geometry.cx), + "cy": compact(geometry.cy), + "rx": compact(geometry.rx), + "ry": compact(geometry.ry), + } + ) + if abs(geometry.rotation) > 0.01: + attributes["transform"] = ( + f"rotate({compact(geometry.rotation)} {compact(geometry.cx)} {compact(geometry.cy)})" + ) + return ET.Element(f"{{{SVG_NS}}}ellipse", attributes) + if isinstance(geometry, PolygonGeometry): + attributes["points"] = " ".join( + f"{compact(point.x)},{compact(point.y)}" for point in geometry.points + ) + return ET.Element(f"{{{SVG_NS}}}polygon", attributes) + assert isinstance(geometry, PathGeometry) + attributes["d"] = path_data(geometry) + if geometry.arrow_start: + attributes["marker-start"] = "url(#editable-arrow)" + if geometry.arrow_end: + attributes["marker-end"] = "url(#editable-arrow)" + if not geometry.closed and fill != "none": + attributes["fill"] = "none" + return ET.Element(f"{{{SVG_NS}}}path", attributes) + + +def path_data(geometry: PathGeometry) -> str: + points = geometry.points + if not geometry.smooth or len(points) < 3: + output = f"M {compact(points[0].x)} {compact(points[0].y)}" + output += "".join(f" L {compact(point.x)} {compact(point.y)}" for point in points[1:]) + else: + output = f"M {compact(points[0].x)} {compact(points[0].y)}" + for current, following in zip(points[1:-1], points[2:]): + midpoint_x = (current.x + following.x) / 2 + midpoint_y = (current.y + following.y) / 2 + output += ( + f" Q {compact(current.x)} {compact(current.y)}" + f" {compact(midpoint_x)} {compact(midpoint_y)}" + ) + output += f" T {compact(points[-1].x)} {compact(points[-1].y)}" + if geometry.closed: + output += " Z" + return output + + +def bitmap_element(bitmap: BitmapRegion, href: str, fallback_href: str) -> ET.Element: + box = bitmap.box + return ET.Element( + f"{{{SVG_NS}}}image", + { + "id": bitmap.id, + "href": encoded_image_href(href), + "x": compact(box.x), + "y": compact(box.y), + "width": compact(box.width), + "height": compact(box.height), + "opacity": compact(bitmap.opacity), + "preserveAspectRatio": "none", + }, + ) + + +def font_rule(family: str, weight: int, italic: bool) -> str: + face = get_face(family, weight, italic) + style = "italic" if italic else "normal" + return ( + f"@font-face {{font-family:'{face.family}';font-style:{style};" + f"font-weight:{face.weight};src:url('{font_data_uri(face.filename)}') format('woff2');}}" + ) + + +def text_element(region: TextRegion) -> ET.Element: + box = region.overlay_box + style = region.style + anchor_x = { + "start": box.x, + "middle": box.x + box.width / 2, + "end": box.x + box.width, + }[style.anchor] + baseline = box.y + box.height * 0.82 + attributes = { + "id": region.id, + "x": compact(anchor_x), + "y": compact(baseline), + "font-family": style.font_family, + "font-size": compact(style.font_size), + "font-weight": str(style.font_weight), + "font-style": "italic" if style.italic else "normal", + "fill": style.fill, + "fill-opacity": compact(style.opacity), + "stroke": style.stroke, + "stroke-width": compact(style.stroke_width), + "letter-spacing": compact(style.letter_spacing), + "text-anchor": style.anchor, + "data-confidence": compact(region.confidence), + "data-source-quad": " ".join(f"{compact(p.x)},{compact(p.y)}" for p in region.source_quad), + } + if abs(region.rotation) > 0.01: + attributes["transform"] = f"rotate({compact(region.rotation)} {compact(anchor_x)} {compact(baseline)})" + element = ET.Element(f"{{{SVG_NS}}}text", attributes) + element.text = region.text + return element + + +def encoded_image_href(href: str) -> str: + if href.startswith("data:image/png;base64,"): + return href + return quote(href, safe="/:%-._~") + + +def compact(value: float) -> str: + return f"{value:.3f}".rstrip("0").rstrip(".")