tonigi commited on
Commit
3c2cd23
·
verified ·
1 Parent(s): 5a139f9

Add complete Gradio interface

Browse files

Expose all conversion controls, downloadable SVG assets, preview, mask, metadata, and CPU Space dependencies.

This view is limited to 50 files because it contains too many changes.   See raw diff
Files changed (50) hide show
  1. .gitattributes +17 -0
  2. README.md +99 -8
  3. app.py +323 -0
  4. examples/scheme1.png +3 -0
  5. pyproject.toml +54 -0
  6. requirements.txt +2 -0
  7. src/editable_image/__init__.py +4 -0
  8. src/editable_image/__pycache__/__init__.cpython-313.pyc +0 -0
  9. src/editable_image/__pycache__/bitmap_extract.cpython-313.pyc +0 -0
  10. src/editable_image/__pycache__/cli.cpython-313.pyc +0 -0
  11. src/editable_image/__pycache__/fonts.cpython-313.pyc +0 -0
  12. src/editable_image/__pycache__/image_io.cpython-313.pyc +0 -0
  13. src/editable_image/__pycache__/inpaint.cpython-313.pyc +0 -0
  14. src/editable_image/__pycache__/masking.cpython-313.pyc +0 -0
  15. src/editable_image/__pycache__/models.cpython-313.pyc +0 -0
  16. src/editable_image/__pycache__/ocr.cpython-313.pyc +0 -0
  17. src/editable_image/__pycache__/pipeline.cpython-313.pyc +0 -0
  18. src/editable_image/__pycache__/sam2.cpython-313.pyc +0 -0
  19. src/editable_image/__pycache__/shape_recognition.cpython-313.pyc +0 -0
  20. src/editable_image/__pycache__/style.cpython-313.pyc +0 -0
  21. src/editable_image/__pycache__/svg_export.cpython-313.pyc +0 -0
  22. src/editable_image/bitmap_extract.py +203 -0
  23. src/editable_image/cli.py +220 -0
  24. src/editable_image/fonts.py +69 -0
  25. src/editable_image/fonts/LICENSES.txt +10 -0
  26. src/editable_image/fonts/Lato-Bold.ttf +3 -0
  27. src/editable_image/fonts/Lato-BoldItalic.ttf +3 -0
  28. src/editable_image/fonts/Lato-Italic.ttf +3 -0
  29. src/editable_image/fonts/Lato-Regular.ttf +3 -0
  30. src/editable_image/fonts/LiberationMono-Bold.ttf +3 -0
  31. src/editable_image/fonts/LiberationMono-BoldItalic.ttf +3 -0
  32. src/editable_image/fonts/LiberationMono-Italic.ttf +3 -0
  33. src/editable_image/fonts/LiberationMono-Regular.ttf +3 -0
  34. src/editable_image/fonts/LiberationSans-Bold.ttf +3 -0
  35. src/editable_image/fonts/LiberationSans-BoldItalic.ttf +3 -0
  36. src/editable_image/fonts/LiberationSans-Italic.ttf +3 -0
  37. src/editable_image/fonts/LiberationSans-Regular.ttf +3 -0
  38. src/editable_image/fonts/LiberationSerif-Bold.ttf +3 -0
  39. src/editable_image/fonts/LiberationSerif-BoldItalic.ttf +3 -0
  40. src/editable_image/fonts/LiberationSerif-Italic.ttf +3 -0
  41. src/editable_image/fonts/LiberationSerif-Regular.ttf +3 -0
  42. src/editable_image/image_io.py +73 -0
  43. src/editable_image/inpaint.py +103 -0
  44. src/editable_image/masking.py +94 -0
  45. src/editable_image/models.py +244 -0
  46. src/editable_image/ocr.py +92 -0
  47. src/editable_image/pipeline.py +153 -0
  48. src/editable_image/sam2.py +112 -0
  49. src/editable_image/shape_recognition.py +769 -0
  50. src/editable_image/style.py +201 -0
.gitattributes CHANGED
@@ -33,3 +33,20 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
33
  *.zip filter=lfs diff=lfs merge=lfs -text
34
  *.zst filter=lfs diff=lfs merge=lfs -text
35
  *tfevents* filter=lfs diff=lfs merge=lfs -text
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
33
  *.zip filter=lfs diff=lfs merge=lfs -text
34
  *.zst filter=lfs diff=lfs merge=lfs -text
35
  *tfevents* filter=lfs diff=lfs merge=lfs -text
36
+ examples/scheme1.png filter=lfs diff=lfs merge=lfs -text
37
+ src/editable_image/fonts/Lato-Bold.ttf filter=lfs diff=lfs merge=lfs -text
38
+ src/editable_image/fonts/Lato-BoldItalic.ttf filter=lfs diff=lfs merge=lfs -text
39
+ src/editable_image/fonts/Lato-Italic.ttf filter=lfs diff=lfs merge=lfs -text
40
+ src/editable_image/fonts/Lato-Regular.ttf filter=lfs diff=lfs merge=lfs -text
41
+ src/editable_image/fonts/LiberationMono-Bold.ttf filter=lfs diff=lfs merge=lfs -text
42
+ src/editable_image/fonts/LiberationMono-BoldItalic.ttf filter=lfs diff=lfs merge=lfs -text
43
+ src/editable_image/fonts/LiberationMono-Italic.ttf filter=lfs diff=lfs merge=lfs -text
44
+ src/editable_image/fonts/LiberationMono-Regular.ttf filter=lfs diff=lfs merge=lfs -text
45
+ src/editable_image/fonts/LiberationSans-Bold.ttf filter=lfs diff=lfs merge=lfs -text
46
+ src/editable_image/fonts/LiberationSans-BoldItalic.ttf filter=lfs diff=lfs merge=lfs -text
47
+ src/editable_image/fonts/LiberationSans-Italic.ttf filter=lfs diff=lfs merge=lfs -text
48
+ src/editable_image/fonts/LiberationSans-Regular.ttf filter=lfs diff=lfs merge=lfs -text
49
+ src/editable_image/fonts/LiberationSerif-Bold.ttf filter=lfs diff=lfs merge=lfs -text
50
+ src/editable_image/fonts/LiberationSerif-BoldItalic.ttf filter=lfs diff=lfs merge=lfs -text
51
+ src/editable_image/fonts/LiberationSerif-Italic.ttf filter=lfs diff=lfs merge=lfs -text
52
+ src/editable_image/fonts/LiberationSerif-Regular.ttf filter=lfs diff=lfs merge=lfs -text
README.md CHANGED
@@ -1,14 +1,105 @@
1
  ---
2
- title: Make Editable Image
3
- emoji: 🐠
4
- colorFrom: green
5
- colorTo: purple
6
  sdk: gradio
7
- sdk_version: 6.26.0
8
- python_version: '3.12'
9
  app_file: app.py
 
10
  pinned: false
11
- short_description: Make editable
12
  ---
13
 
14
- Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
  ---
2
+ title: Editable Image
3
+ emoji: ✏️
4
+ colorFrom: blue
5
+ colorTo: indigo
6
  sdk: gradio
 
 
7
  app_file: app.py
8
+ python_version: "3.11"
9
  pinned: false
 
10
  ---
11
 
12
+ # Editable Image
13
+
14
+ 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.
15
+
16
+ ## Install
17
+
18
+ Python 3.11 or newer and `uv` are recommended.
19
+
20
+ ```bash
21
+ uv sync --extra cpu --extra test
22
+ uv run python make_editable_image.py models install ocr
23
+ ```
24
+
25
+ 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:
26
+
27
+ ```bash
28
+ uv run python make_editable_image.py models install lama
29
+ ```
30
+
31
+ ## Use
32
+
33
+ ```bash
34
+ uv run python make_editable_image.py image.png --output output
35
+ uv run python make_editable_image.py images/ --recursive --output output
36
+ ```
37
+
38
+ Each input produces `<output>/<image-name>/<image-name>-background.png` and `<image-name>-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.
39
+
40
+ ### Gradio / Hugging Face Spaces
41
+
42
+ The repository includes a Gradio interface in `app.py` with controls for every
43
+ conversion option. To run it locally:
44
+
45
+ ```bash
46
+ uv run --with-requirements requirements.txt python app.py
47
+ ```
48
+
49
+ For Hugging Face Spaces, create a Gradio Space and upload the repository. The
50
+ README metadata selects `app.py`, and `requirements.txt` installs the package
51
+ with CPU ONNX Runtime. The default OpenCV cleanup and bitmap extraction work
52
+ without additional model downloads. LaMa and SAM 2 remain available in the UI
53
+ when their optional dependencies and model files are added to the Space.
54
+
55
+ Add `--vectorize-shapes` to promote confident diagram primitives—rounded rectangles, circles, ellipses, simple polygons, and connector-like paths—into editable SVG elements:
56
+
57
+ ```bash
58
+ uv run python make_editable_image.py image.png --output output --vectorize-shapes
59
+ ```
60
+
61
+ Detailed illustrations are preserved as transparent PNG crops under `<image-name>-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.
62
+
63
+ Bitmap masks use OpenCV by default. SAM 2 can optionally refine difficult masks:
64
+
65
+ ```bash
66
+ uv sync --extra cpu --extra sam2
67
+ uv run python make_editable_image.py models install sam2
68
+ uv run python make_editable_image.py image.png --output output \
69
+ --vectorize-shapes --bitmap-extractor sam2
70
+ ```
71
+
72
+ The SAM 2 checkpoint is downloaded only by the explicit install command.
73
+
74
+ 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`.
75
+
76
+ 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.
77
+
78
+ SVGs reference font family names by default to keep files small. Add `--embed-fonts` to embed the used WOFF2 font faces for portable rendering.
79
+
80
+ ## Example output
81
+
82
+ 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.
83
+
84
+ ## Notes
85
+
86
+ - The default OpenCV cleanup is fast and works best on flat or slowly varying backgrounds.
87
+ - LaMa can improve larger removals over textured backgrounds but adds PyTorch and a much larger model.
88
+ - Font recovery is approximate. The app bundles OFL-licensed Lato and Liberation faces; embedding them is optional.
89
+ - The script accepts PNG, JPEG, and WebP images up to 50 MB and 40 megapixels.
90
+ - Shape vectorization targets clean diagram geometry, not photographs, logos, or detailed illustrations.
91
+
92
+ # Method
93
+
94
+ The conversion is a hybrid OCR and image-processing pipeline. It does not use a large generative model for the normal workflow.
95
+
96
+ 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.
97
+ 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.
98
+ 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.
99
+ 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.
100
+ 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.
101
+ 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.
102
+ 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.
103
+ 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.
104
+ 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.
105
+ 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 `<text>` 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.
app.py ADDED
@@ -0,0 +1,323 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from __future__ import annotations
2
+
3
+ import base64
4
+ import html
5
+ import re
6
+ import tempfile
7
+ import zipfile
8
+ from pathlib import Path
9
+
10
+ import gradio as gr
11
+
12
+ from editable_image.models import (
13
+ BitmapExtractor,
14
+ CleanupMode,
15
+ DEFAULT_SNAP_ANGLES,
16
+ Device,
17
+ ProcessOptions,
18
+ )
19
+ from editable_image.pipeline import ProcessResult, export_result, process_path
20
+ from editable_image.svg_export import build_svg
21
+
22
+
23
+ APP_TITLE = "Editable Image"
24
+ APP_DESCRIPTION = (
25
+ "Remove printed text from a bitmap and rebuild it as an editable SVG text layer. "
26
+ "Shape vectorization can also promote confident diagram geometry to SVG elements."
27
+ )
28
+ ANGLE_DEFAULT = ", ".join(f"{angle:g}" for angle in DEFAULT_SNAP_ANGLES)
29
+ CSS = """
30
+ .gradio-container { max-width: 1280px !important; }
31
+ .hero { margin: 0 auto 0.6rem; max-width: 900px; text-align: center; }
32
+ .hero h1 { font-size: 2.15rem; margin-bottom: 0.25rem; }
33
+ .hero p { color: var(--body-text-color-subdued); font-size: 1rem; }
34
+ .svg-preview {
35
+ align-items: center;
36
+ background-color: #f4f4f5;
37
+ background-image: linear-gradient(45deg, #e4e4e7 25%, transparent 25%),
38
+ linear-gradient(-45deg, #e4e4e7 25%, transparent 25%),
39
+ linear-gradient(45deg, transparent 75%, #e4e4e7 75%),
40
+ linear-gradient(-45deg, transparent 75%, #e4e4e7 75%);
41
+ background-position: 0 0, 0 8px, 8px -8px, -8px 0;
42
+ background-size: 16px 16px;
43
+ border: 1px solid var(--border-color-primary);
44
+ border-radius: 10px;
45
+ display: flex;
46
+ justify-content: center;
47
+ min-height: 420px;
48
+ overflow: auto;
49
+ padding: 1rem;
50
+ }
51
+ .svg-preview img { display: block; max-height: 72vh; max-width: 100%; object-fit: contain; }
52
+ """
53
+ GRADIO_MAJOR = int(gr.__version__.split(".", maxsplit=1)[0])
54
+
55
+
56
+ def parse_snap_angles(value: str) -> list[float]:
57
+ try:
58
+ angles = [float(item.strip()) for item in value.split(",") if item.strip()]
59
+ except (AttributeError, ValueError) as exc:
60
+ raise ValueError("Snap angles must be comma-separated numbers.") from exc
61
+ if not angles:
62
+ raise ValueError("Enter at least one snap angle.")
63
+ if any(angle < -180 or angle > 180 for angle in angles):
64
+ raise ValueError("Snap angles must be between -180 and 180 degrees.")
65
+ return angles
66
+
67
+
68
+ def safe_stem(path: Path) -> str:
69
+ stem = re.sub(r"[^A-Za-z0-9._-]+", "-", path.stem).strip("-._")
70
+ return stem[:100] or "image"
71
+
72
+
73
+ def file_data_uri(path: Path, media_type: str) -> str:
74
+ payload = base64.b64encode(path.read_bytes()).decode("ascii")
75
+ return f"data:{media_type};base64,{payload}"
76
+
77
+
78
+ def preview_html(
79
+ result: ProcessResult,
80
+ background_path: Path,
81
+ asset_paths: list[Path],
82
+ embed_fonts: bool,
83
+ ) -> str:
84
+ background_uri = file_data_uri(background_path, "image/png")
85
+ asset_by_name = {path.name: path for path in asset_paths}
86
+ bitmap_hrefs: dict[str, tuple[str, str]] = {}
87
+ for region in result.document.bitmap_regions:
88
+ asset_path = asset_by_name.get(region.asset_name)
89
+ if asset_path is not None:
90
+ uri = file_data_uri(asset_path, "image/png")
91
+ bitmap_hrefs[region.id] = (uri, uri)
92
+ svg = build_svg(
93
+ result.document,
94
+ background_href=background_uri,
95
+ background_fallback_href=background_uri,
96
+ embed_fonts=embed_fonts,
97
+ bitmap_hrefs=bitmap_hrefs,
98
+ )
99
+ encoded = base64.b64encode(svg.encode("utf-8")).decode("ascii")
100
+ alt = html.escape(f"Editable reconstruction of {result.document.source_name}", quote=True)
101
+ return (
102
+ '<div class="svg-preview">'
103
+ f'<img src="data:image/svg+xml;base64,{encoded}" alt="{alt}">'
104
+ "</div>"
105
+ )
106
+
107
+
108
+ def make_bundle(output_directory: Path, stem: str) -> Path:
109
+ archive = output_directory.parent / f"{stem}-editable.zip"
110
+ with zipfile.ZipFile(archive, "w", compression=zipfile.ZIP_DEFLATED) as bundle:
111
+ for path in sorted(output_directory.rglob("*")):
112
+ if path.is_file():
113
+ bundle.write(path, arcname=path.relative_to(output_directory.parent))
114
+ return archive
115
+
116
+
117
+ def convert_image(
118
+ image_path: str | None,
119
+ confidence: float,
120
+ cleanup: str,
121
+ device: str,
122
+ max_pixels: float,
123
+ snap_angles: str,
124
+ snap_tolerance: float,
125
+ snap_font_sizes: bool,
126
+ embed_fonts: bool,
127
+ vectorize_shapes: bool,
128
+ bitmap_extractor: str,
129
+ ) -> tuple[str, object, object, dict[str, object], str, str]:
130
+ if not image_path:
131
+ raise gr.Error("Upload a PNG, JPEG, or WebP image first.")
132
+
133
+ source = Path(image_path)
134
+ try:
135
+ options = ProcessOptions(
136
+ confidence=confidence,
137
+ cleanup=CleanupMode(cleanup),
138
+ device=Device(device),
139
+ max_pixels=int(max_pixels),
140
+ snap_angles=parse_snap_angles(snap_angles),
141
+ snap_tolerance=snap_tolerance,
142
+ snap_font_sizes=snap_font_sizes,
143
+ vectorize_shapes=vectorize_shapes,
144
+ bitmap_extractor=BitmapExtractor(bitmap_extractor),
145
+ )
146
+ result = process_path(source, options)
147
+
148
+ job_root = Path(tempfile.mkdtemp(prefix="editable-image-"))
149
+ stem = safe_stem(source)
150
+ output_directory = job_root / stem
151
+ background_path, overlay_path, asset_paths = export_result(
152
+ result,
153
+ output_directory,
154
+ stem,
155
+ embed_fonts=embed_fonts,
156
+ )
157
+ document_path = output_directory / f"{stem}-document.json"
158
+ document_path.write_text(
159
+ result.document.model_dump_json(indent=2),
160
+ encoding="utf-8",
161
+ )
162
+ bundle_path = make_bundle(output_directory, stem)
163
+ rendered = preview_html(result, background_path, asset_paths, embed_fonts)
164
+ except Exception as exc:
165
+ # Gradio displays this as a concise toast instead of exposing a server traceback.
166
+ raise gr.Error(str(exc) or type(exc).__name__) from exc
167
+
168
+ metadata: dict[str, object] = {
169
+ "source": result.document.source_name,
170
+ "dimensions": f"{result.image.width} × {result.image.height}",
171
+ "recognized_text_regions": len(result.document.regions),
172
+ "vector_shapes": len(result.document.shapes),
173
+ "bitmap_assets": len(result.bitmap_assets),
174
+ "cleanup": options.cleanup.value,
175
+ "device": options.device.value,
176
+ "document_version": result.document.version,
177
+ }
178
+ return (
179
+ rendered,
180
+ result.background,
181
+ result.mask,
182
+ metadata,
183
+ str(overlay_path),
184
+ str(bundle_path),
185
+ )
186
+
187
+
188
+ def build_demo() -> gr.Blocks:
189
+ blocks_options: dict[str, object] = {"title": APP_TITLE}
190
+ if GRADIO_MAJOR < 6:
191
+ blocks_options.update(theme=gr.themes.Soft(), css=CSS)
192
+ with gr.Blocks(**blocks_options) as demo:
193
+ gr.HTML(
194
+ '<section class="hero">'
195
+ f"<h1>{APP_TITLE}</h1><p>{APP_DESCRIPTION}</p>"
196
+ "</section>"
197
+ )
198
+ with gr.Row(equal_height=False):
199
+ with gr.Column(scale=5):
200
+ source = gr.Image(
201
+ label="Source image",
202
+ type="filepath",
203
+ image_mode="RGBA",
204
+ sources=["upload", "clipboard"],
205
+ height=430,
206
+ )
207
+ convert = gr.Button("Make editable", variant="primary", size="lg")
208
+ gr.Markdown(
209
+ "PNG, JPEG, and WebP are supported. Keep the downloaded bundle intact: "
210
+ "its SVG links to the sibling background and any extracted bitmap assets."
211
+ )
212
+ with gr.Column(scale=4):
213
+ with gr.Accordion("OCR and orientation", open=True):
214
+ confidence = gr.Slider(
215
+ 0,
216
+ 1,
217
+ value=0.5,
218
+ step=0.01,
219
+ label="Minimum OCR confidence",
220
+ )
221
+ snap_angles = gr.Textbox(
222
+ value=ANGLE_DEFAULT,
223
+ label="Snap angles (degrees)",
224
+ info="Comma-separated targets between -180 and 180.",
225
+ )
226
+ snap_tolerance = gr.Slider(
227
+ 0,
228
+ 180,
229
+ value=6,
230
+ step=0.5,
231
+ label="Angle snap tolerance (degrees)",
232
+ )
233
+ snap_font_sizes = gr.Checkbox(
234
+ value=False,
235
+ label="Snap font sizes to detected distribution peaks",
236
+ )
237
+ embed_fonts = gr.Checkbox(
238
+ value=False,
239
+ label="Embed used WOFF2 fonts in the SVG",
240
+ )
241
+ with gr.Accordion("Cleanup and vectorization", open=True):
242
+ cleanup = gr.Dropdown(
243
+ choices=[item.value for item in CleanupMode],
244
+ value=CleanupMode.OPENCV.value,
245
+ label="Cleanup method",
246
+ info="LaMa requires the optional dependency and installed model.",
247
+ )
248
+ vectorize_shapes = gr.Checkbox(
249
+ value=False,
250
+ label="Vectorize confident diagram shapes",
251
+ )
252
+ bitmap_extractor = gr.Dropdown(
253
+ choices=[item.value for item in BitmapExtractor],
254
+ value=BitmapExtractor.OPENCV.value,
255
+ label="Bitmap layer extractor",
256
+ info="Used during shape vectorization. SAM 2 requires its optional model.",
257
+ )
258
+ with gr.Accordion("Runtime and limits", open=False):
259
+ device = gr.Dropdown(
260
+ choices=[item.value for item in Device],
261
+ value=Device.AUTO.value,
262
+ label="Compute device",
263
+ )
264
+ max_pixels = gr.Number(
265
+ value=40_000_000,
266
+ minimum=1,
267
+ maximum=200_000_000,
268
+ step=1,
269
+ precision=0,
270
+ label="Maximum decoded pixels",
271
+ )
272
+
273
+ gr.Markdown("## Result")
274
+ with gr.Tabs():
275
+ with gr.Tab("Editable preview"):
276
+ rendered_preview = gr.HTML(
277
+ '<div class="svg-preview">Upload an image and choose “Make editable”.</div>',
278
+ label="Reconstructed SVG",
279
+ )
280
+ with gr.Tab("Cleared background"):
281
+ background = gr.Image(label="Background PNG", interactive=False)
282
+ with gr.Tab("Removal mask"):
283
+ mask = gr.Image(label="Applied mask", interactive=False)
284
+ with gr.Tab("Details"):
285
+ metadata = gr.JSON(label="Conversion details")
286
+ with gr.Row():
287
+ overlay = gr.File(label="Editable SVG")
288
+ bundle = gr.File(label="Complete bundle (.zip)")
289
+
290
+ controls = [
291
+ source,
292
+ confidence,
293
+ cleanup,
294
+ device,
295
+ max_pixels,
296
+ snap_angles,
297
+ snap_tolerance,
298
+ snap_font_sizes,
299
+ embed_fonts,
300
+ vectorize_shapes,
301
+ bitmap_extractor,
302
+ ]
303
+ outputs = [rendered_preview, background, mask, metadata, overlay, bundle]
304
+ convert.click(convert_image, inputs=controls, outputs=outputs, api_name="convert")
305
+
306
+ example = Path(__file__).resolve().parent / "examples" / "scheme1.png"
307
+ if example.exists():
308
+ gr.Examples(
309
+ examples=[[str(example)]],
310
+ inputs=[source],
311
+ label="Example",
312
+ )
313
+ return demo
314
+
315
+
316
+ demo = build_demo()
317
+
318
+
319
+ if __name__ == "__main__":
320
+ launch_options: dict[str, object] = {}
321
+ if GRADIO_MAJOR >= 6:
322
+ launch_options.update(theme=gr.themes.Soft(), css=CSS)
323
+ demo.queue(default_concurrency_limit=1).launch(**launch_options)
examples/scheme1.png ADDED

Git LFS Details

  • SHA256: 33322a4fc618424afc3c79ec844b0ecd472d6590824f462ad99a601b18e8c6d1
  • Pointer size: 132 Bytes
  • Size of remote file: 1.43 MB
pyproject.toml ADDED
@@ -0,0 +1,54 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ [build-system]
2
+ requires = ["hatchling>=1.27"]
3
+ build-backend = "hatchling.build"
4
+
5
+ [project]
6
+ name = "editable-image"
7
+ version = "0.1.0"
8
+ description = "Remove bitmap text into an editable SVG overlay"
9
+ readme = "README.md"
10
+ requires-python = ">=3.11"
11
+ license = { text = "MIT" }
12
+ dependencies = [
13
+ "fonttools[woff]>=4.59,<5",
14
+ "numpy>=2.0,<3",
15
+ "opencv-python>=4.10,<6",
16
+ "pillow>=11,<13",
17
+ "platformdirs>=4.3,<5",
18
+ "pydantic>=2.10,<3",
19
+ "rapidocr>=3.9,<4",
20
+ ]
21
+
22
+ [project.optional-dependencies]
23
+ cpu = ["onnxruntime>=1.20,<2"]
24
+ cuda = ["onnxruntime-gpu>=1.20,<2"]
25
+ lama = ["torch>=2.5,<3"]
26
+ sam2 = [
27
+ "huggingface-hub>=0.34,<2",
28
+ "torch>=2.5,<3",
29
+ "torchvision>=0.20,<1",
30
+ "transformers>=4.57,<6",
31
+ ]
32
+ test = [
33
+ "pytest>=8.3,<9",
34
+ "pytest-cov>=6,<8",
35
+ ]
36
+
37
+ [project.scripts]
38
+ editable-image = "editable_image.cli:main"
39
+
40
+ [tool.hatch.build.targets.wheel]
41
+ packages = ["src/editable_image"]
42
+
43
+ [tool.hatch.build]
44
+ exclude = ["src/editable_image/static"]
45
+
46
+ [tool.hatch.build.targets.wheel.force-include]
47
+ "src/editable_image/fonts" = "editable_image/fonts"
48
+
49
+ [tool.pytest.ini_options]
50
+ addopts = "-q"
51
+ testpaths = ["tests"]
52
+
53
+ [tool.uv]
54
+ conflicts = [[{ extra = "cpu" }, { extra = "cuda" }]]
requirements.txt ADDED
@@ -0,0 +1,2 @@
 
 
 
1
+ -e .[cpu]
2
+ gradio>=5,<7
src/editable_image/__init__.py ADDED
@@ -0,0 +1,4 @@
 
 
 
 
 
1
+ """Bitmap text extraction into editable SVG overlays."""
2
+
3
+ __version__ = "0.1.0"
4
+
src/editable_image/__pycache__/__init__.cpython-313.pyc ADDED
Binary file (258 Bytes). View file
 
src/editable_image/__pycache__/bitmap_extract.cpython-313.pyc ADDED
Binary file (14.7 kB). View file
 
src/editable_image/__pycache__/cli.cpython-313.pyc ADDED
Binary file (11.6 kB). View file
 
src/editable_image/__pycache__/fonts.cpython-313.pyc ADDED
Binary file (4.67 kB). View file
 
src/editable_image/__pycache__/image_io.cpython-313.pyc ADDED
Binary file (4.39 kB). View file
 
src/editable_image/__pycache__/inpaint.cpython-313.pyc ADDED
Binary file (7.9 kB). View file
 
src/editable_image/__pycache__/masking.cpython-313.pyc ADDED
Binary file (7.2 kB). View file
 
src/editable_image/__pycache__/models.cpython-313.pyc ADDED
Binary file (12.6 kB). View file
 
src/editable_image/__pycache__/ocr.cpython-313.pyc ADDED
Binary file (6.13 kB). View file
 
src/editable_image/__pycache__/pipeline.cpython-313.pyc ADDED
Binary file (8.12 kB). View file
 
src/editable_image/__pycache__/sam2.cpython-313.pyc ADDED
Binary file (6.13 kB). View file
 
src/editable_image/__pycache__/shape_recognition.cpython-313.pyc ADDED
Binary file (49 kB). View file
 
src/editable_image/__pycache__/style.cpython-313.pyc ADDED
Binary file (14.6 kB). View file
 
src/editable_image/__pycache__/svg_export.cpython-313.pyc ADDED
Binary file (18.6 kB). View file
 
src/editable_image/bitmap_extract.py ADDED
@@ -0,0 +1,203 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from __future__ import annotations
2
+
3
+ import cv2
4
+ import numpy as np
5
+
6
+ from .models import (
7
+ BitmapExtractor,
8
+ BitmapRegion,
9
+ Box,
10
+ DocumentSpec,
11
+ LinearGradientPaint,
12
+ ShapeRegion,
13
+ SolidPaint,
14
+ )
15
+ from .shape_recognition import geometry_mask, shape_source_mask
16
+
17
+
18
+ def extract_bitmap_layers(
19
+ rgba: np.ndarray,
20
+ document: DocumentSpec,
21
+ text_mask: np.ndarray,
22
+ extractor: BitmapExtractor,
23
+ device: str,
24
+ ) -> tuple[list[BitmapRegion], dict[str, np.ndarray]]:
25
+ height, width = rgba.shape[:2]
26
+ rgb = rgba[:, :, :3]
27
+ text_exclusion = cv2.dilate(text_mask, np.ones((5, 5), np.uint8))
28
+ for region in document.regions:
29
+ polygon = np.array(
30
+ [(point.x, point.y) for point in region.source_quad],
31
+ dtype=np.int32,
32
+ )
33
+ cv2.fillPoly(text_exclusion, [polygon], 255)
34
+ vector_foreground = np.zeros((height, width), dtype=np.uint8)
35
+ for shape in document.shapes:
36
+ if not shape.is_container and shape.include_overlay:
37
+ vector_foreground = cv2.bitwise_or(
38
+ vector_foreground,
39
+ geometry_mask(shape.geometry, width, height),
40
+ )
41
+
42
+ regions: list[BitmapRegion] = []
43
+ assets: dict[str, np.ndarray] = {}
44
+ for shape in document.shapes:
45
+ if not shape.is_container or not shape.include_overlay:
46
+ continue
47
+ container_mask = geometry_mask(shape.geometry, width, height)
48
+ predicted = render_fill(shape, width, height)
49
+ source_lab = cv2.cvtColor(rgb, cv2.COLOR_RGB2LAB).astype(np.float32)
50
+ predicted_lab = cv2.cvtColor(predicted, cv2.COLOR_RGB2LAB).astype(np.float32)
51
+ delta = np.linalg.norm(source_lab - predicted_lab, axis=2)
52
+ border_exclusion = cv2.subtract(
53
+ container_mask,
54
+ cv2.erode(container_mask, np.ones((7, 7), np.uint8)),
55
+ )
56
+ eligible = (
57
+ (container_mask > 0)
58
+ & (text_exclusion == 0)
59
+ & (vector_foreground == 0)
60
+ & (border_exclusion == 0)
61
+ )
62
+ values = delta[eligible]
63
+ if values.size < 20:
64
+ continue
65
+ baseline = float(np.percentile(values, 35))
66
+ threshold = max(9.0, baseline + 6.0)
67
+ initial = np.zeros((height, width), dtype=np.uint8)
68
+ initial[eligible & (delta >= threshold * 0.65)] = cv2.GC_PR_FGD
69
+ initial[eligible & (delta >= threshold * 1.35)] = cv2.GC_FGD
70
+ initial[(container_mask > 0) & (initial == 0)] = cv2.GC_BGD
71
+ if np.count_nonzero(initial == cv2.GC_FGD) < 8:
72
+ continue
73
+ mask = refine_opencv(rgb, initial, container_mask)
74
+ if extractor == BitmapExtractor.SAM2:
75
+ from .sam2 import refine_with_sam2
76
+
77
+ mask = refine_with_sam2(rgb, mask, container_mask, device)
78
+ mask[text_exclusion > 0] = 0
79
+ mask[vector_foreground > 0] = 0
80
+ mask = filter_components(mask, max(8, int(width * height * 0.00001)))
81
+ if cv2.countNonZero(mask) < 12:
82
+ continue
83
+ alpha = soft_alpha(delta, mask, threshold)
84
+ x, y, w, h = cv2.boundingRect((alpha > 0).astype(np.uint8))
85
+ if w <= 0 or h <= 0:
86
+ continue
87
+ padding = 2
88
+ x0, y0 = max(0, x - padding), max(0, y - padding)
89
+ x1, y1 = min(width, x + w + padding), min(height, y + h + padding)
90
+ crop = np.dstack([rgb[y0:y1, x0:x1], alpha[y0:y1, x0:x1]])
91
+ asset_name = f"bitmap-{len(regions) + 1:04d}.png"
92
+ region = BitmapRegion(
93
+ id=f"bitmap-{len(regions) + 1:04d}",
94
+ asset_name=asset_name,
95
+ box=Box(x=x0, y=y0, width=x1 - x0, height=y1 - y0),
96
+ z_index=500,
97
+ )
98
+ regions.append(region)
99
+ assets[region.id] = crop
100
+ return regions, assets
101
+
102
+
103
+ def refine_opencv(rgb: np.ndarray, initial: np.ndarray, container_mask: np.ndarray) -> np.ndarray:
104
+ mask = initial.copy()
105
+ try:
106
+ cv2.grabCut(
107
+ cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR),
108
+ mask,
109
+ None,
110
+ np.zeros((1, 65), np.float64),
111
+ np.zeros((1, 65), np.float64),
112
+ 2,
113
+ cv2.GC_INIT_WITH_MASK,
114
+ )
115
+ output = np.isin(mask, [cv2.GC_FGD, cv2.GC_PR_FGD]).astype(np.uint8) * 255
116
+ except cv2.error:
117
+ output = (initial >= cv2.GC_PR_FGD).astype(np.uint8) * 255
118
+ output[container_mask == 0] = 0
119
+ return output
120
+
121
+
122
+ def filter_components(mask: np.ndarray, minimum_area: int) -> np.ndarray:
123
+ count, labels, stats, _ = cv2.connectedComponentsWithStats((mask > 0).astype(np.uint8), 8)
124
+ output = np.zeros_like(mask, dtype=np.uint8)
125
+ for index in range(1, count):
126
+ if stats[index, cv2.CC_STAT_AREA] >= minimum_area:
127
+ output[labels == index] = 255
128
+ return output
129
+
130
+
131
+ def soft_alpha(delta: np.ndarray, mask: np.ndarray, threshold: float) -> np.ndarray:
132
+ strength = np.clip((delta - threshold * 0.45) / max(threshold, 1), 0, 1)
133
+ alpha = np.rint(strength * 255).astype(np.uint8)
134
+ alpha[mask == 0] = 0
135
+ alpha = cv2.GaussianBlur(alpha, (3, 3), 0.65)
136
+ alpha[cv2.dilate(mask, np.ones((3, 3), np.uint8)) == 0] = 0
137
+ return alpha
138
+
139
+
140
+ def render_fill(shape: ShapeRegion, width: int, height: int) -> np.ndarray:
141
+ paint = shape.style.fill
142
+ if isinstance(paint, SolidPaint):
143
+ color = parse_hex(paint.color)
144
+ return np.broadcast_to(color, (height, width, 3)).copy()
145
+ assert isinstance(paint, LinearGradientPaint)
146
+ first, last = paint.stops[0], paint.stops[-1]
147
+ color0 = parse_hex(first.color).astype(np.float32)
148
+ color1 = parse_hex(last.color).astype(np.float32)
149
+ radians = np.deg2rad(paint.angle)
150
+ direction = np.array([np.cos(radians), np.sin(radians)], dtype=np.float32)
151
+ yy, xx = np.mgrid[0:height, 0:width]
152
+ projection = xx * direction[0] + yy * direction[1]
153
+ minimum, maximum = float(projection.min()), float(projection.max())
154
+ t = (projection - minimum) / max(maximum - minimum, 1)
155
+ result = color0[None, None, :] * (1 - t[:, :, None]) + color1[None, None, :] * t[:, :, None]
156
+ return np.clip(np.rint(result), 0, 255).astype(np.uint8)
157
+
158
+
159
+ def reconstruct_background(
160
+ rgba: np.ndarray,
161
+ document: DocumentSpec,
162
+ combined_mask: np.ndarray,
163
+ inpaint_callback,
164
+ ) -> np.ndarray:
165
+ height, width = rgba.shape[:2]
166
+ working = rgba.copy()
167
+ container_union = np.zeros((height, width), dtype=np.uint8)
168
+ for shape in sorted(document.shapes, key=lambda item: geometry_area(item), reverse=True):
169
+ if not shape.is_container or not shape.remove_source:
170
+ continue
171
+ mask = shape_source_mask(shape, width, height)
172
+ if not cv2.countNonZero(mask):
173
+ continue
174
+ ring = cv2.subtract(
175
+ cv2.dilate(mask, np.ones((15, 15), np.uint8)),
176
+ cv2.dilate(mask, np.ones((3, 3), np.uint8)),
177
+ )
178
+ samples = working[:, :, :3][ring > 0]
179
+ fill = np.median(samples, axis=0) if samples.size else np.array([255, 255, 255])
180
+ working[:, :, :3][mask > 0] = np.rint(fill).astype(np.uint8)
181
+ container_union = cv2.bitwise_or(container_union, mask)
182
+ remaining = combined_mask.copy()
183
+ remaining[container_union > 0] = 0
184
+ if cv2.countNonZero(remaining):
185
+ working = inpaint_callback(working, remaining)
186
+ return working
187
+
188
+
189
+ def geometry_area(shape: ShapeRegion) -> float:
190
+ geometry = shape.geometry
191
+ if hasattr(geometry, "box"):
192
+ return geometry.box.width * geometry.box.height
193
+ if hasattr(geometry, "rx"):
194
+ return float(np.pi * geometry.rx * geometry.ry)
195
+ points = np.array([(point.x, point.y) for point in geometry.points], dtype=np.float32)
196
+ return abs(float(cv2.contourArea(points))) if len(points) >= 3 else 0
197
+
198
+
199
+ def parse_hex(value: str) -> np.ndarray:
200
+ value = value.lstrip("#")
201
+ if len(value) == 3:
202
+ value = "".join(character * 2 for character in value)
203
+ return np.array([int(value[index:index + 2], 16) for index in (0, 2, 4)], dtype=np.uint8)
src/editable_image/cli.py ADDED
@@ -0,0 +1,220 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from __future__ import annotations
2
+
3
+ import argparse
4
+ import importlib.util
5
+ import json
6
+ import sys
7
+ from pathlib import Path
8
+
9
+ from .image_io import InvalidImage
10
+ from .inpaint import InpaintUnavailable, install_lama, lama_installed, resolve_torch_device
11
+ from .models import BitmapExtractor, CleanupMode, DEFAULT_SNAP_ANGLES, Device, ProcessOptions
12
+ from .ocr import OcrUnavailable, is_ocr_available
13
+ from .pipeline import export_result, process_path
14
+ from .sam2 import (
15
+ ShapeExtractionUnavailable,
16
+ install_sam2,
17
+ sam2_dependencies_installed,
18
+ sam2_installed,
19
+ )
20
+
21
+
22
+ IMAGE_SUFFIXES = {".png", ".jpg", ".jpeg", ".webp"}
23
+
24
+
25
+ def build_parser() -> argparse.ArgumentParser:
26
+ parser = argparse.ArgumentParser(prog="editable-image")
27
+ subparsers = parser.add_subparsers(dest="command", required=True)
28
+
29
+ convert = subparsers.add_parser("convert", help="convert bitmap images")
30
+ convert.add_argument("inputs", nargs="+", type=Path)
31
+ convert.add_argument("--output", "-o", type=Path, required=True)
32
+ convert.add_argument("--recursive", action="store_true")
33
+ convert.add_argument("--cleanup", choices=[item.value for item in CleanupMode], default="opencv")
34
+ convert.add_argument("--device", choices=[item.value for item in Device], default="auto")
35
+ convert.add_argument("--confidence", type=float, default=0.5)
36
+ convert.add_argument(
37
+ "--snap-angles",
38
+ type=parse_angles,
39
+ default=list(DEFAULT_SNAP_ANGLES),
40
+ metavar="DEGREES",
41
+ help="comma-separated target angles (default: 0,45,90,-45,-90)",
42
+ )
43
+ convert.add_argument(
44
+ "--snap-tolerance",
45
+ type=float,
46
+ default=6.0,
47
+ metavar="DEGREES",
48
+ help="maximum distance from a target angle (default: 6)",
49
+ )
50
+ convert.add_argument(
51
+ "--snap-font-sizes",
52
+ action="store_true",
53
+ help="snap estimated font sizes to peaks in the image's size distribution",
54
+ )
55
+ convert.add_argument(
56
+ "--embed-fonts",
57
+ action="store_true",
58
+ help="embed used font faces in the SVG (default: link system fonts)",
59
+ )
60
+ convert.add_argument(
61
+ "--vectorize-shapes",
62
+ action="store_true",
63
+ help="promote confident diagram primitives into editable SVG shapes",
64
+ )
65
+ convert.add_argument(
66
+ "--bitmap-extractor",
67
+ choices=[item.value for item in BitmapExtractor],
68
+ default=BitmapExtractor.OPENCV.value,
69
+ help="bitmap-layer mask extractor used with --vectorize-shapes (default: opencv)",
70
+ )
71
+ convert.add_argument("--overwrite", action="store_true")
72
+
73
+ models = subparsers.add_parser("models", help="manage optional model files")
74
+ model_commands = models.add_subparsers(dest="model_command", required=True)
75
+ install = model_commands.add_parser("install")
76
+ install.add_argument("model", choices=["ocr", "lama", "sam2"])
77
+ model_commands.add_parser("status")
78
+ return parser
79
+
80
+
81
+ def main(argv: list[str] | None = None) -> int:
82
+ args = build_parser().parse_args(argv)
83
+ if args.command == "models":
84
+ return model_command(args)
85
+ return convert_command(args)
86
+
87
+
88
+ def model_command(args: argparse.Namespace) -> int:
89
+ if args.model_command == "install":
90
+ if args.model == "lama":
91
+ path = install_lama()
92
+ print(f"Installed LaMa at {path}")
93
+ return 0
94
+ if args.model == "sam2":
95
+ path = install_sam2()
96
+ print(f"Installed SAM 2 at {path}")
97
+ return 0
98
+ if not is_ocr_available():
99
+ print("Install the cpu or cuda project extra before installing OCR models", file=sys.stderr)
100
+ return 2
101
+ from .ocr import RapidOcrEngine
102
+
103
+ RapidOcrEngine(Device.CPU)
104
+ print("OCR models are ready")
105
+ return 0
106
+ print(
107
+ json.dumps(
108
+ {
109
+ "ocr": is_ocr_available(),
110
+ "lama": lama_installed(),
111
+ "sam2": sam2_installed(),
112
+ "sam2_dependencies": sam2_dependencies_installed(),
113
+ "devices": available_devices(),
114
+ },
115
+ indent=2,
116
+ )
117
+ )
118
+ return 0
119
+
120
+
121
+ def convert_command(args: argparse.Namespace) -> int:
122
+ if args.bitmap_extractor != BitmapExtractor.OPENCV.value and not args.vectorize_shapes:
123
+ print("--bitmap-extractor requires --vectorize-shapes", file=sys.stderr)
124
+ return 2
125
+ options = ProcessOptions(
126
+ confidence=args.confidence,
127
+ cleanup=CleanupMode(args.cleanup),
128
+ device=Device(args.device),
129
+ snap_angles=args.snap_angles,
130
+ snap_tolerance=args.snap_tolerance,
131
+ snap_font_sizes=args.snap_font_sizes,
132
+ vectorize_shapes=args.vectorize_shapes,
133
+ bitmap_extractor=BitmapExtractor(args.bitmap_extractor),
134
+ )
135
+ inputs = collect_inputs(args.inputs, args.recursive)
136
+ if not inputs:
137
+ print("No supported images found", file=sys.stderr)
138
+ return 2
139
+ failures = 0
140
+ for source in inputs:
141
+ destination = args.output / source.stem
142
+ expected = [destination / f"{source.stem}-background.png", destination / f"{source.stem}-overlay.svg"]
143
+ if not args.overwrite and any(path.exists() for path in expected):
144
+ failures += 1
145
+ print(f"skip {source}: output exists (use --overwrite)", file=sys.stderr)
146
+ continue
147
+ try:
148
+ result = process_path(source, options)
149
+ background, overlay, assets = export_result(
150
+ result,
151
+ destination,
152
+ source.stem,
153
+ embed_fonts=args.embed_fonts,
154
+ )
155
+ print(
156
+ json.dumps(
157
+ {
158
+ "source": str(source),
159
+ "background": str(background),
160
+ "overlay": str(overlay),
161
+ "assets": [str(path) for path in assets],
162
+ }
163
+ )
164
+ )
165
+ except (
166
+ InvalidImage,
167
+ OcrUnavailable,
168
+ InpaintUnavailable,
169
+ ShapeExtractionUnavailable,
170
+ OSError,
171
+ ValueError,
172
+ ) as exc:
173
+ failures += 1
174
+ print(f"failed {source}: {exc}", file=sys.stderr)
175
+ return 1 if failures else 0
176
+
177
+
178
+ def collect_inputs(paths: list[Path], recursive: bool) -> list[Path]:
179
+ output: list[Path] = []
180
+ for path in paths:
181
+ if path.is_file() and path.suffix.lower() in IMAGE_SUFFIXES:
182
+ output.append(path)
183
+ elif path.is_dir():
184
+ iterator = path.rglob("*") if recursive else path.glob("*")
185
+ output.extend(item for item in iterator if item.is_file() and item.suffix.lower() in IMAGE_SUFFIXES)
186
+ return sorted(set(output))
187
+
188
+
189
+ def parse_angles(value: str) -> list[float]:
190
+ try:
191
+ angles = [float(item.strip()) for item in value.split(",") if item.strip()]
192
+ except ValueError as exc:
193
+ raise argparse.ArgumentTypeError("angles must be comma-separated numbers") from exc
194
+ if not angles:
195
+ raise argparse.ArgumentTypeError("at least one snap angle is required")
196
+ if any(angle < -180 or angle > 180 for angle in angles):
197
+ raise argparse.ArgumentTypeError("snap angles must be between -180 and 180")
198
+ return angles
199
+
200
+
201
+ def available_devices() -> list[str]:
202
+ devices = ["cpu"]
203
+ try:
204
+ import onnxruntime as ort
205
+
206
+ if "CUDAExecutionProvider" in ort.get_available_providers():
207
+ devices.append("cuda")
208
+ except ImportError:
209
+ pass
210
+ if importlib.util.find_spec("torch") is not None:
211
+ try:
212
+ if resolve_torch_device(Device.AUTO) == "mps":
213
+ devices.append("mps")
214
+ except InpaintUnavailable:
215
+ pass
216
+ return devices
217
+
218
+
219
+ if __name__ == "__main__":
220
+ raise SystemExit(main())
src/editable_image/fonts.py ADDED
@@ -0,0 +1,69 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from __future__ import annotations
2
+
3
+ import base64
4
+ from io import BytesIO
5
+ from dataclasses import dataclass
6
+ from functools import lru_cache
7
+ from importlib.resources import files
8
+ from pathlib import Path
9
+
10
+
11
+ @dataclass(frozen=True, slots=True)
12
+ class FontFace:
13
+ family: str
14
+ weight: int
15
+ italic: bool
16
+ filename: str
17
+
18
+ @property
19
+ def path(self) -> Path:
20
+ return Path(str(files("editable_image").joinpath("fonts", self.filename)))
21
+
22
+
23
+ FACES = (
24
+ FontFace("Lato", 400, False, "Lato-Regular.ttf"),
25
+ FontFace("Lato", 400, True, "Lato-Italic.ttf"),
26
+ FontFace("Lato", 700, False, "Lato-Bold.ttf"),
27
+ FontFace("Lato", 700, True, "Lato-BoldItalic.ttf"),
28
+ FontFace("Liberation Sans", 400, False, "LiberationSans-Regular.ttf"),
29
+ FontFace("Liberation Sans", 400, True, "LiberationSans-Italic.ttf"),
30
+ FontFace("Liberation Sans", 700, False, "LiberationSans-Bold.ttf"),
31
+ FontFace("Liberation Sans", 700, True, "LiberationSans-BoldItalic.ttf"),
32
+ FontFace("Liberation Serif", 400, False, "LiberationSerif-Regular.ttf"),
33
+ FontFace("Liberation Serif", 400, True, "LiberationSerif-Italic.ttf"),
34
+ FontFace("Liberation Serif", 700, False, "LiberationSerif-Bold.ttf"),
35
+ FontFace("Liberation Serif", 700, True, "LiberationSerif-BoldItalic.ttf"),
36
+ FontFace("Liberation Mono", 400, False, "LiberationMono-Regular.ttf"),
37
+ FontFace("Liberation Mono", 400, True, "LiberationMono-Italic.ttf"),
38
+ FontFace("Liberation Mono", 700, False, "LiberationMono-Bold.ttf"),
39
+ FontFace("Liberation Mono", 700, True, "LiberationMono-BoldItalic.ttf"),
40
+ )
41
+
42
+
43
+ def families() -> list[str]:
44
+ return list(dict.fromkeys(face.family for face in FACES if face.path.exists()))
45
+
46
+
47
+ def get_face(family: str, weight: int, italic: bool) -> FontFace:
48
+ candidates = [face for face in FACES if face.family == family and face.path.exists()]
49
+ if not candidates:
50
+ candidates = [face for face in FACES if face.path.exists()]
51
+ if not candidates:
52
+ raise FileNotFoundError("No bundled fonts are installed")
53
+ return min(
54
+ candidates,
55
+ key=lambda face: abs(face.weight - weight) + (1000 if face.italic != italic else 0),
56
+ )
57
+
58
+
59
+ @lru_cache(maxsize=32)
60
+ def font_data_uri(filename: str) -> str:
61
+ path = Path(str(files("editable_image").joinpath("fonts", filename)))
62
+ from fontTools.ttLib import TTFont
63
+
64
+ font = TTFont(path)
65
+ font.flavor = "woff2"
66
+ output = BytesIO()
67
+ font.save(output)
68
+ data = base64.b64encode(output.getvalue()).decode("ascii")
69
+ return f"data:font/woff2;base64,{data}"
src/editable_image/fonts/LICENSES.txt ADDED
@@ -0,0 +1,10 @@
 
 
 
 
 
 
 
 
 
 
 
1
+ Lato
2
+ Copyright (c) 2010-2014 by tyPoland Lukasz Dziedzic with Reserved Font Name "Lato".
3
+ Licensed under the SIL Open Font License, Version 1.1.
4
+ https://scripts.sil.org/OFL
5
+
6
+ Liberation Fonts
7
+ Copyright (c) 2007 Red Hat, Inc. with Reserved Font Name "Liberation".
8
+ Licensed under the SIL Open Font License, Version 1.1.
9
+ https://scripts.sil.org/OFL
10
+
src/editable_image/fonts/Lato-Bold.ttf ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:bf1b8130069b44b9148eeece35e5423bedac49777ba746615b826b8276574a7b
3
+ size 657188
src/editable_image/fonts/Lato-BoldItalic.ttf ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:e2f9244fdcd20840f2480f91370cea885bf99230c49d153199fa64853ad74b69
3
+ size 699008
src/editable_image/fonts/Lato-Italic.ttf ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:8c863f833b12bc8486671d26e3a58b13b0f9eaff1d6b10398e7c107a9b7ee05a
3
+ size 723544
src/editable_image/fonts/Lato-Regular.ttf ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:6f6940be0835c3ddec9199e5fc42be4cbc61ebcfd58c623fdf719366253f1780
3
+ size 657212
src/editable_image/fonts/LiberationMono-Bold.ttf ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:626655e94dd82f3f42549daf995c921b0915fa8ab1f4b839559e8892ea41d240
3
+ size 308068
src/editable_image/fonts/LiberationMono-BoldItalic.ttf ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:15eb161953e3ecc7fc05a3fec8a59e0f4e0a54a4e375736adf8e98d65981f813
3
+ size 284184
src/editable_image/fonts/LiberationMono-Italic.ttf ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:a71b2c25c89da05cf0e7c4dbba8d473fdead0b181ae56165217747d2c1f39215
3
+ size 281608
src/editable_image/fonts/LiberationMono-Regular.ttf ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:395fa5ab8d40c8eba390ced528744ea75a7f69aabf3e68b6f925ca0e39a27370
3
+ size 319624
src/editable_image/fonts/LiberationSans-Bold.ttf ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:3973aa5054fb467dd5627245d3dc82e37bf16fe075756156a570455871351582
3
+ size 414568
src/editable_image/fonts/LiberationSans-BoldItalic.ttf ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:c80fa7f2ffa0e01d4d8dcd6a6d1e43eda665222d1b4db597dde1174c456006cf
3
+ size 409108
src/editable_image/fonts/LiberationSans-Italic.ttf ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:830c5fa600505fb4c1a271b4271c53c44bae43f492b2a240d0e98a3a7a380121
3
+ size 415920
src/editable_image/fonts/LiberationSans-Regular.ttf ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:4659bc0c58c5028dd488ec928d41d9265db43d9b669fc14ca8b0832daca7b144
3
+ size 410820
src/editable_image/fonts/LiberationSerif-Bold.ttf ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:9e66c25e20868756bf12eb24a520a4552727c9cd108f577f7ae33e3c0110e39c
3
+ size 370196
src/editable_image/fonts/LiberationSerif-BoldItalic.ttf ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:379010e87421a883f7bbfa7936d23dfe0257a54ce850819b74556e9f2c615a8a
3
+ size 376892
src/editable_image/fonts/LiberationSerif-Italic.ttf ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:eaeaf7f2b12544ecde64290c1fa7ec403ff70fd58b313637b265400b20249242
3
+ size 375760
src/editable_image/fonts/LiberationSerif-Regular.ttf ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:705903ae1382f0150f115d6f10a30f5b2f3d6ade649c474afe6366c39174fb65
3
+ size 393692
src/editable_image/image_io.py ADDED
@@ -0,0 +1,73 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from __future__ import annotations
2
+
3
+ from dataclasses import dataclass
4
+ from io import BytesIO
5
+ from pathlib import Path
6
+
7
+ import numpy as np
8
+ from PIL import Image, ImageOps, UnidentifiedImageError
9
+
10
+
11
+ ALLOWED_FORMATS = {"PNG", "JPEG", "WEBP"}
12
+ MAX_FILE_BYTES = 50 * 1024 * 1024
13
+
14
+
15
+ class InvalidImage(ValueError):
16
+ pass
17
+
18
+
19
+ @dataclass(slots=True)
20
+ class LoadedImage:
21
+ rgba: np.ndarray
22
+ icc_profile: bytes | None
23
+ source_format: str
24
+
25
+ @property
26
+ def width(self) -> int:
27
+ return int(self.rgba.shape[1])
28
+
29
+ @property
30
+ def height(self) -> int:
31
+ return int(self.rgba.shape[0])
32
+
33
+
34
+ def load_image_bytes(data: bytes, max_pixels: int = 40_000_000) -> LoadedImage:
35
+ if not data:
36
+ raise InvalidImage("The image is empty")
37
+ if len(data) > MAX_FILE_BYTES:
38
+ raise InvalidImage("The image exceeds the 50 MB limit")
39
+
40
+ try:
41
+ with Image.open(BytesIO(data)) as source:
42
+ if source.format not in ALLOWED_FORMATS:
43
+ raise InvalidImage("Only PNG, JPEG, and WebP images are supported")
44
+ if getattr(source, "is_animated", False):
45
+ raise InvalidImage("Animated images are not supported")
46
+ source.load()
47
+ image = ImageOps.exif_transpose(source)
48
+ if image.width * image.height > max_pixels:
49
+ raise InvalidImage(f"The decoded image exceeds {max_pixels:,} pixels")
50
+ rgba = np.asarray(image.convert("RGBA"), dtype=np.uint8).copy()
51
+ return LoadedImage(
52
+ rgba=rgba,
53
+ icc_profile=source.info.get("icc_profile"),
54
+ source_format=source.format,
55
+ )
56
+ except (UnidentifiedImageError, OSError) as exc:
57
+ raise InvalidImage("The uploaded file is not a readable image") from exc
58
+
59
+
60
+ def load_image_path(path: Path, max_pixels: int = 40_000_000) -> LoadedImage:
61
+ try:
62
+ return load_image_bytes(path.read_bytes(), max_pixels=max_pixels)
63
+ except OSError as exc:
64
+ raise InvalidImage(f"Could not read {path}") from exc
65
+
66
+
67
+ def save_png(image: LoadedImage | np.ndarray, path: Path, icc_profile: bytes | None = None) -> None:
68
+ rgba = image.rgba if isinstance(image, LoadedImage) else image
69
+ if isinstance(image, LoadedImage):
70
+ icc_profile = image.icc_profile
71
+ target = Image.fromarray(rgba, "RGBA")
72
+ target.save(path, format="PNG", icc_profile=icc_profile, optimize=True)
73
+
src/editable_image/inpaint.py ADDED
@@ -0,0 +1,103 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from __future__ import annotations
2
+
3
+ import hashlib
4
+ import importlib.util
5
+ import urllib.request
6
+ from pathlib import Path
7
+
8
+ import cv2
9
+ import numpy as np
10
+ from platformdirs import user_cache_path
11
+
12
+ from .models import CleanupMode, Device
13
+
14
+
15
+ LAMA_URL = "https://github.com/Sanster/models/releases/download/add_big_lama/big-lama.pt"
16
+ LAMA_MD5 = "e3aa4aaa15225a33ec84f9f4bc47e500"
17
+
18
+
19
+ class InpaintUnavailable(RuntimeError):
20
+ pass
21
+
22
+
23
+ def lama_path() -> Path:
24
+ return user_cache_path("editable-image") / "models" / "big-lama.pt"
25
+
26
+
27
+ def lama_installed() -> bool:
28
+ path = lama_path()
29
+ return path.exists() and file_md5(path) == LAMA_MD5
30
+
31
+
32
+ def file_md5(path: Path) -> str:
33
+ digest = hashlib.md5(usedforsecurity=False)
34
+ with path.open("rb") as stream:
35
+ for chunk in iter(lambda: stream.read(1024 * 1024), b""):
36
+ digest.update(chunk)
37
+ return digest.hexdigest()
38
+
39
+
40
+ def install_lama() -> Path:
41
+ target = lama_path()
42
+ target.parent.mkdir(parents=True, exist_ok=True)
43
+ temporary = target.with_suffix(".download")
44
+ urllib.request.urlretrieve(LAMA_URL, temporary)
45
+ if file_md5(temporary) != LAMA_MD5:
46
+ temporary.unlink(missing_ok=True)
47
+ raise InpaintUnavailable("The downloaded LaMa model failed checksum verification")
48
+ temporary.replace(target)
49
+ return target
50
+
51
+
52
+ def inpaint(rgba: np.ndarray, mask: np.ndarray, mode: CleanupMode, device: Device) -> np.ndarray:
53
+ if not np.any(mask):
54
+ return rgba.copy()
55
+ if mode == CleanupMode.LAMA:
56
+ return lama_inpaint(rgba, mask, device)
57
+ return opencv_inpaint(rgba, mask)
58
+
59
+
60
+ def opencv_inpaint(rgba: np.ndarray, mask: np.ndarray) -> np.ndarray:
61
+ rgb = rgba[:, :, :3]
62
+ alpha = rgba[:, :, 3]
63
+ radius = max(2, min(9, int(round(np.sqrt(np.count_nonzero(mask)) / 120))))
64
+ restored_rgb = cv2.inpaint(cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR), mask, radius, cv2.INPAINT_TELEA)
65
+ restored_alpha = cv2.inpaint(alpha, mask, radius, cv2.INPAINT_TELEA)
66
+ return np.dstack([cv2.cvtColor(restored_rgb, cv2.COLOR_BGR2RGB), restored_alpha])
67
+
68
+
69
+ def resolve_torch_device(requested: Device) -> str:
70
+ if importlib.util.find_spec("torch") is None:
71
+ raise InpaintUnavailable("Install the lama project extra to use neural cleanup")
72
+ import torch
73
+
74
+ if requested == Device.CPU:
75
+ return "cpu"
76
+ if requested in {Device.AUTO, Device.CUDA} and torch.cuda.is_available():
77
+ return "cuda"
78
+ if requested in {Device.AUTO, Device.MPS} and torch.backends.mps.is_available():
79
+ return "mps"
80
+ if requested in {Device.CUDA, Device.MPS}:
81
+ raise InpaintUnavailable(f"The requested {requested.value} device is unavailable")
82
+ return "cpu"
83
+
84
+
85
+ def lama_inpaint(rgba: np.ndarray, mask: np.ndarray, device: Device) -> np.ndarray:
86
+ if not lama_installed():
87
+ raise InpaintUnavailable("LaMa is not installed. Run 'editable-image models install lama'.")
88
+ import torch
89
+
90
+ target_device = resolve_torch_device(device)
91
+ model = torch.jit.load(str(lama_path()), map_location=target_device).eval()
92
+ height, width = rgba.shape[:2]
93
+ pad_height = (8 - height % 8) % 8
94
+ pad_width = (8 - width % 8) % 8
95
+ rgb = np.pad(rgba[:, :, :3], ((0, pad_height), (0, pad_width), (0, 0)), mode="reflect")
96
+ padded_mask = np.pad(mask, ((0, pad_height), (0, pad_width)), mode="constant")
97
+ image_tensor = torch.from_numpy(rgb.transpose(2, 0, 1)).float().div(255).unsqueeze(0).to(target_device)
98
+ mask_tensor = torch.from_numpy((padded_mask > 0).astype(np.float32)).unsqueeze(0).unsqueeze(0).to(target_device)
99
+ with torch.inference_mode():
100
+ restored = model(image_tensor, mask_tensor)[0].permute(1, 2, 0).clamp(0, 1)
101
+ result = (restored.cpu().numpy()[:height, :width] * 255).astype(np.uint8)
102
+ return np.dstack([result, rgba[:, :, 3]])
103
+
src/editable_image/masking.py ADDED
@@ -0,0 +1,94 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from __future__ import annotations
2
+
3
+ import cv2
4
+ import numpy as np
5
+
6
+ from .models import DocumentSpec, MaskStroke, Point, TextRegion
7
+ from .shape_recognition import shape_source_mask
8
+
9
+
10
+ def build_mask(rgba: np.ndarray, document: DocumentSpec) -> np.ndarray:
11
+ height, width = rgba.shape[:2]
12
+ combined = np.zeros((height, width), dtype=np.uint8)
13
+ for region in document.regions:
14
+ if region.remove_text:
15
+ combined = cv2.bitwise_or(combined, region_mask(rgba, region))
16
+ for shape in document.shapes:
17
+ if shape.remove_source:
18
+ shape_mask = shape_source_mask(shape, width, height)
19
+ padding = max(0, int(round(shape.mask_padding)))
20
+ if padding:
21
+ kernel = cv2.getStructuringElement(
22
+ cv2.MORPH_ELLIPSE,
23
+ (padding * 2 + 1, padding * 2 + 1),
24
+ )
25
+ shape_mask = cv2.dilate(shape_mask, kernel)
26
+ combined = cv2.bitwise_or(combined, shape_mask)
27
+ for stroke in document.mask_strokes:
28
+ apply_stroke(combined, stroke)
29
+ return combined
30
+
31
+
32
+ def region_mask(rgba: np.ndarray, region: TextRegion) -> np.ndarray:
33
+ height, width = rgba.shape[:2]
34
+ polygon = np.array([[point.x, point.y] for point in region.source_quad], dtype=np.int32)
35
+ bounds = cv2.boundingRect(polygon)
36
+ x, y, w, h = bounds
37
+ x0, y0 = max(0, x), max(0, y)
38
+ x1, y1 = min(width, x + w), min(height, y + h)
39
+ output = np.zeros((height, width), dtype=np.uint8)
40
+ if x1 <= x0 or y1 <= y0:
41
+ return output
42
+
43
+ crop = rgba[y0:y1, x0:x1, :3]
44
+ lab = cv2.cvtColor(crop, cv2.COLOR_RGB2LAB).astype(np.float32)
45
+ local_polygon = polygon - np.array([x0, y0])
46
+ polygon_mask = np.zeros(crop.shape[:2], dtype=np.uint8)
47
+ cv2.fillPoly(polygon_mask, [local_polygon], 255)
48
+
49
+ perimeter = cv2.morphologyEx(polygon_mask, cv2.MORPH_GRADIENT, np.ones((3, 3), np.uint8))
50
+ border_pixels = lab[perimeter > 0]
51
+ background = np.median(border_pixels, axis=0) if border_pixels.size else np.median(lab.reshape(-1, 3), axis=0)
52
+ delta = np.linalg.norm(lab - background, axis=2)
53
+ inside = delta[polygon_mask > 0]
54
+ if inside.size:
55
+ otsu_threshold, _ = cv2.threshold(
56
+ np.clip(inside, 0, 255).astype(np.uint8),
57
+ 0,
58
+ 255,
59
+ cv2.THRESH_BINARY + cv2.THRESH_OTSU,
60
+ )
61
+ threshold = max(8.0, float(otsu_threshold) * 0.72)
62
+ else:
63
+ threshold = 8.0
64
+ foreground = ((delta >= threshold) & (polygon_mask > 0)).astype(np.uint8) * 255
65
+
66
+ count, labels, stats, _ = cv2.connectedComponentsWithStats(foreground, 8)
67
+ filtered = np.zeros_like(foreground)
68
+ max_area = max(9, int(w * h * 0.22))
69
+ min_area = 1
70
+ for index in range(1, count):
71
+ area = stats[index, cv2.CC_STAT_AREA]
72
+ component_height = stats[index, cv2.CC_STAT_HEIGHT]
73
+ if min_area <= area <= max_area and component_height >= 1:
74
+ filtered[labels == index] = 255
75
+ if cv2.countNonZero(filtered) < max(3, w * h * 0.005):
76
+ filtered = foreground
77
+
78
+ padding = max(1, int(round(region.mask_padding)))
79
+ kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (padding * 2 + 1, padding * 2 + 1))
80
+ filtered = cv2.dilate(filtered, kernel)
81
+ output[y0:y1, x0:x1] = filtered
82
+ return output
83
+
84
+
85
+ def apply_stroke(mask: np.ndarray, stroke: MaskStroke) -> None:
86
+ value = 255 if stroke.operation == "add" else 0
87
+ radius = max(1, int(round(stroke.radius)))
88
+ points = [(int(round(point.x)), int(round(point.y))) for point in stroke.points]
89
+ if len(points) == 1:
90
+ cv2.circle(mask, points[0], radius, value, -1, lineType=cv2.LINE_AA)
91
+ return
92
+ for start, end in zip(points, points[1:]):
93
+ cv2.line(mask, start, end, value, radius * 2, lineType=cv2.LINE_AA)
94
+ cv2.circle(mask, end, radius, value, -1, lineType=cv2.LINE_AA)
src/editable_image/models.py ADDED
@@ -0,0 +1,244 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from __future__ import annotations
2
+
3
+ from enum import StrEnum
4
+ from typing import Annotated, Literal
5
+
6
+ from pydantic import BaseModel, ConfigDict, Field, field_validator
7
+
8
+
9
+ Pixel = Annotated[float, Field(ge=0)]
10
+ DEFAULT_SNAP_ANGLES = (0.0, 45.0, 90.0, -45.0, -90.0)
11
+
12
+
13
+ class CleanupMode(StrEnum):
14
+ OPENCV = "opencv"
15
+ LAMA = "lama"
16
+
17
+
18
+ class Device(StrEnum):
19
+ AUTO = "auto"
20
+ CPU = "cpu"
21
+ CUDA = "cuda"
22
+ MPS = "mps"
23
+
24
+
25
+ class BitmapExtractor(StrEnum):
26
+ OPENCV = "opencv"
27
+ SAM2 = "sam2"
28
+
29
+
30
+ class Point(BaseModel):
31
+ x: Pixel
32
+ y: Pixel
33
+
34
+
35
+ class Box(BaseModel):
36
+ x: Pixel
37
+ y: Pixel
38
+ width: Annotated[float, Field(gt=0)]
39
+ height: Annotated[float, Field(gt=0)]
40
+
41
+
42
+ class TextStyle(BaseModel):
43
+ font_family: str = "Lato"
44
+ font_size: Annotated[float, Field(ge=4, le=1000)] = 16
45
+ font_weight: Literal[400, 700] = 400
46
+ italic: bool = False
47
+ fill: str = "#111111"
48
+ opacity: Annotated[float, Field(ge=0, le=1)] = 1
49
+ stroke: str = "none"
50
+ stroke_width: Annotated[float, Field(ge=0, le=50)] = 0
51
+ letter_spacing: Annotated[float, Field(ge=-20, le=100)] = 0
52
+ anchor: Literal["start", "middle", "end"] = "start"
53
+
54
+ @field_validator("fill", "stroke")
55
+ @classmethod
56
+ def validate_color(cls, value: str) -> str:
57
+ if value == "none" or (
58
+ len(value) in {4, 7, 9} and value.startswith("#")
59
+ ):
60
+ return value
61
+ raise ValueError("expected a CSS hexadecimal color or 'none'")
62
+
63
+
64
+ class TextRegion(BaseModel):
65
+ model_config = ConfigDict(extra="forbid")
66
+
67
+ id: str
68
+ text: str
69
+ confidence: Annotated[float, Field(ge=0, le=1)]
70
+ source_quad: Annotated[list[Point], Field(min_length=4, max_length=4)]
71
+ overlay_box: Box
72
+ rotation: Annotated[float, Field(ge=-180, le=180)] = 0
73
+ remove_text: bool = True
74
+ include_overlay: bool = True
75
+ mask_padding: Annotated[float, Field(ge=0, le=100)] = 3
76
+ style: TextStyle = Field(default_factory=TextStyle)
77
+ z_index: int = 1000
78
+
79
+
80
+ class SolidPaint(BaseModel):
81
+ model_config = ConfigDict(extra="forbid")
82
+
83
+ type: Literal["solid"] = "solid"
84
+ color: str
85
+
86
+ @field_validator("color")
87
+ @classmethod
88
+ def validate_color(cls, value: str) -> str:
89
+ if len(value) in {4, 7, 9} and value.startswith("#"):
90
+ return value
91
+ raise ValueError("expected a CSS hexadecimal color")
92
+
93
+
94
+ class GradientStop(BaseModel):
95
+ offset: Annotated[float, Field(ge=0, le=1)]
96
+ color: str
97
+ opacity: Annotated[float, Field(ge=0, le=1)] = 1
98
+
99
+ @field_validator("color")
100
+ @classmethod
101
+ def validate_color(cls, value: str) -> str:
102
+ if len(value) in {4, 7, 9} and value.startswith("#"):
103
+ return value
104
+ raise ValueError("expected a CSS hexadecimal color")
105
+
106
+
107
+ class LinearGradientPaint(BaseModel):
108
+ model_config = ConfigDict(extra="forbid")
109
+
110
+ type: Literal["linear-gradient"] = "linear-gradient"
111
+ angle: Annotated[float, Field(ge=-180, le=180)] = 0
112
+ stops: Annotated[list[GradientStop], Field(min_length=2, max_length=8)]
113
+
114
+
115
+ Paint = SolidPaint | LinearGradientPaint
116
+
117
+
118
+ class ShapeStyle(BaseModel):
119
+ fill: Paint = Field(default_factory=lambda: SolidPaint(color="#000000"))
120
+ fill_opacity: Annotated[float, Field(ge=0, le=1)] = 1
121
+ stroke: str = "none"
122
+ stroke_width: Annotated[float, Field(ge=0, le=100)] = 0
123
+ stroke_opacity: Annotated[float, Field(ge=0, le=1)] = 1
124
+ linecap: Literal["butt", "round", "square"] = "round"
125
+ linejoin: Literal["miter", "round", "bevel"] = "round"
126
+
127
+ @field_validator("stroke")
128
+ @classmethod
129
+ def validate_stroke(cls, value: str) -> str:
130
+ if value == "none" or (len(value) in {4, 7, 9} and value.startswith("#")):
131
+ return value
132
+ raise ValueError("expected a CSS hexadecimal color or 'none'")
133
+
134
+
135
+ class RectGeometry(BaseModel):
136
+ type: Literal["rect"] = "rect"
137
+ box: Box
138
+ rx: Annotated[float, Field(ge=0)] = 0
139
+ ry: Annotated[float, Field(ge=0)] = 0
140
+ rotation: Annotated[float, Field(ge=-180, le=180)] = 0
141
+
142
+
143
+ class EllipseGeometry(BaseModel):
144
+ type: Literal["ellipse"] = "ellipse"
145
+ cx: Pixel
146
+ cy: Pixel
147
+ rx: Annotated[float, Field(gt=0)]
148
+ ry: Annotated[float, Field(gt=0)]
149
+ rotation: Annotated[float, Field(ge=-180, le=180)] = 0
150
+
151
+
152
+ class PolygonGeometry(BaseModel):
153
+ type: Literal["polygon"] = "polygon"
154
+ points: Annotated[list[Point], Field(min_length=3)]
155
+
156
+
157
+ class PathGeometry(BaseModel):
158
+ type: Literal["path"] = "path"
159
+ points: Annotated[list[Point], Field(min_length=2)]
160
+ closed: bool = False
161
+ smooth: bool = False
162
+ arrow_start: bool = False
163
+ arrow_end: bool = False
164
+
165
+
166
+ ShapeGeometry = RectGeometry | EllipseGeometry | PolygonGeometry | PathGeometry
167
+
168
+
169
+ class ShapeRegion(BaseModel):
170
+ model_config = ConfigDict(extra="forbid")
171
+
172
+ id: str
173
+ confidence: Annotated[float, Field(ge=0, le=1)]
174
+ geometry: ShapeGeometry
175
+ source_contours: list[list[Point]] = Field(default_factory=list)
176
+ style: ShapeStyle = Field(default_factory=ShapeStyle)
177
+ remove_source: bool = True
178
+ include_overlay: bool = True
179
+ mask_padding: Annotated[float, Field(ge=0, le=100)] = 1
180
+ z_index: int = 100
181
+ is_container: bool = False
182
+ semantic_role: Literal["shape", "arrow"] = "shape"
183
+
184
+
185
+ class BitmapRegion(BaseModel):
186
+ model_config = ConfigDict(extra="forbid")
187
+
188
+ id: str
189
+ asset_name: str
190
+ box: Box
191
+ opacity: Annotated[float, Field(ge=0, le=1)] = 1
192
+ z_index: int = 500
193
+
194
+
195
+ class MaskStroke(BaseModel):
196
+ operation: Literal["add", "erase"]
197
+ radius: Annotated[float, Field(gt=0, le=500)]
198
+ points: Annotated[list[Point], Field(min_length=1)]
199
+
200
+
201
+ class Canvas(BaseModel):
202
+ width: Annotated[int, Field(gt=0)]
203
+ height: Annotated[int, Field(gt=0)]
204
+
205
+
206
+ class DocumentSpec(BaseModel):
207
+ model_config = ConfigDict(extra="forbid")
208
+
209
+ version: Literal[1, 2] = 1
210
+ source_name: str
211
+ canvas: Canvas
212
+ regions: list[TextRegion]
213
+ shapes: list[ShapeRegion] = Field(default_factory=list)
214
+ bitmap_regions: list[BitmapRegion] = Field(default_factory=list)
215
+ mask_strokes: list[MaskStroke] = Field(default_factory=list)
216
+
217
+
218
+ class ProcessOptions(BaseModel):
219
+ confidence: Annotated[float, Field(ge=0, le=1)] = 0.5
220
+ device: Device = Device.AUTO
221
+ cleanup: CleanupMode = CleanupMode.OPENCV
222
+ max_pixels: Annotated[int, Field(gt=0)] = 40_000_000
223
+ snap_angles: Annotated[
224
+ list[Annotated[float, Field(ge=-180, le=180)]],
225
+ Field(min_length=1),
226
+ ] = Field(default_factory=lambda: list(DEFAULT_SNAP_ANGLES))
227
+ snap_tolerance: Annotated[float, Field(ge=0, le=180)] = 6.0
228
+ snap_font_sizes: bool = False
229
+ vectorize_shapes: bool = False
230
+ bitmap_extractor: BitmapExtractor = BitmapExtractor.OPENCV
231
+
232
+
233
+ class JobResponse(BaseModel):
234
+ id: str
235
+ document: DocumentSpec
236
+ original_url: str
237
+ background_url: str
238
+
239
+
240
+ class Capabilities(BaseModel):
241
+ ocr_available: bool
242
+ lama_installed: bool
243
+ devices: list[str]
244
+ fonts: list[str]
src/editable_image/ocr.py ADDED
@@ -0,0 +1,92 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from __future__ import annotations
2
+
3
+ import importlib.util
4
+ import math
5
+ from dataclasses import dataclass
6
+
7
+ import numpy as np
8
+
9
+ from .models import Device, Point
10
+
11
+
12
+ class OcrUnavailable(RuntimeError):
13
+ pass
14
+
15
+
16
+ @dataclass(slots=True)
17
+ class OcrLine:
18
+ text: str
19
+ confidence: float
20
+ quad: list[Point]
21
+
22
+
23
+ def is_ocr_available() -> bool:
24
+ return importlib.util.find_spec("rapidocr") is not None and (
25
+ importlib.util.find_spec("onnxruntime") is not None
26
+ )
27
+
28
+
29
+ class RapidOcrEngine:
30
+ def __init__(self, device: Device = Device.AUTO) -> None:
31
+ if not is_ocr_available():
32
+ raise OcrUnavailable(
33
+ "ONNX Runtime is missing. Install the cpu or cuda project extra."
34
+ )
35
+ from rapidocr import RapidOCR
36
+
37
+ # RapidOCR 3.9 defaults to PP-OCRv6-small detection and recognition.
38
+ self._engine = RapidOCR(params={"Global.use_cls": False})
39
+
40
+ def recognize(self, rgba: np.ndarray, confidence: float) -> list[OcrLine]:
41
+ rgb = rgba[:, :, :3]
42
+ result = self._engine(rgb)
43
+ boxes = getattr(result, "boxes", None)
44
+ txts = getattr(result, "txts", None)
45
+ scores = getattr(result, "scores", None)
46
+ if boxes is None or txts is None or scores is None:
47
+ return []
48
+
49
+ lines: list[OcrLine] = []
50
+ for box, text, score in zip(boxes, txts, scores, strict=True):
51
+ value = float(score)
52
+ clean = str(text).strip()
53
+ if value < confidence or not clean:
54
+ continue
55
+ quad = [Point(x=max(0, float(x)), y=max(0, float(y))) for x, y in box]
56
+ if len(quad) != 4 or polygon_area(quad) < 4:
57
+ continue
58
+ lines.append(OcrLine(text=clean, confidence=value, quad=quad))
59
+ return sorted(lines, key=lambda line: (quad_box(line.quad)[1], quad_box(line.quad)[0]))
60
+
61
+
62
+ def polygon_area(points: list[Point]) -> float:
63
+ return abs(
64
+ sum(
65
+ point.x * points[(index + 1) % len(points)].y
66
+ - points[(index + 1) % len(points)].x * point.y
67
+ for index, point in enumerate(points)
68
+ )
69
+ / 2
70
+ )
71
+
72
+
73
+ def quad_box(quad: list[Point]) -> tuple[float, float, float, float]:
74
+ xs = [point.x for point in quad]
75
+ ys = [point.y for point in quad]
76
+ return min(xs), min(ys), max(xs) - min(xs), max(ys) - min(ys)
77
+
78
+
79
+ def quad_rotation(quad: list[Point]) -> float:
80
+ first, second = quad[0], quad[1]
81
+ return math.degrees(math.atan2(second.y - first.y, second.x - first.x))
82
+
83
+
84
+ def snap_rotation(angle: float, targets: list[float], tolerance: float) -> float:
85
+ nearest = min(targets, key=lambda target: angular_distance(angle, target))
86
+ if angular_distance(angle, nearest) <= tolerance:
87
+ return float(nearest)
88
+ return angle
89
+
90
+
91
+ def angular_distance(first: float, second: float) -> float:
92
+ return abs((first - second + 180) % 360 - 180)
src/editable_image/pipeline.py ADDED
@@ -0,0 +1,153 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from __future__ import annotations
2
+
3
+ from dataclasses import dataclass, field
4
+ from pathlib import Path
5
+ from uuid import uuid4
6
+
7
+ import numpy as np
8
+
9
+ from .image_io import LoadedImage, load_image_path, save_png
10
+ from .bitmap_extract import extract_bitmap_layers, reconstruct_background
11
+ from .inpaint import inpaint
12
+ from .masking import build_mask
13
+ from .models import Canvas, DocumentSpec, ProcessOptions, TextRegion
14
+ from .ocr import RapidOcrEngine, snap_rotation
15
+ from .shape_recognition import detect_shapes
16
+ from .style import estimate_style, snap_sizes_to_distribution_peaks
17
+ from .svg_export import build_svg
18
+
19
+
20
+ @dataclass(slots=True)
21
+ class ProcessResult:
22
+ image: LoadedImage
23
+ document: DocumentSpec
24
+ mask: np.ndarray
25
+ background: np.ndarray
26
+ bitmap_assets: dict[str, np.ndarray] = field(default_factory=dict)
27
+
28
+
29
+ class EditableImagePipeline:
30
+ def __init__(self, options: ProcessOptions) -> None:
31
+ self.options = options
32
+
33
+ def analyze(self, image: LoadedImage, source_name: str) -> ProcessResult:
34
+ engine = RapidOcrEngine(self.options.device)
35
+ lines = engine.recognize(image.rgba, self.options.confidence)
36
+ regions: list[TextRegion] = []
37
+ for line in lines:
38
+ estimated = estimate_style(image.rgba, line.text, line.quad)
39
+ regions.append(
40
+ TextRegion(
41
+ id=f"text-{uuid4().hex[:12]}",
42
+ text=line.text,
43
+ confidence=line.confidence,
44
+ source_quad=line.quad,
45
+ overlay_box=estimated.box,
46
+ rotation=snap_rotation(
47
+ estimated.rotation,
48
+ self.options.snap_angles,
49
+ self.options.snap_tolerance,
50
+ ),
51
+ style=estimated.style,
52
+ )
53
+ )
54
+ if self.options.snap_font_sizes:
55
+ snapped_sizes = snap_sizes_to_distribution_peaks(
56
+ [region.style.font_size for region in regions]
57
+ )
58
+ for region, font_size in zip(regions, snapped_sizes, strict=True):
59
+ region.style.font_size = font_size
60
+ document = DocumentSpec(
61
+ version=2 if self.options.vectorize_shapes else 1,
62
+ source_name=source_name,
63
+ canvas=Canvas(width=image.width, height=image.height),
64
+ regions=regions,
65
+ )
66
+ if self.options.vectorize_shapes:
67
+ text_mask = build_mask(image.rgba, document)
68
+ document.shapes = detect_shapes(image.rgba, text_mask)
69
+ return self.rebuild(image, document)
70
+
71
+ def rebuild(self, image: LoadedImage, document: DocumentSpec) -> ProcessResult:
72
+ if (document.canvas.width, document.canvas.height) != (image.width, image.height):
73
+ raise ValueError("Document dimensions do not match the source image")
74
+ text_only = document.model_copy(deep=True)
75
+ text_only.shapes = []
76
+ text_only.bitmap_regions = []
77
+ text_mask = build_mask(image.rgba, text_only)
78
+ bitmap_assets: dict[str, np.ndarray] = {}
79
+ if document.shapes:
80
+ bitmap_regions, bitmap_assets = extract_bitmap_layers(
81
+ image.rgba,
82
+ document,
83
+ text_mask,
84
+ self.options.bitmap_extractor,
85
+ self.options.device.value,
86
+ )
87
+ document.bitmap_regions = bitmap_regions
88
+ mask = build_mask(image.rgba, document)
89
+ if document.shapes:
90
+ background = reconstruct_background(
91
+ image.rgba,
92
+ document,
93
+ mask,
94
+ lambda source, remaining: inpaint(
95
+ source,
96
+ remaining,
97
+ self.options.cleanup,
98
+ self.options.device,
99
+ ),
100
+ )
101
+ else:
102
+ background = inpaint(image.rgba, mask, self.options.cleanup, self.options.device)
103
+ return ProcessResult(
104
+ image=image,
105
+ document=document,
106
+ mask=mask,
107
+ background=background,
108
+ bitmap_assets=bitmap_assets,
109
+ )
110
+
111
+
112
+ def export_result(
113
+ result: ProcessResult,
114
+ directory: Path,
115
+ stem: str,
116
+ *,
117
+ embed_fonts: bool = False,
118
+ ) -> tuple[Path, Path, list[Path]]:
119
+ directory.mkdir(parents=True, exist_ok=True)
120
+ background_path = directory / f"{stem}-background.png"
121
+ overlay_path = directory / f"{stem}-overlay.svg"
122
+ asset_directory = directory / f"{stem}-assets"
123
+ save_png(result.background, background_path, result.image.icc_profile)
124
+ asset_paths: list[Path] = []
125
+ bitmap_hrefs: dict[str, tuple[str, str]] = {}
126
+ if result.bitmap_assets:
127
+ asset_directory.mkdir(parents=True, exist_ok=True)
128
+ regions = {region.id: region for region in result.document.bitmap_regions}
129
+ for region_id, pixels in result.bitmap_assets.items():
130
+ region = regions[region_id]
131
+ asset_path = asset_directory / region.asset_name
132
+ save_png(pixels, asset_path)
133
+ asset_paths.append(asset_path)
134
+ bitmap_hrefs[region_id] = (
135
+ f"{asset_directory.name}/{region.asset_name}",
136
+ asset_path.resolve().as_uri(),
137
+ )
138
+ overlay_path.write_text(
139
+ build_svg(
140
+ result.document,
141
+ background_href=background_path.name,
142
+ background_fallback_href=background_path.name,
143
+ embed_fonts=embed_fonts,
144
+ bitmap_hrefs=bitmap_hrefs,
145
+ ),
146
+ encoding="utf-8",
147
+ )
148
+ return background_path, overlay_path, asset_paths
149
+
150
+
151
+ def process_path(path: Path, options: ProcessOptions) -> ProcessResult:
152
+ image = load_image_path(path, max_pixels=options.max_pixels)
153
+ return EditableImagePipeline(options).analyze(image, path.name)
src/editable_image/sam2.py ADDED
@@ -0,0 +1,112 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from __future__ import annotations
2
+
3
+ import importlib.util
4
+ from functools import lru_cache
5
+ from pathlib import Path
6
+
7
+ import cv2
8
+ import numpy as np
9
+ from platformdirs import user_cache_path
10
+
11
+ from .models import Device
12
+
13
+
14
+ SAM2_REPOSITORY = "facebook/sam2.1-hiera-tiny"
15
+
16
+
17
+ class ShapeExtractionUnavailable(RuntimeError):
18
+ pass
19
+
20
+
21
+ def sam2_path() -> Path:
22
+ return user_cache_path("editable-image") / "models" / "sam2.1-hiera-tiny"
23
+
24
+
25
+ def sam2_installed() -> bool:
26
+ target = sam2_path()
27
+ return (target / "config.json").exists() and (target / "model.safetensors").exists()
28
+
29
+
30
+ def sam2_dependencies_installed() -> bool:
31
+ return all(importlib.util.find_spec(name) is not None for name in ("torch", "transformers"))
32
+
33
+
34
+ def install_sam2() -> Path:
35
+ if importlib.util.find_spec("huggingface_hub") is None:
36
+ raise ShapeExtractionUnavailable(
37
+ "Install the sam2 project extra before downloading the SAM 2 model"
38
+ )
39
+ from huggingface_hub import snapshot_download
40
+
41
+ target = sam2_path()
42
+ target.mkdir(parents=True, exist_ok=True)
43
+ snapshot_download(
44
+ repo_id=SAM2_REPOSITORY,
45
+ local_dir=target,
46
+ allow_patterns=["*.json", "*.safetensors", "*.txt"],
47
+ )
48
+ if not sam2_installed():
49
+ raise ShapeExtractionUnavailable("The downloaded SAM 2 model is incomplete")
50
+ return target
51
+
52
+
53
+ def refine_with_sam2(
54
+ rgb: np.ndarray,
55
+ residual_mask: np.ndarray,
56
+ container_mask: np.ndarray,
57
+ device: str,
58
+ ) -> np.ndarray:
59
+ if not sam2_dependencies_installed() or not sam2_installed():
60
+ raise ShapeExtractionUnavailable(
61
+ "SAM 2 is unavailable. Install the sam2 extra and run "
62
+ "'editable-image models install sam2'."
63
+ )
64
+ generator = mask_generator(resolve_device(Device(device)))
65
+ outputs = generator(rgb, points_per_batch=32)
66
+ masks = outputs.get("masks", [])
67
+ selected = np.zeros_like(residual_mask, dtype=np.uint8)
68
+ container_area = max(cv2.countNonZero(container_mask), 1)
69
+ for raw in masks:
70
+ candidate = np.asarray(raw, dtype=np.uint8)
71
+ if candidate.shape != residual_mask.shape:
72
+ candidate = cv2.resize(
73
+ candidate,
74
+ (residual_mask.shape[1], residual_mask.shape[0]),
75
+ interpolation=cv2.INTER_NEAREST,
76
+ )
77
+ candidate = (candidate > 0).astype(np.uint8) * 255
78
+ area = cv2.countNonZero(candidate)
79
+ if area == 0 or area > container_area * 0.8:
80
+ continue
81
+ inside = cv2.countNonZero(cv2.bitwise_and(candidate, container_mask)) / area
82
+ residual_overlap = cv2.countNonZero(cv2.bitwise_and(candidate, residual_mask)) / area
83
+ if inside >= 0.8 and residual_overlap >= 0.12:
84
+ selected = cv2.bitwise_or(selected, candidate)
85
+ selected[container_mask == 0] = 0
86
+ return selected if cv2.countNonZero(selected) else residual_mask
87
+
88
+
89
+ def resolve_device(requested: Device) -> str:
90
+ import torch
91
+
92
+ if requested == Device.CPU:
93
+ return "cpu"
94
+ if requested in {Device.AUTO, Device.CUDA} and torch.cuda.is_available():
95
+ return "cuda"
96
+ if requested in {Device.AUTO, Device.MPS} and torch.backends.mps.is_available():
97
+ return "mps"
98
+ if requested in {Device.CUDA, Device.MPS}:
99
+ raise ShapeExtractionUnavailable(f"The requested {requested.value} device is unavailable")
100
+ return "cpu"
101
+
102
+
103
+ @lru_cache(maxsize=3)
104
+ def mask_generator(device: str):
105
+ from transformers import pipeline
106
+
107
+ pipeline_device: int | str = 0 if device == "cuda" else device
108
+ return pipeline(
109
+ "mask-generation",
110
+ model=str(sam2_path()),
111
+ device=pipeline_device,
112
+ )
src/editable_image/shape_recognition.py ADDED
@@ -0,0 +1,769 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from __future__ import annotations
2
+
3
+ from dataclasses import dataclass
4
+ from math import atan2, degrees
5
+
6
+ import cv2
7
+ import numpy as np
8
+
9
+ from .models import (
10
+ Box,
11
+ EllipseGeometry,
12
+ GradientStop,
13
+ LinearGradientPaint,
14
+ PathGeometry,
15
+ Point,
16
+ PolygonGeometry,
17
+ RectGeometry,
18
+ ShapeRegion,
19
+ ShapeStyle,
20
+ SolidPaint,
21
+ )
22
+
23
+
24
+ @dataclass(slots=True)
25
+ class Candidate:
26
+ shape: ShapeRegion
27
+ mask: np.ndarray
28
+ score: float
29
+
30
+
31
+ def detect_shapes(rgba: np.ndarray, protected_mask: np.ndarray) -> list[ShapeRegion]:
32
+ """Recognize low-complexity diagram primitives without tracing illustrations."""
33
+ rgb = rgba[:, :, :3]
34
+ height, width = rgb.shape[:2]
35
+ clusters = color_clusters(rgb, protected_mask)
36
+ candidates: list[Candidate] = smoothed_container_candidates(rgb, protected_mask)
37
+ candidates.extend(
38
+ arrow_candidates(
39
+ rgb,
40
+ protected_mask,
41
+ [*clusters, *arrow_color_masks(rgb, protected_mask)],
42
+ )
43
+ )
44
+ for cluster_index, cluster_mask in enumerate(clusters):
45
+ cleaned = cv2.morphologyEx(cluster_mask, cv2.MORPH_OPEN, np.ones((2, 2), np.uint8))
46
+ cleaned = cv2.morphologyEx(cleaned, cv2.MORPH_CLOSE, np.ones((5, 5), np.uint8))
47
+ contours, _ = cv2.findContours(cleaned, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
48
+ for contour in contours:
49
+ candidate = fit_cluster_contour(
50
+ rgb,
51
+ cleaned,
52
+ contour,
53
+ protected_mask,
54
+ width,
55
+ height,
56
+ cluster_index,
57
+ )
58
+ if candidate is not None:
59
+ candidates.append(candidate)
60
+ candidates.extend(edge_rectangle_candidates(rgb, protected_mask))
61
+ accepted = deduplicate(candidates)
62
+ containers = [candidate for candidate in accepted if candidate.shape.is_container]
63
+ accepted = [
64
+ candidate
65
+ for candidate in accepted
66
+ if not (
67
+ isinstance(candidate.shape.geometry, EllipseGeometry)
68
+ and any(
69
+ cv2.countNonZero(cv2.bitwise_and(candidate.mask, container.mask))
70
+ / max(cv2.countNonZero(candidate.mask), 1)
71
+ >= 0.9
72
+ for container in containers
73
+ )
74
+ and cv2.countNonZero(cv2.bitwise_and(candidate.mask, protected_mask)) < 3
75
+ )
76
+ ]
77
+ accepted = filter_arrow_context(accepted)
78
+ accepted.sort(
79
+ key=lambda item: (-item.shape.is_container, -int(item.mask.sum()), item.shape.id)
80
+ )
81
+ for index, candidate in enumerate(accepted, 1):
82
+ candidate.shape.id = f"shape-{index:04d}"
83
+ candidate.shape.z_index = 100 if candidate.shape.is_container else 700
84
+ return [candidate.shape for candidate in accepted]
85
+
86
+
87
+ def color_clusters(rgb: np.ndarray, protected_mask: np.ndarray, count: int = 10) -> list[np.ndarray]:
88
+ lab = cv2.cvtColor(rgb, cv2.COLOR_RGB2LAB).astype(np.float32)
89
+ pixels = lab.reshape(-1, 3)
90
+ allowed = protected_mask.reshape(-1) == 0
91
+ indices = np.flatnonzero(allowed)
92
+ if indices.size == 0:
93
+ return []
94
+ stride = max(1, indices.size // 50_000)
95
+ sample = pixels[indices[::stride]].copy()
96
+ count = min(count, max(2, sample.shape[0] // 100))
97
+ cv2.setRNGSeed(1729)
98
+ _, _, centers = cv2.kmeans(
99
+ sample,
100
+ count,
101
+ None,
102
+ (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.4),
103
+ 3,
104
+ cv2.KMEANS_PP_CENTERS,
105
+ )
106
+ distances = np.linalg.norm(pixels[:, None, :] - centers[None, :, :], axis=2)
107
+ labels = np.argmin(distances, axis=1).reshape(lab.shape[:2])
108
+ output: list[np.ndarray] = []
109
+ for index in range(count):
110
+ mask = ((labels == index) & (protected_mask == 0)).astype(np.uint8) * 255
111
+ output.append(mask)
112
+ return output
113
+
114
+
115
+ def arrow_candidates(
116
+ rgb: np.ndarray,
117
+ protected_mask: np.ndarray,
118
+ clusters: list[np.ndarray],
119
+ ) -> list[Candidate]:
120
+ """Find connected, concave silhouettes with a supported arrow tip and tail."""
121
+ height, width = rgb.shape[:2]
122
+ image_area = width * height
123
+ output: list[Candidate] = []
124
+ for cluster_index, cluster_mask in enumerate(clusters):
125
+ cleaned = cv2.morphologyEx(cluster_mask, cv2.MORPH_OPEN, np.ones((2, 2), np.uint8))
126
+ cleaned = cv2.morphologyEx(cleaned, cv2.MORPH_CLOSE, np.ones((3, 3), np.uint8))
127
+ contours, _ = cv2.findContours(cleaned, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
128
+ for contour in contours:
129
+ area = float(cv2.contourArea(contour))
130
+ x, y, w, h = cv2.boundingRect(contour)
131
+ aspect = max(w / max(h, 1), h / max(w, 1))
132
+ if (
133
+ area < max(45, image_area * 0.000025)
134
+ or w * h < 600
135
+ or max(w, h) < 24
136
+ # Compact arrowheads and short curved arrows can be almost square.
137
+ or aspect < 1.03
138
+ ):
139
+ continue
140
+ perimeter = cv2.arcLength(contour, True)
141
+ if perimeter <= 0:
142
+ continue
143
+ approx = cv2.approxPolyDP(contour, max(1.2, perimeter * 0.008), True)
144
+ if not 5 <= len(approx) <= 48:
145
+ continue
146
+ hull = cv2.convexHull(contour)
147
+ hull_area = float(cv2.contourArea(hull))
148
+ solidity = area / max(hull_area, 1)
149
+ fill_ratio = area / max(w * h, 1)
150
+ if (
151
+ solidity >= 0.93
152
+ or fill_ratio > 0.55
153
+ or (len(approx) > 14 and fill_ratio >= 0.28)
154
+ or (area > image_area * 0.003 and aspect < 4)
155
+ ):
156
+ continue
157
+ tip = supported_arrow_tip(approx)
158
+ if tip is None:
159
+ continue
160
+ tip_index, tip_score = tip
161
+ points = [Point(x=float(item[0][0]), y=float(item[0][1])) for item in approx]
162
+ geometry = PathGeometry(points=points, closed=True)
163
+ fitted_mask = np.zeros((height, width), dtype=np.uint8)
164
+ cv2.drawContours(fitted_mask, [contour], -1, 255, -1)
165
+ sample_mask = cv2.bitwise_and(cleaned, fitted_mask)
166
+ sample_mask[protected_mask > 0] = 0
167
+ style = estimate_style(rgb, sample_mask, geometry)
168
+ confidence = min(0.99, 0.72 + tip_score * 0.2 + (1 - solidity) * 0.1)
169
+ shape = ShapeRegion(
170
+ id=f"arrow-{cluster_index}-{tip_index}",
171
+ confidence=confidence,
172
+ geometry=geometry,
173
+ source_contours=[contour_points(contour)],
174
+ style=style,
175
+ semantic_role="arrow",
176
+ )
177
+ output.append(Candidate(shape=shape, mask=fitted_mask, score=confidence))
178
+ return output
179
+
180
+
181
+ def arrow_color_masks(rgb: np.ndarray, protected_mask: np.ndarray) -> list[np.ndarray]:
182
+ hsv = cv2.cvtColor(rgb, cv2.COLOR_RGB2HSV)
183
+ masks: list[np.ndarray] = []
184
+ gray = (
185
+ (hsv[:, :, 1] < 55)
186
+ & (hsv[:, :, 2] >= 45)
187
+ & (hsv[:, :, 2] <= 220)
188
+ & (protected_mask == 0)
189
+ )
190
+ masks.append(gray.astype(np.uint8) * 255)
191
+ # Overlapping bins keep anti-aliased and shaded arrow pixels connected while
192
+ # still separating unrelated diagram colors.
193
+ for start in range(0, 165, 15):
194
+ colored = (
195
+ (hsv[:, :, 0] >= start)
196
+ & (hsv[:, :, 0] < start + 30)
197
+ & (hsv[:, :, 1] >= 55)
198
+ & (hsv[:, :, 2] >= 35)
199
+ & (hsv[:, :, 2] <= 235)
200
+ & (protected_mask == 0)
201
+ )
202
+ if np.count_nonzero(colored) >= 40:
203
+ masks.append(colored.astype(np.uint8) * 255)
204
+ red = (
205
+ ((hsv[:, :, 0] < 15) | (hsv[:, :, 0] >= 165))
206
+ & (hsv[:, :, 1] >= 55)
207
+ & (hsv[:, :, 2] >= 35)
208
+ & (hsv[:, :, 2] <= 235)
209
+ & (protected_mask == 0)
210
+ )
211
+ if np.count_nonzero(red) >= 40:
212
+ masks.append(red.astype(np.uint8) * 255)
213
+ return masks
214
+
215
+
216
+ def filter_arrow_context(candidates: list[Candidate]) -> list[Candidate]:
217
+ containers = [candidate for candidate in candidates if candidate.shape.is_container]
218
+ arrow_groups: dict[int, list[Candidate]] = {index: [] for index in range(len(containers))}
219
+ outside: set[int] = set()
220
+ for candidate in candidates:
221
+ if candidate.shape.semantic_role != "arrow":
222
+ continue
223
+ area = max(cv2.countNonZero(candidate.mask), 1)
224
+ containing = [
225
+ index
226
+ for index, container in enumerate(containers)
227
+ if cv2.countNonZero(cv2.bitwise_and(candidate.mask, container.mask)) / area >= 0.9
228
+ ]
229
+ if not containing:
230
+ outside.add(id(candidate))
231
+ else:
232
+ arrow_groups[containing[0]].append(candidate)
233
+
234
+ accepted_inside: set[int] = set()
235
+ for container_index, group in arrow_groups.items():
236
+ container = containers[container_index]
237
+ eligible = [
238
+ candidate
239
+ for candidate in group
240
+ if color_distance(candidate.shape, container.shape) >= 25
241
+ ]
242
+ peer_counts = {
243
+ id(candidate): sum(
244
+ color_distance(candidate.shape, other.shape) <= 75
245
+ and similar_arrow_scale(candidate.mask, other.mask)
246
+ for other in eligible
247
+ )
248
+ for candidate in eligible
249
+ }
250
+ stable = [candidate for candidate in eligible if peer_counts[id(candidate)] >= 4]
251
+ for candidate in stable:
252
+ stable_peers = sum(
253
+ color_distance(candidate.shape, other.shape) <= 75
254
+ and similar_arrow_scale(candidate.mask, other.mask)
255
+ for other in stable
256
+ )
257
+ # A coherent repeated set is strong evidence inside illustration
258
+ # panels; isolated arrow-like glyphs are usually raster detail.
259
+ if stable_peers >= 4:
260
+ accepted_inside.add(id(candidate))
261
+ return [
262
+ candidate
263
+ for candidate in candidates
264
+ if candidate.shape.semantic_role != "arrow"
265
+ or id(candidate) in outside
266
+ or id(candidate) in accepted_inside
267
+ ]
268
+
269
+
270
+ def color_distance(first: ShapeRegion, second: ShapeRegion) -> float:
271
+ if not isinstance(first.style.fill, SolidPaint) or not isinstance(second.style.fill, SolidPaint):
272
+ return float("inf")
273
+ return float(np.linalg.norm(hex_rgb(first.style.fill.color) - hex_rgb(second.style.fill.color)))
274
+
275
+
276
+ def similar_arrow_scale(first: np.ndarray, second: np.ndarray) -> bool:
277
+ first_area = max(cv2.countNonZero(first), 1)
278
+ second_area = max(cv2.countNonZero(second), 1)
279
+ ratio = max(first_area, second_area) / min(first_area, second_area)
280
+ return ratio <= 3.5
281
+
282
+
283
+ def hex_rgb(value: str) -> np.ndarray:
284
+ value = value.lstrip("#")
285
+ if len(value) == 3:
286
+ value = "".join(character * 2 for character in value)
287
+ return np.array([int(value[index:index + 2], 16) for index in (0, 2, 4)], dtype=np.float32)
288
+
289
+
290
+ def supported_arrow_tip(approx: np.ndarray) -> tuple[int, float] | None:
291
+ points = approx[:, 0, :].astype(np.float32)
292
+ centroid = np.mean(points, axis=0)
293
+ best: tuple[int, float] | None = None
294
+ for index, tip in enumerate(points):
295
+ previous = points[(index - 1) % len(points)]
296
+ following = points[(index + 1) % len(points)]
297
+ side_a = previous - tip
298
+ side_b = following - tip
299
+ norm_a, norm_b = np.linalg.norm(side_a), np.linalg.norm(side_b)
300
+ if min(norm_a, norm_b) < 2:
301
+ continue
302
+ cosine = float(np.clip(np.dot(side_a, side_b) / (norm_a * norm_b), -1, 1))
303
+ angle = float(np.degrees(np.arccos(cosine)))
304
+ if angle > 95:
305
+ continue
306
+ base_midpoint = (previous + following) / 2
307
+ direction = tip - base_midpoint
308
+ direction_norm = np.linalg.norm(direction)
309
+ if direction_norm < 2:
310
+ continue
311
+ direction /= direction_norm
312
+ projections = (points - centroid) @ direction
313
+ tip_projection = float((tip - centroid) @ direction)
314
+ projection_span = float(projections.max() - projections.min())
315
+ if projection_span < 8 or tip_projection < float(projections.max()) - projection_span * 0.08:
316
+ continue
317
+ head_width = float(np.linalg.norm(previous - following))
318
+ if head_width < 4 or projection_span < head_width * 0.65:
319
+ continue
320
+ base_projection = float((base_midpoint - centroid) @ direction)
321
+ tail_points = int(np.count_nonzero(projections < base_projection - max(2, projection_span * 0.08)))
322
+ if tail_points < 2:
323
+ continue
324
+ symmetry = 1 - min(1.0, abs(norm_a - norm_b) / max(norm_a, norm_b))
325
+ sharpness = max(0.0, 1 - angle / 100)
326
+ extremity = min(1.0, (tip_projection - float(np.median(projections))) / max(projection_span * 0.45, 1))
327
+ score = sharpness * 0.4 + symmetry * 0.25 + extremity * 0.35
328
+ if score >= 0.54 and (best is None or score > best[1]):
329
+ best = (index, score)
330
+ return best
331
+
332
+
333
+ def smoothed_container_candidates(
334
+ rgb: np.ndarray,
335
+ protected_mask: np.ndarray,
336
+ ) -> list[Candidate]:
337
+ """Find large low-frequency panels after suppressing their foreground details."""
338
+ height, width = rgb.shape[:2]
339
+ scale = min(1.0, 500 / max(width, height))
340
+ small_width = max(32, int(round(width * scale)))
341
+ small_height = max(32, int(round(height * scale)))
342
+ small = cv2.resize(rgb, (small_width, small_height), interpolation=cv2.INTER_AREA)
343
+ kernel = min(21, min(small_width, small_height) // 2 * 2 - 1)
344
+ kernel = max(5, kernel)
345
+ smooth = cv2.medianBlur(small, kernel)
346
+ lab = cv2.cvtColor(smooth, cv2.COLOR_RGB2LAB).astype(np.float32)
347
+ pixels = lab.reshape(-1, 3)
348
+ cv2.setRNGSeed(2718)
349
+ _, labels, centers = cv2.kmeans(
350
+ pixels,
351
+ 12,
352
+ None,
353
+ (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.35),
354
+ 4,
355
+ cv2.KMEANS_PP_CENTERS,
356
+ )
357
+ labels = labels.reshape(small_height, small_width)
358
+ output: list[Candidate] = []
359
+ for index, center in enumerate(centers):
360
+ mask = (labels == index).astype(np.uint8) * 255
361
+ mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, np.ones((5, 5), np.uint8))
362
+ contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
363
+ for contour in contours:
364
+ x, y, w, h = cv2.boundingRect(contour)
365
+ area = cv2.contourArea(contour)
366
+ if not (
367
+ width * 0.12 <= w / scale <= width * 0.28
368
+ and height * 0.45 <= h / scale <= height * 0.78
369
+ and area / max(w * h, 1) >= 0.55
370
+ ):
371
+ continue
372
+ box = Box(
373
+ x=x / scale,
374
+ y=y / scale,
375
+ width=min(width - x / scale, w / scale),
376
+ height=min(height - y / scale, h / scale),
377
+ )
378
+ box = refine_container_box(rgb, box)
379
+ geometry = RectGeometry(
380
+ box=box,
381
+ rx=min(box.width, box.height) * 0.035,
382
+ ry=min(box.width, box.height) * 0.035,
383
+ )
384
+ fitted = geometry_mask(geometry, width, height)
385
+ full_lab = cv2.cvtColor(rgb, cv2.COLOR_RGB2LAB).astype(np.float32)
386
+ close = np.linalg.norm(full_lab - center[None, None, :], axis=2) < 12
387
+ sample = ((fitted > 0) & close & (protected_mask == 0)).astype(np.uint8) * 255
388
+ style = estimate_style(rgb, sample, geometry)
389
+ corners = [
390
+ Point(x=box.x, y=box.y),
391
+ Point(x=box.x + box.width, y=box.y),
392
+ Point(x=box.x + box.width, y=box.y + box.height),
393
+ Point(x=box.x, y=box.y + box.height),
394
+ ]
395
+ confidence = min(0.98, 0.82 + area / max(w * h, 1) * 0.16)
396
+ shape = ShapeRegion(
397
+ id=f"smooth-{index}",
398
+ confidence=confidence,
399
+ geometry=geometry,
400
+ source_contours=[corners],
401
+ style=style,
402
+ is_container=True,
403
+ )
404
+ output.append(Candidate(shape=shape, mask=fitted, score=confidence))
405
+ return output
406
+
407
+
408
+ def refine_container_box(rgb: np.ndarray, box: Box) -> Box:
409
+ height, width = rgb.shape[:2]
410
+ lab = cv2.cvtColor(rgb, cv2.COLOR_RGB2LAB).astype(np.float32)
411
+ border = np.concatenate(
412
+ [lab[:8].reshape(-1, 3), lab[-8:].reshape(-1, 3), lab[:, :8].reshape(-1, 3), lab[:, -8:].reshape(-1, 3)]
413
+ )
414
+ delta = np.linalg.norm(lab - np.median(border, axis=0), axis=2)
415
+ x0, y0 = max(0, int(box.x)), max(0, int(box.y))
416
+ x1 = min(width, int(np.ceil(box.x + box.width)))
417
+ y1 = min(height, int(np.ceil(box.y + box.height)))
418
+ row_signal = np.median(delta[y0:y1, x0:x1], axis=1) > 2.5
419
+ row_run = longest_run(row_signal)
420
+ if row_run and row_run[1] - row_run[0] >= box.height * 0.5:
421
+ y0, y1 = y0 + row_run[0], y0 + row_run[1]
422
+ column_signal = np.median(delta[y0:y1, x0:x1], axis=0) > 2.5
423
+ column_run = longest_run(column_signal)
424
+ if column_run and column_run[1] - column_run[0] >= box.width * 0.65:
425
+ x0, x1 = x0 + column_run[0], x0 + column_run[1]
426
+ return Box(x=x0, y=y0, width=max(1, x1 - x0), height=max(1, y1 - y0))
427
+
428
+
429
+ def longest_run(values: np.ndarray) -> tuple[int, int] | None:
430
+ mask = (values.astype(np.uint8) * 255)[:, None]
431
+ mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, np.ones((7, 1), np.uint8))[:, 0] > 0
432
+ padded = np.pad(mask.astype(np.int8), (1, 1))
433
+ changes = np.diff(padded)
434
+ starts = np.flatnonzero(changes == 1)
435
+ ends = np.flatnonzero(changes == -1)
436
+ if not len(starts):
437
+ return None
438
+ index = int(np.argmax(ends - starts))
439
+ return int(starts[index]), int(ends[index])
440
+
441
+
442
+ def fit_cluster_contour(
443
+ rgb: np.ndarray,
444
+ cluster_mask: np.ndarray,
445
+ contour: np.ndarray,
446
+ protected_mask: np.ndarray,
447
+ width: int,
448
+ height: int,
449
+ cluster_index: int,
450
+ ) -> Candidate | None:
451
+ area = float(cv2.contourArea(contour))
452
+ image_area = width * height
453
+ if area < max(45, image_area * 0.00003):
454
+ return None
455
+ perimeter = cv2.arcLength(contour, True)
456
+ if perimeter <= 0:
457
+ return None
458
+ approx = cv2.approxPolyDP(contour, max(1.5, perimeter * 0.012), True)
459
+ x, y, w, h = cv2.boundingRect(contour)
460
+ if w < 5 or h < 5:
461
+ return None
462
+ contour_mask = np.zeros((height, width), dtype=np.uint8)
463
+ cv2.drawContours(contour_mask, [contour], -1, 255, -1)
464
+ circularity = 4 * np.pi * area / max(perimeter * perimeter, 1)
465
+ geometry: RectGeometry | EllipseGeometry | PolygonGeometry | PathGeometry
466
+ fitted_mask = np.zeros_like(contour_mask)
467
+ rectangularity = area / max(w * h, 1)
468
+ if len(approx) == 4 and rectangularity >= 0.58:
469
+ rotated = cv2.minAreaRect(contour)
470
+ (cx, cy), (rw, rh), angle = rotated
471
+ if rw < rh:
472
+ rw, rh = rh, rw
473
+ angle += 90
474
+ angle = normalize_angle(angle)
475
+ if cx - rw / 2 < 0 or cy - rh / 2 < 0:
476
+ return None
477
+ radius = min(rw, rh) * (0.045 if rectangularity < 0.97 else 0)
478
+ geometry = RectGeometry(
479
+ box=Box(x=cx - rw / 2, y=cy - rh / 2, width=rw, height=rh),
480
+ rx=radius,
481
+ ry=radius,
482
+ rotation=angle,
483
+ )
484
+ fitted_mask = geometry_mask(geometry, width, height)
485
+ elif circularity >= 0.72 and len(contour) >= 5:
486
+ (cx, cy), (ew, eh), angle = cv2.fitEllipse(contour)
487
+ if ew < 5 or eh < 5:
488
+ return None
489
+ geometry = EllipseGeometry(cx=cx, cy=cy, rx=ew / 2, ry=eh / 2, rotation=angle)
490
+ fitted_mask = geometry_mask(geometry, width, height)
491
+ elif 3 <= len(approx) <= 10:
492
+ points = [Point(x=float(item[0][0]), y=float(item[0][1])) for item in approx]
493
+ geometry = PolygonGeometry(points=points)
494
+ fitted_mask = geometry_mask(geometry, width, height)
495
+ elif len(approx) <= 20 and (w / h >= 3.5 or h / w >= 3.5):
496
+ points = [Point(x=float(item[0][0]), y=float(item[0][1])) for item in approx]
497
+ geometry = PathGeometry(points=points, closed=True)
498
+ fitted_mask = geometry_mask(geometry, width, height)
499
+ else:
500
+ return None
501
+
502
+ large_container = area >= image_area * 0.012 and isinstance(geometry, RectGeometry)
503
+ if not large_container:
504
+ if isinstance(geometry, RectGeometry) and area < image_area * 0.00045:
505
+ return None
506
+ if isinstance(geometry, EllipseGeometry) and min(geometry.rx, geometry.ry) < 12:
507
+ return None
508
+ if isinstance(geometry, PolygonGeometry):
509
+ aspect = max(w / h, h / w)
510
+ if area < image_area * 0.00018 or not (
511
+ len(geometry.points) == 3 or aspect >= 1.8 or area >= image_area * 0.0015
512
+ ):
513
+ return None
514
+
515
+ fitted_area = cv2.countNonZero(fitted_mask)
516
+ if fitted_area == 0:
517
+ return None
518
+ support = cv2.countNonZero(cv2.bitwise_and(cluster_mask, fitted_mask)) / fitted_area
519
+ overlap = cv2.countNonZero(cv2.bitwise_and(contour_mask, fitted_mask))
520
+ union = cv2.countNonZero(cv2.bitwise_or(contour_mask, fitted_mask))
521
+ iou = overlap / max(union, 1)
522
+ minimum_support = 0.38 if large_container else 0.64
523
+ if support < minimum_support or iou < (0.58 if large_container else 0.68):
524
+ return None
525
+
526
+ sample_mask = cv2.bitwise_and(cluster_mask, fitted_mask)
527
+ sample_mask[protected_mask > 0] = 0
528
+ style = estimate_style(rgb, sample_mask, geometry)
529
+ source = contour_points(contour)
530
+ confidence = min(0.99, 0.45 * support + 0.45 * iou + 0.1)
531
+ shape = ShapeRegion(
532
+ id=f"cluster-{cluster_index}",
533
+ confidence=confidence,
534
+ geometry=geometry,
535
+ source_contours=[source],
536
+ style=style,
537
+ is_container=large_container,
538
+ )
539
+ return Candidate(shape=shape, mask=fitted_mask, score=confidence)
540
+
541
+
542
+ def edge_rectangle_candidates(rgb: np.ndarray, protected_mask: np.ndarray) -> list[Candidate]:
543
+ height, width = rgb.shape[:2]
544
+ gray = cv2.cvtColor(rgb, cv2.COLOR_RGB2GRAY)
545
+ edges = cv2.Canny(gray, 40, 120)
546
+ edges[cv2.dilate(protected_mask, np.ones((3, 3), np.uint8)) > 0] = 0
547
+ edges = cv2.morphologyEx(edges, cv2.MORPH_CLOSE, np.ones((5, 5), np.uint8))
548
+ contours, _ = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
549
+ output: list[Candidate] = []
550
+ image_area = width * height
551
+ for contour in contours:
552
+ area = float(cv2.contourArea(contour))
553
+ if area < image_area * 0.001:
554
+ continue
555
+ perimeter = cv2.arcLength(contour, True)
556
+ approx = cv2.approxPolyDP(contour, max(2, perimeter * 0.02), True)
557
+ x, y, w, h = cv2.boundingRect(contour)
558
+ if not 4 <= len(approx) <= 8 or w < 20 or h < 20:
559
+ continue
560
+ rectangularity = area / max(w * h, 1)
561
+ circularity = 4 * np.pi * area / max(perimeter * perimeter, 1)
562
+ if 0.8 <= w / h <= 1.2 and circularity >= 0.68:
563
+ continue
564
+ if rectangularity < 0.72:
565
+ continue
566
+ geometry = RectGeometry(
567
+ box=Box(x=x, y=y, width=w, height=h),
568
+ rx=min(w, h) * (0.045 if rectangularity < 0.97 else 0),
569
+ ry=min(w, h) * (0.045 if rectangularity < 0.97 else 0),
570
+ )
571
+ mask = geometry_mask(geometry, width, height)
572
+ interior = cv2.erode(mask, np.ones((5, 5), np.uint8))
573
+ interior[protected_mask > 0] = 0
574
+ style = estimate_style(rgb, interior, geometry)
575
+ ring = cv2.subtract(mask, cv2.erode(mask, np.ones((5, 5), np.uint8)))
576
+ ring_pixels = rgb[ring > 0]
577
+ if ring_pixels.size:
578
+ style.stroke = rgb_hex(np.median(ring_pixels, axis=0))
579
+ style.stroke_width = 2
580
+ confidence = min(0.94, 0.55 + rectangularity * 0.4)
581
+ shape = ShapeRegion(
582
+ id="edge-rect",
583
+ confidence=confidence,
584
+ geometry=geometry,
585
+ source_contours=[contour_points(contour)],
586
+ style=style,
587
+ is_container=area >= image_area * 0.012,
588
+ )
589
+ output.append(Candidate(shape=shape, mask=mask, score=confidence))
590
+ return output
591
+
592
+
593
+ def estimate_style(
594
+ rgb: np.ndarray,
595
+ sample_mask: np.ndarray,
596
+ geometry: RectGeometry | EllipseGeometry | PolygonGeometry | PathGeometry,
597
+ ) -> ShapeStyle:
598
+ ys, xs = np.nonzero(sample_mask)
599
+ if xs.size == 0:
600
+ return ShapeStyle(fill=SolidPaint(color="#000000"))
601
+ values = rgb[ys, xs].astype(np.float32)
602
+ median = np.median(values, axis=0)
603
+ solid_error = float(np.mean((values - median) ** 2))
604
+ if xs.size >= 100 and solid_error > 5:
605
+ x_span = max(float(xs.max() - xs.min()), 1)
606
+ y_span = max(float(ys.max() - ys.min()), 1)
607
+ tx = (xs - xs.min()) / x_span
608
+ ty = (ys - ys.min()) / y_span
609
+ fit_x, error_x = linear_color_fit(values, tx)
610
+ fit_y, error_y = linear_color_fit(values, ty)
611
+ fit, error, angle = (fit_x, error_x, 0.0) if error_x <= error_y else (fit_y, error_y, 90.0)
612
+ if error < solid_error * 0.65 and np.linalg.norm(fit[1] - fit[0]) >= 5:
613
+ return ShapeStyle(
614
+ fill=LinearGradientPaint(
615
+ angle=angle,
616
+ stops=[
617
+ GradientStop(offset=0, color=rgb_hex(fit[0])),
618
+ GradientStop(offset=1, color=rgb_hex(fit[1])),
619
+ ],
620
+ )
621
+ )
622
+ return ShapeStyle(fill=SolidPaint(color=rgb_hex(median)))
623
+
624
+
625
+ def linear_color_fit(values: np.ndarray, t: np.ndarray) -> tuple[np.ndarray, float]:
626
+ design = np.column_stack([np.ones_like(t), t])
627
+ coefficients, _, _, _ = np.linalg.lstsq(design, values, rcond=None)
628
+ predicted = design @ coefficients
629
+ endpoints = np.clip(np.vstack([coefficients[0], coefficients[0] + coefficients[1]]), 0, 255)
630
+ return endpoints, float(np.mean((values - predicted) ** 2))
631
+
632
+
633
+ def deduplicate(candidates: list[Candidate]) -> list[Candidate]:
634
+ accepted: list[Candidate] = []
635
+ # Preserve semantic interpretations when their silhouette overlaps a more
636
+ # generic polygon/path fit. Containers still establish the layout first.
637
+ for candidate in sorted(
638
+ candidates,
639
+ key=lambda item: (
640
+ item.shape.is_container,
641
+ item.shape.semantic_role == "arrow",
642
+ item.score,
643
+ ),
644
+ reverse=True,
645
+ ):
646
+ if (
647
+ isinstance(candidate.shape.geometry, RectGeometry)
648
+ and not candidate.shape.is_container
649
+ and any(
650
+ existing.shape.is_container
651
+ and cv2.countNonZero(cv2.bitwise_and(candidate.mask, existing.mask))
652
+ / max(cv2.countNonZero(candidate.mask), 1)
653
+ >= 0.9
654
+ for existing in accepted
655
+ )
656
+ ):
657
+ continue
658
+ duplicate = False
659
+ for existing in accepted:
660
+ intersection = cv2.countNonZero(cv2.bitwise_and(candidate.mask, existing.mask))
661
+ union = cv2.countNonZero(cv2.bitwise_or(candidate.mask, existing.mask))
662
+ smaller = min(cv2.countNonZero(candidate.mask), cv2.countNonZero(existing.mask))
663
+ container_containment = (
664
+ candidate.shape.is_container
665
+ and existing.shape.is_container
666
+ and smaller > 0
667
+ and intersection / smaller >= 0.78
668
+ )
669
+ arrow_containment = (
670
+ candidate.shape.semantic_role == "arrow"
671
+ and existing.shape.semantic_role == "arrow"
672
+ and smaller > 0
673
+ and intersection / smaller >= 0.7
674
+ )
675
+ if union and (
676
+ intersection / union >= 0.78
677
+ or container_containment
678
+ or arrow_containment
679
+ ):
680
+ duplicate = True
681
+ break
682
+ if not duplicate:
683
+ accepted.append(candidate)
684
+ return accepted
685
+
686
+
687
+ def geometry_mask(
688
+ geometry: RectGeometry | EllipseGeometry | PolygonGeometry | PathGeometry,
689
+ width: int,
690
+ height: int,
691
+ ) -> np.ndarray:
692
+ mask = np.zeros((height, width), dtype=np.uint8)
693
+ if isinstance(geometry, RectGeometry):
694
+ box = geometry.box
695
+ center = (box.x + box.width / 2, box.y + box.height / 2)
696
+ if abs(geometry.rotation) > 0.01:
697
+ points = cv2.boxPoints((center, (box.width, box.height), geometry.rotation)).astype(np.int32)
698
+ cv2.fillPoly(mask, [points], 255)
699
+ else:
700
+ p1 = (int(round(box.x)), int(round(box.y)))
701
+ p2 = (int(round(box.x + box.width)), int(round(box.y + box.height)))
702
+ radius = int(round(max(geometry.rx, geometry.ry)))
703
+ if radius > 0:
704
+ cv2.rectangle(mask, (p1[0] + radius, p1[1]), (p2[0] - radius, p2[1]), 255, -1)
705
+ cv2.rectangle(mask, (p1[0], p1[1] + radius), (p2[0], p2[1] - radius), 255, -1)
706
+ for cx, cy in (
707
+ (p1[0] + radius, p1[1] + radius),
708
+ (p2[0] - radius, p1[1] + radius),
709
+ (p1[0] + radius, p2[1] - radius),
710
+ (p2[0] - radius, p2[1] - radius),
711
+ ):
712
+ cv2.circle(mask, (cx, cy), radius, 255, -1)
713
+ else:
714
+ cv2.rectangle(mask, p1, p2, 255, -1)
715
+ elif isinstance(geometry, EllipseGeometry):
716
+ cv2.ellipse(
717
+ mask,
718
+ (int(round(geometry.cx)), int(round(geometry.cy))),
719
+ (int(round(geometry.rx)), int(round(geometry.ry))),
720
+ geometry.rotation,
721
+ 0,
722
+ 360,
723
+ 255,
724
+ -1,
725
+ )
726
+ else:
727
+ points = np.array([(point.x, point.y) for point in geometry.points], dtype=np.int32)
728
+ if isinstance(geometry, PolygonGeometry) or geometry.closed:
729
+ cv2.fillPoly(mask, [points], 255)
730
+ else:
731
+ cv2.polylines(mask, [points], False, 255, 2, cv2.LINE_AA)
732
+ return mask
733
+
734
+
735
+ def shape_source_mask(shape: ShapeRegion, width: int, height: int) -> np.ndarray:
736
+ if not shape.source_contours:
737
+ return geometry_mask(shape.geometry, width, height)
738
+ mask = np.zeros((height, width), dtype=np.uint8)
739
+ contours = [
740
+ np.array([(point.x, point.y) for point in contour], dtype=np.int32)
741
+ for contour in shape.source_contours
742
+ if len(contour) >= 3
743
+ ]
744
+ if contours:
745
+ cv2.fillPoly(mask, contours, 255)
746
+ return mask
747
+
748
+
749
+ def contour_points(contour: np.ndarray) -> list[Point]:
750
+ simplified = cv2.approxPolyDP(contour, max(1, cv2.arcLength(contour, True) * 0.004), True)
751
+ return [Point(x=float(item[0][0]), y=float(item[0][1])) for item in simplified]
752
+
753
+
754
+ def rgb_hex(value: np.ndarray) -> str:
755
+ red, green, blue = np.clip(np.rint(value), 0, 255).astype(np.uint8)
756
+ return f"#{red:02x}{green:02x}{blue:02x}"
757
+
758
+
759
+ def normalize_angle(value: float) -> float:
760
+ while value > 90:
761
+ value -= 180
762
+ while value < -90:
763
+ value += 180
764
+ return float(value)
765
+
766
+
767
+ def path_angle(points: list[Point]) -> float:
768
+ start, end = points[0], points[-1]
769
+ return degrees(atan2(end.y - start.y, end.x - start.x))
src/editable_image/style.py ADDED
@@ -0,0 +1,201 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from __future__ import annotations
2
+
3
+ from dataclasses import dataclass
4
+
5
+ import cv2
6
+ import numpy as np
7
+ from PIL import Image, ImageDraw, ImageFont
8
+
9
+ from .fonts import FACES
10
+ from .models import Box, Point, TextStyle
11
+ from .ocr import quad_box, quad_rotation
12
+
13
+
14
+ @dataclass(slots=True)
15
+ class StyleResult:
16
+ box: Box
17
+ rotation: float
18
+ style: TextStyle
19
+
20
+
21
+ def estimate_style(rgba: np.ndarray, text: str, quad: list[Point]) -> StyleResult:
22
+ x, y, width, height = quad_box(quad)
23
+ max_height, max_width = rgba.shape[:2]
24
+ left = max(0, int(x))
25
+ top = max(0, int(y))
26
+ right = min(max_width, int(np.ceil(x + width)))
27
+ bottom = min(max_height, int(np.ceil(y + height)))
28
+ crop = rgba[top:bottom, left:right, :3]
29
+ font_size = max(6.0, height * 0.92)
30
+ foreground = estimate_foreground_color(crop)
31
+ face, adjusted_size = match_font(crop, text, font_size, width)
32
+ style = TextStyle(
33
+ font_family=face.family,
34
+ font_size=adjusted_size,
35
+ font_weight=700 if face.weight >= 700 else 400,
36
+ italic=face.italic,
37
+ fill=foreground,
38
+ )
39
+ return StyleResult(
40
+ box=Box(x=x, y=y, width=max(1, width), height=max(1, height)),
41
+ rotation=max(-180, min(180, quad_rotation(quad))),
42
+ style=style,
43
+ )
44
+
45
+
46
+ def estimate_foreground_color(crop: np.ndarray) -> str:
47
+ if crop.size == 0:
48
+ return "#111111"
49
+ pixels = crop.reshape(-1, 3).astype(np.float32)
50
+ border = np.concatenate(
51
+ [crop[0, :, :3], crop[-1, :, :3], crop[:, 0, :3], crop[:, -1, :3]], axis=0
52
+ ).astype(np.float32)
53
+ background = np.median(border, axis=0)
54
+ distances = np.linalg.norm(pixels - background, axis=1)
55
+ threshold = max(24.0, float(np.percentile(distances, 75)))
56
+ candidates = pixels[distances >= threshold]
57
+ if len(candidates) < 3:
58
+ candidates = pixels[np.argsort(distances)[-max(3, len(pixels) // 10) :]]
59
+ color = np.median(candidates, axis=0).clip(0, 255).astype(int)
60
+ return "#{:02x}{:02x}{:02x}".format(*color)
61
+
62
+
63
+ def match_font(crop: np.ndarray, text: str, size: float, target_width: float):
64
+ available = [face for face in FACES if face.family == "Lato" and face.path.exists()]
65
+ if not available:
66
+ raise FileNotFoundError("Bundled fonts are missing")
67
+ best_face = available[0]
68
+ best_size = max(6, int(round(size)))
69
+ best_score = float("inf")
70
+ observed = foreground_binary(crop)
71
+ likely_italic = estimate_italic(observed)
72
+ available = [face for face in available if face.italic == likely_italic]
73
+ for face in available:
74
+ candidate_size = max(6, int(round(size)))
75
+ font = ImageFont.truetype(str(face.path), candidate_size)
76
+ bounds = font.getbbox(text or " ")
77
+ rendered_width = max(1, bounds[2] - bounds[0])
78
+ scaled_size = int(np.clip(candidate_size * target_width / rendered_width, 6, 1000))
79
+ mask = render_text_mask(text, face.path, scaled_size, crop.shape[:2])
80
+ score = mask_distance(observed, mask)
81
+ score += 0.002 if face.weight >= 700 else 0
82
+ if score < best_score:
83
+ best_score = score
84
+ best_face = face
85
+ best_size = scaled_size
86
+ return best_face, float(best_size)
87
+
88
+
89
+ def foreground_binary(crop: np.ndarray) -> np.ndarray:
90
+ if crop.size == 0:
91
+ return np.zeros((1, 1), dtype=np.uint8)
92
+ gray = cv2.cvtColor(crop, cv2.COLOR_RGB2GRAY)
93
+ border = np.concatenate([gray[0], gray[-1], gray[:, 0], gray[:, -1]])
94
+ background = float(np.median(border))
95
+ delta = cv2.absdiff(gray, np.full_like(gray, int(background)))
96
+ _, mask = cv2.threshold(delta, 18, 255, cv2.THRESH_BINARY)
97
+ return mask
98
+
99
+
100
+ def render_text_mask(text: str, path, size: int, shape: tuple[int, int]) -> np.ndarray:
101
+ height, width = max(1, shape[0]), max(1, shape[1])
102
+ image = Image.new("L", (width, height), 0)
103
+ draw = ImageDraw.Draw(image)
104
+ font = ImageFont.truetype(str(path), size)
105
+ bounds = draw.textbbox((0, 0), text, font=font)
106
+ y = max(-bounds[1], (height - (bounds[3] - bounds[1])) // 2 - bounds[1])
107
+ draw.text((0, y), text, font=font, fill=255)
108
+ return np.asarray(image)
109
+
110
+
111
+ def mask_distance(observed: np.ndarray, candidate: np.ndarray) -> float:
112
+ if observed.shape != candidate.shape:
113
+ candidate = cv2.resize(candidate, (observed.shape[1], observed.shape[0]))
114
+ observed_pixels = observed > 0
115
+ candidate_pixels = candidate > 0
116
+ if not np.any(observed_pixels) or not np.any(candidate_pixels):
117
+ return 1.0
118
+ distance_to_observed = cv2.distanceTransform(
119
+ (~observed_pixels).astype(np.uint8), cv2.DIST_L2, 3
120
+ )
121
+ distance_to_candidate = cv2.distanceTransform(
122
+ (~candidate_pixels).astype(np.uint8), cv2.DIST_L2, 3
123
+ )
124
+ chamfer = (
125
+ float(distance_to_observed[candidate_pixels].mean())
126
+ + float(distance_to_candidate[observed_pixels].mean())
127
+ ) / (2 * max(1, observed.shape[0]))
128
+ density_delta = abs(
129
+ float(candidate_pixels.mean()) - float(observed_pixels.mean())
130
+ )
131
+ return chamfer + density_delta * 0.45
132
+
133
+
134
+ def estimate_italic(mask: np.ndarray) -> bool:
135
+ if mask.size == 0 or np.count_nonzero(mask) < 12:
136
+ return False
137
+ height, width = mask.shape
138
+ shears = np.linspace(-0.3, 0.3, 13)
139
+ scores: list[float] = []
140
+ for shear in shears:
141
+ transform = np.float32([[1, shear, -shear * height / 2], [0, 1, 0]])
142
+ deskewed = cv2.warpAffine(mask, transform, (width, height)) > 0
143
+ columns = deskewed.sum(axis=0).astype(np.float32)
144
+ scores.append(float(np.square(columns).sum() / max(1, columns.sum())))
145
+ return float(shears[int(np.argmax(scores))]) >= 0.12
146
+
147
+
148
+ def snap_sizes_to_distribution_peaks(sizes: list[float]) -> list[float]:
149
+ if len(sizes) < 2:
150
+ return list(sizes)
151
+ values = np.asarray(sizes, dtype=np.float64)
152
+ bandwidth = max(0.75, float(np.median(values)) * 0.06)
153
+ peaks = font_size_distribution_peaks(values, bandwidth)
154
+ if not peaks:
155
+ return list(sizes)
156
+
157
+ snap_distance = bandwidth * 2
158
+ snapped: list[float] = []
159
+ for value in values:
160
+ nearest = min(peaks, key=lambda peak: abs(peak - value))
161
+ snapped.append(nearest if abs(nearest - value) <= snap_distance else float(value))
162
+ return snapped
163
+
164
+
165
+ def font_size_distribution_peaks(values: np.ndarray, bandwidth: float) -> list[float]:
166
+ lower = float(values.min())
167
+ upper = float(values.max())
168
+ if upper - lower < 1e-6:
169
+ return [lower]
170
+
171
+ sample_count = min(
172
+ 4096,
173
+ max(64, int(np.ceil((upper - lower) / (bandwidth / 12))) + 1),
174
+ )
175
+ grid = np.linspace(lower, upper, sample_count)
176
+ offsets = (grid[:, None] - values[None, :]) / bandwidth
177
+ density = np.exp(-0.5 * np.square(offsets)).sum(axis=1)
178
+
179
+ candidates: list[int] = []
180
+ if density[0] >= density[1]:
181
+ candidates.append(0)
182
+ candidates.extend(
183
+ index
184
+ for index in range(1, sample_count - 1)
185
+ if density[index] >= density[index - 1]
186
+ and density[index] >= density[index + 1]
187
+ )
188
+ if density[-1] >= density[-2]:
189
+ candidates.append(sample_count - 1)
190
+
191
+ peaks: list[float] = []
192
+ for index in sorted(candidates, key=lambda item: density[item], reverse=True):
193
+ center = grid[index]
194
+ distances = np.abs(values - center)
195
+ supported = values[distances <= bandwidth * 2]
196
+ if len(supported) < 2:
197
+ continue
198
+ peak = round(float(np.mean(supported)) * 2) / 2
199
+ if all(abs(peak - existing) > bandwidth * 2 for existing in peaks):
200
+ peaks.append(peak)
201
+ return sorted(peaks)