From e1a6ea9e6760ea8f3f2cb47e8c7b84ecf8c68181 Mon Sep 17 00:00:00 2001 From: "Github Action (authored by pmalacho-mit)" Date: Fri, 15 May 2026 05:04:12 +0000 Subject: [PATCH 1/6] upgrading to python 3.12. removing playsound dependency (didn't seem necessary?). adding in devcontainer config --- .devcontainer/devcontainer.json | 7 ++++--- docker/backend.Dockerfile | 2 +- docker/backend.Dockerfile.dockerignore | 1 + requirements.txt | 8 +++++--- 4 files changed, 11 insertions(+), 7 deletions(-) diff --git a/.devcontainer/devcontainer.json b/.devcontainer/devcontainer.json index 3b2dbbb..cb41494 100644 --- a/.devcontainer/devcontainer.json +++ b/.devcontainer/devcontainer.json @@ -41,7 +41,8 @@ "mounts": [ "source=${localEnv:HOME}${localEnv:USERPROFILE}/.ssh,target=/home/vscode/.ssh,type=bind,consistency=cached" ], - "postCreateCommand": "sudo apt-get install -y inotify-tools && (grep -q 'IgnoreUnknown UseKeychain' ~/.ssh/config 2>/dev/null || sed -i '/UseKeychain/i IgnoreUnknown UseKeychain' ~/.ssh/config 2>/dev/null) && git config core.editor '/usr/bin/vim' && git config pull.rebase false && sudo curl -fsSL https://gist.githubusercontent.com/pmalacho-mit/d64caa8e16b0b0fdd5e58cc37a0ce242/raw/8f0f620d4a65f6f3d930ea14e7c2fdf382ef1bf1/codeblockify.sh -o /usr/local/bin/codeblockify && sudo chmod +x /usr/local/bin/codeblockify && curl -fsSL https://raw.githubusercontent.com/pmalacho-mit/devcontainers/refs/heads/main/src/difftastic/install.sh | sudo install -m 0755 /dev/stdin /usr/local/bin/difftastic && git config diff.external difftastic", + "postCreateCommand": "sudo apt-get install -y inotify-tools && sudo apt-get update && sudo apt-get install -y libasound-dev libportaudio2 libportaudiocpp0 portaudio19-dev && sudo apt-get install -y gcc && (grep -q 'IgnoreUnknown UseKeychain' ~/.ssh/config 2>/dev/null || sed -i '/UseKeychain/i IgnoreUnknown UseKeychain' ~/.ssh/config 2>/dev/null) && git config core.editor '/usr/bin/vim' && git config pull.rebase false && sudo curl -fsSL https://gist.githubusercontent.com/pmalacho-mit/d64caa8e16b0b0fdd5e58cc37a0ce242/raw/8f0f620d4a65f6f3d930ea14e7c2fdf382ef1bf1/codeblockify.sh -o /usr/local/bin/codeblockify && sudo chmod +x /usr/local/bin/codeblockify && curl -fsSL https://raw.githubusercontent.com/pmalacho-mit/devcontainers/refs/heads/main/src/difftastic/install.sh | sudo install -m 0755 /dev/stdin /usr/local/bin/difftastic && git config diff.external difftastic", "workspaceFolder": "${localWorkspaceFolder}", - "workspaceMount": "source=${localWorkspaceFolder},target=${localWorkspaceFolder},type=bind" -} + "workspaceMount": "source=${localWorkspaceFolder},target=${localWorkspaceFolder},type=bind", + "postStartCommand": "pip install --upgrade -r requirements.txt" +} \ No newline at end of file diff --git a/docker/backend.Dockerfile b/docker/backend.Dockerfile index 5f8c4c7..f3ce42b 100644 --- a/docker/backend.Dockerfile +++ b/docker/backend.Dockerfile @@ -1,4 +1,4 @@ -FROM python:3.9 +FROM python:3.12 # Install dependencies for pyaudio RUN apt-get update diff --git a/docker/backend.Dockerfile.dockerignore b/docker/backend.Dockerfile.dockerignore index a283ba9..9384d00 100644 --- a/docker/backend.Dockerfile.dockerignore +++ b/docker/backend.Dockerfile.dockerignore @@ -1,6 +1,7 @@ cli docker frontend +.devcontainer # Tests **/__tests__/ diff --git a/requirements.txt b/requirements.txt index 801622e..28edfca 100644 --- a/requirements.txt +++ b/requirements.txt @@ -3,11 +3,13 @@ uvicorn>=0.15.0 openai>=1.0.0 azure-cognitiveservices-speech>=1.31.0 pyaudio>=0.2.11 -playsound>=1.3.0 python-dotenv>=0.19.0 sounddevice>=0.4.6 -scipy>=1.11.0 -numpy>=1.24.0 +scikit-image==0.26.0 +scipy==1.17.1 +networkx==3.6.1 +numpy==2.4.4 +pillow==12.2.0 python-multipart>=0.0.6 pydantic>=2.0.0 aiofiles>=23.0.0 From 0e60c21f03fd5cede67a4b34b45ad178f42f1a68 Mon Sep 17 00:00:00 2001 From: "Github Action (authored by pmalacho-mit)" Date: Fri, 15 May 2026 05:55:43 +0000 Subject: [PATCH 2/6] adding in arc library and new endpoint --- .devcontainer/devcontainer.json | 2 +- .../.dependencies/requirements.txt | 5 + .../workflows/subrepo-pull-into-main.yml | 184 +++ arc_line_vectorization_suede/.gitignore | 4 + arc_line_vectorization_suede/.gitrepo | 12 + arc_line_vectorization_suede/__init__.py | 115 ++ arc_line_vectorization_suede/commands.py | 66 + .../graph/__init__.py | 621 +++++++++ .../segment/__init__.py | 135 ++ .../segment/_helpers.py | 239 ++++ .../segment/fusion.py | 622 +++++++++ .../segment/repair.py | 558 ++++++++ arc_line_vectorization_suede/segment/trace.py | 156 +++ .../skeletonize/__init__.py | 137 ++ .../skeletonize/cleanup.py | 162 +++ .../skeletonize/crossings.py | 505 +++++++ .../vectorize/__init__.py | 0 .../vectorize/high_geometry/__init__.py | 604 ++++++++ .../vectorize/high_geometry/consolidate.py | 1223 +++++++++++++++++ .../vectorize/high_geometry/consolidate_v1.py | 783 +++++++++++ .../vectorize/high_geometry/junction_graph.py | 162 +++ .../vectorize/low_geometry/__init__.py | 427 ++++++ .../vectorize/low_geometry/beautify.py | 260 ++++ .../vectorize/low_geometry/fitting.py | 846 ++++++++++++ .../vectorize/low_geometry/manifest.py | 240 ++++ .../vectorize/low_geometry/primitives.py | 212 +++ .../vectorize/low_geometry/residuals.py | 311 +++++ .../vectorize/low_geometry/routing.py | 427 ++++++ .../vectorize/low_geometry/solve.py | 641 +++++++++ arc_line_vectorization_suede/visualize.py | 505 +++++++ main.py | 219 ++- 31 files changed, 10329 insertions(+), 54 deletions(-) create mode 100644 arc_line_vectorization_suede/.dependencies/requirements.txt create mode 100644 arc_line_vectorization_suede/.github/workflows/subrepo-pull-into-main.yml create mode 100644 arc_line_vectorization_suede/.gitignore create mode 100644 arc_line_vectorization_suede/.gitrepo create mode 100644 arc_line_vectorization_suede/__init__.py create mode 100644 arc_line_vectorization_suede/commands.py create mode 100644 arc_line_vectorization_suede/graph/__init__.py create mode 100644 arc_line_vectorization_suede/segment/__init__.py create mode 100644 arc_line_vectorization_suede/segment/_helpers.py create mode 100644 arc_line_vectorization_suede/segment/fusion.py create mode 100644 arc_line_vectorization_suede/segment/repair.py create mode 100644 arc_line_vectorization_suede/segment/trace.py create mode 100644 arc_line_vectorization_suede/skeletonize/__init__.py create mode 100644 arc_line_vectorization_suede/skeletonize/cleanup.py create mode 100644 arc_line_vectorization_suede/skeletonize/crossings.py create mode 100644 arc_line_vectorization_suede/vectorize/__init__.py create mode 100644 arc_line_vectorization_suede/vectorize/high_geometry/__init__.py create mode 100644 arc_line_vectorization_suede/vectorize/high_geometry/consolidate.py create mode 100644 arc_line_vectorization_suede/vectorize/high_geometry/consolidate_v1.py create mode 100644 arc_line_vectorization_suede/vectorize/high_geometry/junction_graph.py create mode 100644 arc_line_vectorization_suede/vectorize/low_geometry/__init__.py create mode 100644 arc_line_vectorization_suede/vectorize/low_geometry/beautify.py create mode 100644 arc_line_vectorization_suede/vectorize/low_geometry/fitting.py create mode 100644 arc_line_vectorization_suede/vectorize/low_geometry/manifest.py create mode 100644 arc_line_vectorization_suede/vectorize/low_geometry/primitives.py create mode 100644 arc_line_vectorization_suede/vectorize/low_geometry/residuals.py create mode 100644 arc_line_vectorization_suede/vectorize/low_geometry/routing.py create mode 100644 arc_line_vectorization_suede/vectorize/low_geometry/solve.py create mode 100644 arc_line_vectorization_suede/visualize.py diff --git a/.devcontainer/devcontainer.json b/.devcontainer/devcontainer.json index cb41494..32794b8 100644 --- a/.devcontainer/devcontainer.json +++ b/.devcontainer/devcontainer.json @@ -20,7 +20,7 @@ "editor.formatOnSave": true, "editor.tabSize": 2, "python.analysis.typeCheckingMode": "standard", - "python.formatting.blackArgs": [ + "python.formatting.blackArgs": [do "--line-length=88" ], "python.formatting.provider": "black" diff --git a/arc_line_vectorization_suede/.dependencies/requirements.txt b/arc_line_vectorization_suede/.dependencies/requirements.txt new file mode 100644 index 0000000..f005909 --- /dev/null +++ b/arc_line_vectorization_suede/.dependencies/requirements.txt @@ -0,0 +1,5 @@ +networkx==3.6.1 +numpy==2.4.4 +pillow==12.2.0 +scikit-image==0.26.0 +scipy==1.17.1 \ No newline at end of file diff --git a/arc_line_vectorization_suede/.github/workflows/subrepo-pull-into-main.yml b/arc_line_vectorization_suede/.github/workflows/subrepo-pull-into-main.yml new file mode 100644 index 0000000..5dfa356 --- /dev/null +++ b/arc_line_vectorization_suede/.github/workflows/subrepo-pull-into-main.yml @@ -0,0 +1,184 @@ +name: subrepo-pull-into-main +on: + push: + branches: [release] + +permissions: + contents: write + pull-requests: write + +jobs: + pull-to-main: + if: ${{ github.run_number > 1 }} + runs-on: ubuntu-latest + steps: + - name: Set up Git user + run: | + git config --global user.email "action@github.com" + git config --global user.name "GitHub Action" + + - name: Install git-subrepo + run: | + git clone https://github.com/ingydotnet/git-subrepo ~/.git-subrepo + source ~/.git-subrepo/.rc + git subrepo --version + + - uses: actions/checkout@v4 + with: + fetch-depth: 0 + ref: release + + - name: Revert (non-merge) commits from this push + id: revert-release + run: | + set -e + + # Commit at the tip of release when this workflow runs (i.e., AFTER) + TRIGGER_COMMIT=$(git rev-parse HEAD) + echo "Trigger commit: $TRIGGER_COMMIT" + echo "trigger-commit=$TRIGGER_COMMIT" >> "$GITHUB_OUTPUT" + + # Get BEFORE and AFTER from the push event + BEFORE="${{ github.event.before }}" + AFTER="${{ github.sha }}" + echo "Before: $BEFORE" + echo "After: $AFTER" + + # List non-merge commits introduced by this push (newest-first) + COMMITS=$(git rev-list --no-merges "$BEFORE..$AFTER" || true) + + if [ -z "$COMMITS" ]; then + echo "No non-merge commits to revert between $BEFORE and $AFTER." + fi + + echo "Non-merge commits found:" + echo "$COMMITS" + + # ---- Collect all commit messages (newest → oldest) ---- + ALL_MSGS="" + for c in $COMMITS; do + MSG=$(git log -1 --format=%B "$c") + ALL_MSGS="$ALL_MSGS + + --- + Commit: $c + $MSG + " + done + + echo "all-commit-messages<> "$GITHUB_OUTPUT" + echo "$ALL_MSGS" >> "$GITHUB_OUTPUT" + echo "EOF" >> "$GITHUB_OUTPUT" + + # ---- Revert the commits (newest-first) ---- + for c in $COMMITS; do + echo "Reverting $c" + git revert --no-edit "$c" + done + + # Capture the commit that represents the final revert state + REVERT_COMMIT=$(git rev-parse HEAD) + echo "revert-commit=$REVERT_COMMIT" >> "$GITHUB_OUTPUT" + + # Push the revert(s) back to release + git push origin release + + - uses: actions/checkout@v4 + with: + fetch-depth: 0 + ref: main + + - name: Check for release/ folder + run: | + if [ ! -d "release" ]; then + echo "release/ folder not found on main branch. Exiting early." + exit 0 + fi + + - name: Update ./release to latest + run: | + set -e + source ~/.git-subrepo/.rc + git subrepo pull ./release + + - name: Capture .gitrepo content + id: capture-gitrepo + run: | + set -e + # Save the current .gitrepo content + GITREPO_CONTENT=$(cat ./release/.gitrepo) + echo "gitrepo-content<> $GITHUB_OUTPUT + echo "$GITREPO_CONTENT" >> $GITHUB_OUTPUT + echo "EOF" >> $GITHUB_OUTPUT + + - name: Determine repository owner and name + id: repo-info + run: | + set -e + OWNER="${{ github.repository_owner }}" + NAME="${{ github.event.repository.name }}" + echo "owner=$OWNER" >> $GITHUB_OUTPUT + echo "name=$NAME" >> $GITHUB_OUTPUT + echo "Repository: $OWNER/$NAME" + + - name: Delete release folder + run: | + rm -rf ./release + echo "Deleted release/ folder" + + - name: Install trigger commit using degit.sh + run: | + set -e + # Download and execute degit.sh to install the trigger commit + bash <(curl -fsSL https://raw.githubusercontent.com/pmalacho-mit/suede/refs/heads/main/scripts/utils/degit.sh) \ + --repo "${{ steps.repo-info.outputs.owner }}/${{ steps.repo-info.outputs.name }}" \ + --commit "${{ steps.revert-release.outputs.trigger-commit }}" \ + --destination release + echo "Installed trigger commit into release/ folder" + + - name: Recreate .gitrepo file + run: | + set -e + # Write the original .gitrepo content + cat > ./release/.gitrepo << 'EOF' + ${{ steps.capture-gitrepo.outputs.gitrepo-content }} + EOF + + echo "Updated .gitrepo file:" + cat ./release/.gitrepo + + - name: Generate branch name + id: generate-branch-name + run: | + TIMESTAMP=$(date +%Y_%m_%d-%H_%M_%S_%Z) + BRANCH_NAME="chore/update-release-${TIMESTAMP}" + echo "branch-name=$BRANCH_NAME" >> $GITHUB_OUTPUT + echo "Using branch name: $BRANCH_NAME" + + - name: Commit and push changes + run: | + set -e + git switch -c "${{ steps.generate-branch-name.outputs.branch-name }}" + git add release/ + git commit -m "chore(suede): update release to trigger commits:\n ${{ steps.revert-release.outputs.trigger-commit }} + + Original commit messages: + ${{ steps.revert-release.outputs.all-commit-messages }}" + git push -u origin HEAD + + - name: Create PR + env: + GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} + run: | + gh pr create \ + --base main \ + --head "${{ steps.generate-branch-name.outputs.branch-name }}" \ + --title "chore(suede): pull latest upstream release into main" \ + --body "Automated subrepo pull of origin:release into ./release + + **Trigger commit:** ${{ steps.revert-release.outputs.trigger-commit }} + + **Original commit messages:** + \`\`\` + ${{ steps.revert-release.outputs.all-commit-messages }} + \`\`\`" diff --git a/arc_line_vectorization_suede/.gitignore b/arc_line_vectorization_suede/.gitignore new file mode 100644 index 0000000..badebb4 --- /dev/null +++ b/arc_line_vectorization_suede/.gitignore @@ -0,0 +1,4 @@ +.dependencies/* +!.dependencies/*.gitrepo +!.dependencies/package.json +!.dependencies/requirements.txt diff --git a/arc_line_vectorization_suede/.gitrepo b/arc_line_vectorization_suede/.gitrepo new file mode 100644 index 0000000..ed71f30 --- /dev/null +++ b/arc_line_vectorization_suede/.gitrepo @@ -0,0 +1,12 @@ +; DO NOT EDIT (unless you know what you are doing) +; +; This subdirectory is a git "subrepo", and this file is maintained by the +; git-subrepo command. See https://github.com/ingydotnet/git-subrepo#readme +; +[subrepo] + remote = https://github.com/mitmedialab/arc-line-vectorization-suede + branch = release + commit = b57054125d4e38c35462839e719384d72ebafc03 + parent = e1a6ea9e6760ea8f3f2cb47e8c7b84ecf8c68181 + method = merge + cmdver = 0.4.9 diff --git a/arc_line_vectorization_suede/__init__.py b/arc_line_vectorization_suede/__init__.py new file mode 100644 index 0000000..3889b67 --- /dev/null +++ b/arc_line_vectorization_suede/__init__.py @@ -0,0 +1,115 @@ +from .skeletonize import Skeletonize, ImageSource +from .segment import Segment +from .graph import StrokeGraph +from .vectorize.low_geometry import Vectorize as LowGeometryVectorize +from .vectorize.high_geometry import Vectorize as HighGeometryVectorize + +import numpy as np + + +def default_pipeline(source: ImageSource): + skeleton = Skeletonize( + source, + Skeletonize.Config.Binarize(threshold=0.5), + Skeletonize.Config.Skeletonize(method="zhang"), + Skeletonize.Config.Collapse( + skeletonize_method="lee", + max_hole_area=10, + max_thin_thickness=3.0, + reskeletonize=True, + ), + detect_config={ + "local_tau_radius": 40, + "fat_ratio": 1.3, + "min_fat_area": 8, + "group_dilate": 15, + "skel_ring_dilate": 5, + "pairing_tangent_steps": 8, + "pairing_threshold": 1.2, + "min_chromosome_skel_length": 15, + }, + ) + + junction_tol = 2.5 + tangent_sample = 10 + + segment = Segment( + skeleton.uncrossed, + skeleton.binary, + Segment.Config.Segment(min_length=10.0), + Segment.Config.Fuse( + max_path_length=20, + lookback=10, + min_tangent_score=0.5, + gap_penalty=0.05, + curvature_penalty=3.0, + ), + Segment.Config.Repair( + junction_tol=junction_tol, + stable_skip=2, + stable_sample=6, + max_junction_region_length=20, + min_output_polyline_length=2, + min_tangent_spread_deg=15.0, + interp_max_spacing=1.0, + min_curvature_spike_ratio=2.0, + curvature_context_window=8, + ), + Segment.Config.PostRepairFuse( + junction_tol=junction_tol, + tangent_skip=2, + tangent_sample=tangent_sample, + min_tangent_score=0.6, + curvature_penalty=1.0, + ), + ) + + graph = StrokeGraph( + segment.fused_post_repair, + StrokeGraph.Config.Build( + junction_tol=junction_tol, # match Repair.junction_tol + terminal_tangent_window=10, # should match fuse.lookback + crossing_tangent_skip=2, # baseline; dynamic walk handles arbitrary bridges + crossing_tangent_half_window=6, + cusp_angle_threshold_deg=50.0, # raise to ~50 to handle bikelove's cusp-like junction + cluster_merge_centroid_distance=10.0, + cluster_merge_index_gap=10, + ), + ) + + start_pos = np.array([0.0, 0.0]) + start_heading = 0.0 + low_geometry = LowGeometryVectorize( + graph, + start_pos=start_pos, + start_heading=start_heading, + ) + high_geometry = HighGeometryVectorize( + segment.fused_post_repair, + start_pos=start_pos, + start_heading=start_heading, + commands=HighGeometryVectorize.Config.ToCommands( + sigma=2.0, + corner_threshold=0.25, + max_fit_residual=5.0, + ), + consolidate=HighGeometryVectorize.Config.Consolidate( + center_tol_rel=0.25, + radius_tol_rel=0.25, + center_tol_abs=3.0, + radius_tol_abs=3.0, + max_endpoint_snap_rel=0.15, + max_endpoint_snap_abs=6.0, + proximity_min_radius_ratio=0.4, + line_angle_tol_deg=6.0, + line_offset_tol_abs=5.0, + min_line_length=5.0, + max_line_endpoint_snap_abs=5.0, + junction_epsilon=3.0, + merge_arcs=True, + merge_lines=True, + return_report=False, + ), + ) + + return skeleton, segment, graph, low_geometry, high_geometry diff --git a/arc_line_vectorization_suede/commands.py b/arc_line_vectorization_suede/commands.py new file mode 100644 index 0000000..5783e4c --- /dev/null +++ b/arc_line_vectorization_suede/commands.py @@ -0,0 +1,66 @@ +"""Type definitions for the drawing-robot pipeline. + +The output JSON shape mirrors the user's TypeScript discriminated union: + + type DrawingCommand = + | { kind: "line", distance: number, penDown: boolean } + | { kind: "spin", degrees: number } + | { kind: "arc", radius: number, degrees: number } + +Conventions: +- Coordinates are image-space pixel coordinates (x right, y down). +- Angles are measured in this frame: positive = counter-clockwise in the + (x-right, y-down) frame, which is clockwise when rendered to a screen. + This is consistent with SVG, so the renderer round-trips correctly. +- For arcs, `radius` is always positive; the sign of `degrees` indicates + direction (positive = curve to the left of current heading). +""" + +from __future__ import annotations +from dataclasses import dataclass +from typing import List, Literal, TypedDict, Union + +import numpy as np + + +# ---- Output JSON command types (match the TS discriminated union) ----------- + +class LineCommand(TypedDict): + kind: Literal["line"] + distance: float + penDown: bool + + +class SpinCommand(TypedDict): + kind: Literal["spin"] + degrees: float + + +class ArcCommand(TypedDict): + kind: Literal["arc"] + radius: float + degrees: float + + +DrawingCommand = Union[LineCommand, SpinCommand, ArcCommand] + + +# ---- Internal geometric primitives (used during fitting) -------------------- + +@dataclass +class LinePrimitive: + start: np.ndarray # shape (2,) + end: np.ndarray # shape (2,) + + +@dataclass +class ArcPrimitive: + center: np.ndarray # shape (2,) + radius: float + start: np.ndarray # shape (2,) — point on the circle + end: np.ndarray # shape (2,) — point on the circle + ccw: bool # direction of traversal from start to end + + +Primitive = Union[LinePrimitive, ArcPrimitive] +Stroke = List[Primitive] # connected chain of primitives sharing endpoints diff --git a/arc_line_vectorization_suede/graph/__init__.py b/arc_line_vectorization_suede/graph/__init__.py new file mode 100644 index 0000000..98294c2 --- /dev/null +++ b/arc_line_vectorization_suede/graph/__init__.py @@ -0,0 +1,621 @@ +"""Build a stroke graph from repaired polylines.""" + +from __future__ import annotations +from dataclasses import dataclass +from enum import Enum +from typing import Dict, List, Optional, Tuple, TypedDict + +import numpy as np +from numpy.typing import NDArray + + +class Role(Enum): + TERMINAL = "terminal" + CROSSING = "crossing" + CUSP = "cusp" + + +@dataclass +class Participation: + polyline_index: int + point_index: int + role: Role + terminal_end: Optional[int] + arc_length_t: float + tangent_in: NDArray[np.float64] + tangent_out: Optional[NDArray[np.float64]] + + +@dataclass +class Junction: + location: NDArray[np.float64] + participants: List[Participation] + + @property + def is_terminal(self): + return any(p.role == Role.TERMINAL for p in self.participants) + + @property + def is_anchored(self): + return any(p.role in (Role.TERMINAL, Role.CUSP) for p in self.participants) + + @property + def has_crossing(self): + return any(p.role == Role.CROSSING for p in self.participants) + + def pair_angles_deg(self): + def directional(p): + if p.role == Role.CROSSING and p.tangent_out is not None: + v = p.tangent_out - p.tangent_in + n = float(np.linalg.norm(v)) + if n > 1e-9: + return v / n + return p.tangent_in + + out = [] + n = len(self.participants) + for i in range(n): + ti = directional(self.participants[i]) + for j in range(i + 1, n): + tj = directional(self.participants[j]) + cos = float(np.clip(abs(np.dot(ti, tj)), 0.0, 1.0)) + deg = float(np.degrees(np.arccos(cos))) + out.append((i, j, deg)) + return out + + +class BuildConfig(TypedDict): + junction_tol: float + """Max cross-polyline distance for clustering into one junction.""" + + terminal_tangent_window: int + """Number of points to PCA-fit for the tangent at a TERMINAL endpoint.""" + + crossing_tangent_skip: int + """Baseline number of polyline points to skip on each side of an + interior rep before sampling for the tangent PCA. The sampling + point is then advanced FURTHER past any bridge points + (fractional-coord interpolation that ``Segment.Config.Repair`` + adds around an LS-solved jp) until the next real skeleton point + is reached — see ``_crossing_tangents``. The dynamic-bridge walk + means this baseline can stay small (e.g. 2-3) and the function + still handles arbitrarily long cascaded-repair bridges; the + baseline is for skipping past non-bridge local noise (a small + skeletonization jog right next to a real X crossing, etc). + Set to 0 to disable the baseline skip entirely (dynamic bridge + walk still applies).""" + + crossing_tangent_half_window: int + """Half-width of the tangent window on each side of a crossing + rep, applied AFTER ``crossing_tangent_skip`` is applied. The two + windows that PCA-fit ``tangent_in`` and ``tangent_out`` each + contain up to ``half_window + 1`` polyline points; larger = more + robust to noise but less local. ~6 is reasonable.""" + + cusp_angle_threshold_deg: float + """Deflection (degrees) above which an interior junction point + flips from CROSSING to CUSP. Default ~30°: catches real artistic + corners (typically 60-120°) while accepting smooth crossings (0°) + and modest noise.""" + + cluster_merge_centroid_distance: float + """Max centroid-to-centroid distance for two spatial clusters to + be considered as merge candidates by the post-clustering merge + pass. A single logical X-crossing whose polylines pass within + ``junction_tol`` at the entrance, briefly diverge to more than + ``junction_tol`` in the middle, then re-converge within + ``junction_tol`` at the exit ends up split into TWO clusters by + strict spatial clustering — this happens often after + ``resolve_crossings`` draws two close-running Bresenham lines + through a shallow-angle ribbon collapse. The merge pass identifies + such pairs and unions them so the graph reports one junction. + Set to 0 to disable. Default 10.0 (~4x junction_tol).""" + + cluster_merge_index_gap: int + """Max polyline-index gap between two clusters' index ranges for + them to be merged, applied PER shared polyline. Two clusters are + only unioned if every shared polyline has a small gap, which is + the test that distinguishes "one logical junction whose cluster + fragmented" (small gaps; the polylines hardly traverse any real + geometry between the fragments) from "two genuinely distinct + junctions involving the same polyline pair" (large gaps; the + polylines traverse real geometry between, e.g. a whisker that + crosses the head outline at two separate points). Default 10.""" + + +class StrokeGraph: + class Config: + class Build(BuildConfig): + pass + + polylines: List[NDArray[np.float64]] + junctions: List[Junction] + polyline_to_junctions: Dict[int, List[Tuple[int, int]]] + + def __init__(self, polylines, build): + self.polylines = [np.asarray(p, dtype=float) for p in polylines] + self.junctions = _build_junctions(self.polylines, build) + self.polyline_to_junctions = _index_polyline_to_junctions(self.junctions) + + def junctions_for(self, polyline_index): + out = [] + for ji, pi in self.polyline_to_junctions.get(polyline_index, []): + j = self.junctions[ji] + out.append((j, j.participants[pi])) + return out + + def is_closed(self, polyline_index): + return _is_closed_polyline(self.polylines[polyline_index]) + + def closed_polyline_indices(self): + return [i for i in range(len(self.polylines)) if self.is_closed(i)] + + def stats(self): + n_term_only = sum( + 1 for j in self.junctions if j.is_terminal and not j.has_crossing + ) + n_cross_only = sum( + 1 for j in self.junctions if j.has_crossing and not j.is_terminal + ) + n_mixed = sum(1 for j in self.junctions if j.is_terminal and j.has_crossing) + return ( + f"{len(self.polylines)} polylines, " + f"{len(self.junctions)} junctions " + f"({n_term_only} pure-terminal, " + f"{n_cross_only} pure-crossing, " + f"{n_mixed} mixed)" + ) + + +def _build_junctions(polylines, config): + clusters = _cluster_cross_polyline_points(polylines, config["junction_tol"]) + # Merge pass. When two clusters represent fragments of one logical + # junction (typically created by close-running Bresenham lines from + # resolve_crossings), strict spatial clustering leaves them + # separate; this stitches them back together. Disabled by setting + # cluster_merge_centroid_distance to 0. + merge_dist = float(config.get("cluster_merge_centroid_distance", 0.0)) + if merge_dist > 0 and len(clusters) > 1: + clusters = _merge_clusters_by_polyline_proximity( + clusters, + polylines, + max_centroid_distance=merge_dist, + max_index_gap=int(config.get("cluster_merge_index_gap", 10)), + ) + junctions = [] + for cluster in clusters: + cluster_xy = np.array([polylines[pi][ii] for pi, ii in cluster]) + centroid = cluster_xy.mean(axis=0) + per_poly_rep = {} + per_poly_rep_is_endpoint = {} + for pi, ii in cluster: + n = len(polylines[pi]) + is_endpoint = ii == 0 or ii == n - 1 + d2 = float(np.sum((polylines[pi][ii] - centroid) ** 2)) + if pi not in per_poly_rep: + per_poly_rep[pi] = ii + per_poly_rep_is_endpoint[pi] = is_endpoint + continue + cur_is_endpoint = per_poly_rep_is_endpoint[pi] + cur = per_poly_rep[pi] + cur_d2 = float(np.sum((polylines[pi][cur] - centroid) ** 2)) + if is_endpoint and not cur_is_endpoint: + per_poly_rep[pi] = ii + per_poly_rep_is_endpoint[pi] = True + elif is_endpoint == cur_is_endpoint and d2 < cur_d2: + per_poly_rep[pi] = ii + if len(per_poly_rep) < 2: + continue + + participants = [ + _make_participation(polylines[pi], pi, ii, config) + for pi, ii in per_poly_rep.items() + ] + rep_xy = np.array( + [polylines[p.polyline_index][p.point_index] for p in participants] + ) + location = rep_xy.mean(axis=0) + junctions.append(Junction(location=location, participants=participants)) + return junctions + + +def _index_polyline_to_junctions(junctions): + p2j = {} + for ji, j in enumerate(junctions): + for pi_local, part in enumerate(j.participants): + p2j.setdefault(part.polyline_index, []).append((ji, pi_local)) + return p2j + + +def _make_participation(poly, polyline_index, point_index, config): + n = len(poly) + is_closed = _is_closed_polyline(poly) + is_start = point_index == 0 + is_end = point_index == n - 1 + is_endpoint_index = is_start or is_end + + if is_endpoint_index and not is_closed: + terminal_end = 0 if is_start else 1 + tangent_in = _terminal_tangent( + poly, terminal_end, config["terminal_tangent_window"] + ) + return Participation( + polyline_index=polyline_index, + point_index=point_index, + role=Role.TERMINAL, + terminal_end=terminal_end, + arc_length_t=0.0 if is_start else 1.0, + tangent_in=tangent_in, + tangent_out=None, + ) + + skip = config.get("crossing_tangent_skip", 0) + if is_endpoint_index and is_closed: + tangent_in, tangent_out = _closed_loop_tangents( + poly, skip, config["crossing_tangent_half_window"] + ) + arc_length_t = 0.0 + else: + tangent_in, tangent_out = _crossing_tangents( + poly, point_index, skip, config["crossing_tangent_half_window"] + ) + arc_length_t = _arc_length_t(poly, point_index) + + if _is_cusp(tangent_in, tangent_out, config["cusp_angle_threshold_deg"]): + role = Role.CUSP + else: + role = Role.CROSSING + + return Participation( + polyline_index=polyline_index, + point_index=point_index, + role=role, + terminal_end=None, + arc_length_t=arc_length_t, + tangent_in=tangent_in, + tangent_out=tangent_out, + ) + + +def _is_cusp(tangent_in, tangent_out, threshold_deg): + incoming_dir = -tangent_in + cos = float(np.clip(np.dot(incoming_dir, tangent_out), -1.0, 1.0)) + deflection_deg = float(np.degrees(np.arccos(cos))) + return deflection_deg >= threshold_deg + + +def _cluster_cross_polyline_points(polylines, tol): + entries = [] + coords_list = [] + for pi, poly in enumerate(polylines): + for ii in range(len(poly)): + entries.append((pi, ii)) + coords_list.append(poly[ii]) + if not coords_list: + return [] + coords = np.asarray(coords_list, dtype=float) + + cell = max(tol, 1e-9) + grid = {} + for idx in range(len(coords)): + key = (int(coords[idx, 0] // cell), int(coords[idx, 1] // cell)) + grid.setdefault(key, []).append(idx) + + parent = list(range(len(coords))) + + def find(x): + while parent[x] != x: + parent[x] = parent[parent[x]] + x = parent[x] + return x + + def union(a, b): + ra, rb = find(a), find(b) + if ra != rb: + parent[ra] = rb + + tol2 = tol * tol + for (cx, cy), idxs in list(grid.items()): + candidates = [] + for dx in (-1, 0, 1): + for dy in (-1, 0, 1): + key = (cx + dx, cy + dy) + if key in grid: + candidates.extend(grid[key]) + for i in idxs: + pi_i = entries[i][0] + xi, yi = coords[i, 0], coords[i, 1] + for j in candidates: + if j <= i: + continue + if entries[j][0] == pi_i: + continue + dxv = xi - coords[j, 0] + dyv = yi - coords[j, 1] + if dxv * dxv + dyv * dyv < tol2: + union(i, j) + + groups = {} + for idx in range(len(coords)): + groups.setdefault(find(idx), []).append(idx) + out = [] + for members in groups.values(): + cluster = [entries[i] for i in members] + polys_in_cluster = {pi for pi, _ in cluster} + if len(polys_in_cluster) >= 2: + out.append(cluster) + return out + + +_CLOSED_TOL = 1.5 + + +def _merge_clusters_by_polyline_proximity( + clusters, + polylines, + max_centroid_distance, + max_index_gap, +): + """Union spatial clusters that fragment a single logical junction. + + A single X-crossing whose polylines briefly diverge to more than + ``junction_tol`` in the middle (typical of resolve_crossings' + close-running Bresenham lines through a shallow-angle ribbon + collapse) ends up split into two clusters by strict spatial + clustering. This pass identifies such pairs by three checks and + unions them: + + 1. They share ≥2 polylines (so it's the same pair of strokes + in both, not unrelated junctions). + 2. Their centroids are within ``max_centroid_distance``. + 3. For EVERY shared polyline, the polyline-index gap between + the two clusters' ranges in that polyline is within + ``max_index_gap``. This is the test that distinguishes + fragmented one-junction (small gaps) from genuinely two + junctions involving the same pair of polylines (large + gaps; the polyline traverses real geometry between them). + + Clusters that don't merge with anything pass through unchanged. + """ + if len(clusters) <= 1: + return clusters + + # Pre-compute centroid and per-polyline (min_idx, max_idx) ranges. + centroids = [] + ranges = [] + for cl in clusters: + xy = np.array([polylines[pi][ii] for pi, ii in cl]) + centroids.append(xy.mean(axis=0)) + per_poly = {} + for pi, ii in cl: + per_poly.setdefault(pi, []).append(ii) + ranges.append({pi: (min(ix), max(ix)) for pi, ix in per_poly.items()}) + + parent = list(range(len(clusters))) + + def find(x): + while parent[x] != x: + parent[x] = parent[parent[x]] + x = parent[x] + return x + + def union(a, b): + ra, rb = find(a), find(b) + if ra != rb: + parent[ra] = rb + + max_dist_sq = max_centroid_distance * max_centroid_distance + n = len(clusters) + for i in range(n): + for j in range(i + 1, n): + shared = set(ranges[i].keys()) & set(ranges[j].keys()) + if len(shared) < 2: + continue + dx = centroids[i][0] - centroids[j][0] + dy = centroids[i][1] - centroids[j][1] + if dx * dx + dy * dy > max_dist_sq: + continue + ok = True + for pi in shared: + lo_i, hi_i = ranges[i][pi] + lo_j, hi_j = ranges[j][pi] + if lo_i > hi_j: + gap = lo_i - hi_j + elif lo_j > hi_i: + gap = lo_j - hi_i + else: + gap = 0 # ranges overlap + if gap > max_index_gap: + ok = False + break + if ok: + union(i, j) + + groups = {} + for i in range(n): + groups.setdefault(find(i), []).append(i) + + merged = [] + for indices in groups.values(): + combined = [] + for i in indices: + combined.extend(clusters[i]) + merged.append(combined) + return merged + + +def _is_closed_polyline(poly): + if len(poly) < 3: + return False + return float(np.linalg.norm(poly[0] - poly[-1])) < _CLOSED_TOL + + +def _arc_length_t(polyline, idx): + if len(polyline) <= 1: + return 0.0 + diffs = np.diff(polyline, axis=0) + seglens = np.linalg.norm(diffs, axis=1) + cum = np.concatenate([[0.0], np.cumsum(seglens)]) + total = float(cum[-1]) + if total <= 1e-9: + return 0.0 + return float(cum[idx] / total) + + +def _pca_unit_tangent(samples, orient_toward): + if len(samples) < 2: + return np.array([1.0, 0.0]) + centered = samples - samples.mean(axis=0) + _, _, Vt = np.linalg.svd(centered, full_matrices=False) + t = Vt[0] + n = float(np.linalg.norm(t)) + if n < 1e-9: + return np.array([1.0, 0.0]) + t = t / n + if float(np.dot(t, orient_toward)) < 0: + t = -t + return t + + +def _terminal_tangent(polyline, terminal_end, window): + n = len(polyline) + if n < 2: + return np.array([1.0, 0.0]) + if terminal_end == 0: + samples = polyline[: min(n, window + 1)] + orient = samples[-1] - samples[0] + else: + samples = polyline[max(0, n - window - 1) :] + orient = samples[0] - samples[-1] + return _pca_unit_tangent(samples, orient) + + +def _is_integer_point(p, tol: float = 0.05) -> bool: + """True if a polyline point's coords are at integer pixel positions. + + Polylines coming out of ``trace_skeleton`` sit at integer pixel + coords (the skeleton is a boolean image). ``repair_junctions`` + REPLACES the strict spatial-cluster region of each polyline with + a single LS-solved jp (fractional, since it's the algebraic + intersection of approach lines) plus 1px-spaced linear + interpolation back to the polyline's real interior (also + fractional). So "is this point at integer coords" is a reliable + test for "is this point original skeleton geometry" vs "is this + point repair-imposed". + """ + return abs(p[0] - round(p[0])) < tol and abs(p[1] - round(p[1])) < tol + + +def _crossing_tangents(polyline, idx, skip, half_window): + """Outward unit tangents on either side of an interior crossing point. + + Returns ``(tangent_in, tangent_out)`` — both pointing AWAY from + ``polyline[idx]``; ``tangent_in`` toward lower indices, + ``tangent_out`` toward higher indices. + + Bridge handling (this is most of the function's reason to exist): + after ``Segment.Config.Repair`` runs, polyline indices near a + junction can be bridge interpolation between the polyline's real + interior and the LS-solved jp — they are 1px-spaced and nearly + colinear with the repair-imposed direction, so a PCA window that + includes them is biased away from the polyline's natural geometry. + A cascade of nearby repairs can produce a long stretch of bridge + points on either side of an interior rep (up to 20+ points per + side), so a fixed ``skip`` is unreliable. Instead, this function: + + 1. Steps ``skip`` points outward from the rep on each side (the + user-configured baseline to clear short bridges and any + immediate noise). + 2. Then keeps walking outward past any further bridge points, + detected by their fractional coordinates (real skeleton + pixels are at integer positions, bridge interpolation is + not — see ``_is_integer_point``). + 3. Samples up to ``half_window + 1`` points beyond that. + + The ``skip`` argument is still useful for skipping past + NON-bridge local noise (e.g. small skeletonization jogs at an X + crossing). For bridge-only contamination ``skip=0`` is enough + because the dynamic walk handles the bridge. + """ + n = len(polyline) + if n < 2: + return np.array([1.0, 0.0]), np.array([1.0, 0.0]) + here = polyline[idx] + + def find_real_sample_anchor(direction: int) -> int: + """Walk outward from idx in `direction` (+1 or -1): apply the + baseline skip, then continue walking past any bridge (non- + integer) points until a real skeleton point is found or the + polyline boundary is reached.""" + cur = idx + direction * skip + # Clip into polyline bounds. + cur = max(0, min(n - 1, cur)) + # Skip further bridge points (non-integer). + while 0 < cur < n - 1 and not _is_integer_point(polyline[cur]): + nxt = cur + direction + if nxt < 0 or nxt >= n: + break + cur = nxt + return cur + + # In side (toward lower indices). + in_hi = find_real_sample_anchor(-1) + in_lo = max(0, in_hi - half_window) + if in_hi > in_lo: + samples_in = polyline[in_lo : in_hi + 1] + orient_in = polyline[in_lo] - here + tangent_in = _pca_unit_tangent(samples_in, orient_in) + else: + tangent_in = np.array([1.0, 0.0]) + + # Out side (toward higher indices). + out_lo = find_real_sample_anchor(+1) + out_hi = min(n - 1, out_lo + half_window) + if out_hi > out_lo: + samples_out = polyline[out_lo : out_hi + 1] + orient_out = polyline[out_hi] - here + tangent_out = _pca_unit_tangent(samples_out, orient_out) + else: + tangent_out = np.array([1.0, 0.0]) + + return tangent_in, tangent_out + + +def _closed_loop_tangents(polyline, skip, half_window): + """Tangents at the closure point of a closed polyline. + + Same intent and bridge-handling logic as ``_crossing_tangents``, + but the "in" side wraps around to the END of the polyline (the + duplicate of index 0 lives at index n-1) so we walk backward from + n-2 rather than from idx-1. + """ + n = len(polyline) + if n < 4: + return np.array([1.0, 0.0]), np.array([1.0, 0.0]) + here = polyline[0] + + # Forward side: from index 0 going up, stop short of n-1. + fwd_anchor = min(max(0, n - 2), skip + 1) + while fwd_anchor < n - 2 and not _is_integer_point(polyline[fwd_anchor]): + fwd_anchor += 1 + fwd_hi = min(n - 2, fwd_anchor + half_window) + if fwd_hi > fwd_anchor: + fwd_samples = polyline[fwd_anchor : fwd_hi + 1] + orient_fwd = fwd_samples[-1] - here + tangent_out = _pca_unit_tangent(fwd_samples, orient_fwd) + else: + tangent_out = np.array([1.0, 0.0]) + + # Backward side: from index n-2 going down. + bwd_anchor = max(1, n - 2 - skip) + while bwd_anchor > 1 and not _is_integer_point(polyline[bwd_anchor]): + bwd_anchor -= 1 + bwd_lo = max(1, bwd_anchor - half_window) + if bwd_anchor > bwd_lo: + bwd_samples = polyline[bwd_lo : bwd_anchor + 1] + orient_bwd = bwd_samples[0] - here + tangent_in = _pca_unit_tangent(bwd_samples, orient_bwd) + else: + tangent_in = np.array([1.0, 0.0]) + + return tangent_in, tangent_out diff --git a/arc_line_vectorization_suede/segment/__init__.py b/arc_line_vectorization_suede/segment/__init__.py new file mode 100644 index 0000000..ac948dd --- /dev/null +++ b/arc_line_vectorization_suede/segment/__init__.py @@ -0,0 +1,135 @@ +"""Split a skeleton into segments at branch points and visualize them. + +A "segment" is a maximal run of degree-2 skeleton pixels bounded by node +pixels at each end. A node pixel is one of: + - an endpoint: a skeleton pixel with exactly one 8-connected skeleton + neighbour (degree 1). + - a branch: a skeleton pixel with three or more 8-connected skeleton + neighbours (degree >= 3). Adjacent branch pixels are merged into a + single super-junction so that thick junctions in the skeleton don't + spawn spurious tiny segments. + +``segment_skeleton`` is a thin wrapper around ``trace_skeleton`` that also +applies an optional minimum-length filter. ``visualize_segments`` produces +a PNG where each segment is drawn in a distinct colour using golden-ratio +hue spacing so adjacent segments are easy to tell apart visually. + +The ``Segment`` class runs the full pipeline: + + skeleton + binary + | + v + [trace] self.segments + | + v + [fuse 1st pass] self.fused (with self.accepted, self.all) + | A* through binary; bridges noisy junctions. + v + [repair] self.repaired + | Per-junction LS solve, replaces noisy regions. + v + [resolve self.cleaned (with self.collapsed_chromosomes) + chromosomes] Collapses skeletonization "centromere" artifacts + | where two strokes overlapped along a length. + v + [fuse 2nd pass] self.joined (with self.joined_accepted) + Pairwise tangent matching at clean junctions; + no path-finding, just cluster-and-match. +""" + +from __future__ import annotations + +from typing import TypedDict, List + +import numpy as np +from numpy.typing import NDArray + +from .trace import trace_skeleton +from .fusion import ( + fuse_segments, + FuseConfig, + fuse_post_repair, + PostRepairFuseConfig, +) +from .repair import repair_junctions, RepairConfig + + +def filter_short_polylines( + polylines: List[np.ndarray], + min_length: float = 3.0, +) -> List[np.ndarray]: + """Drop polylines whose total arc length is below `min_length` pixels. + + These are almost always skeletonization artifacts at junctions. + """ + out = [] + for p in polylines: + if len(p) < 2: + continue + diffs = np.diff(p, axis=0) + length = float(np.sum(np.linalg.norm(diffs, axis=1))) + if length >= min_length: + out.append(p) + return out + + +def segment_skeleton( + skel: NDArray[np.bool_], + min_length: float = 0.0, +) -> List[NDArray[np.float64]]: + """Split a skeleton into polyline segments at every node (branch / endpoint). + + Each returned segment is an (N, 2) float array of (x, y) pixel + coordinates. The first and last point of every segment is a node + pixel; every interior point is degree-2 (exactly two 8-connected + skeleton neighbours). Segments share their endpoint pixels with the + other segments that meet at the same node — i.e. the segments form + the edge set of a graph whose vertices are the node pixels. + + Args: + skel: boolean skeleton image (True = ink). + min_length: drop segments whose total arc length is below this + value in pixels. Useful for filtering occasional junction + noise. Default 0 keeps everything. + """ + polys = trace_skeleton(skel.astype(bool)) + return ( + filter_short_polylines(polys, min_length=min_length) + if min_length > 0 + else polys + ) + + +class Segment: + class Config: + class Segment(TypedDict): + min_length: float + + class Fuse(FuseConfig): + pass + + class Repair(RepairConfig): + pass + + class PostRepairFuse(PostRepairFuseConfig): + pass + + def __init__( + self, + skel: NDArray[np.bool_], + binary: NDArray[np.bool_], + segment: Config.Segment, + fuse: Config.Fuse, + repair: Config.Repair, + post_repair_fuse: Config.PostRepairFuse, + ): + self.segments = segment_skeleton(skel, min_length=segment["min_length"]) + + self.fused_pre_repair, self.accepted, self.all = fuse_segments( + self.segments, binary=binary, config=fuse + ) + self.repaired = repair_junctions(self.fused_pre_repair, config=repair) + + self.fused_post_repair, self.joined_accepted = fuse_post_repair( + self.repaired, config=post_repair_fuse + ) diff --git a/arc_line_vectorization_suede/segment/_helpers.py b/arc_line_vectorization_suede/segment/_helpers.py new file mode 100644 index 0000000..d9f80eb --- /dev/null +++ b/arc_line_vectorization_suede/segment/_helpers.py @@ -0,0 +1,239 @@ +"""Shared spatial helpers for segment processing modules. + +Used by: + +- ``repair.py``: cluster ALL polyline points to discover junctions, + including the interior points where a polyline passes smoothly + through another polyline's endpoint. + +- ``chromosome.py``: cluster polyline ENDPOINTS to find the + multi-polyline junction clusters that bracket a candidate + centromere. + +- ``fusion.py`` (post-repair pass): cluster polyline ENDPOINTS to + find polylines that share a clean post-repair junction location, + which are then matched pairwise by tangent alignment. + +The leading underscore in the module name indicates these are package- +private; downstream consumers should not import from this module. +""" + +from __future__ import annotations +from typing import Dict, List, Tuple + +import numpy as np +from numpy.typing import NDArray + +# ============================================================================ +# Point gathering +# ============================================================================ + + +def gather_points( + polylines: List[NDArray[np.float64]], +) -> Tuple[List[Tuple[int, int]], NDArray[np.float64]]: + """Flatten every point of every polyline into one coords array, + paired with metadata of ``(poly_idx, point_idx)`` tuples. + + Used by junction repair, which has to discover crossings where the + junction sits at an INTERIOR index of one or more polylines (a + whisker piercing a head outline; the head outline is a long + polyline whose junction point isn't an endpoint). + """ + meta: List[Tuple[int, int]] = [] + coords_list: List[NDArray[np.float64]] = [] + for pi, poly in enumerate(polylines): + for ii in range(len(poly)): + meta.append((pi, ii)) + coords_list.append(poly[ii]) + if not coords_list: + return meta, np.empty((0, 2), dtype=float) + return meta, np.asarray(coords_list, dtype=float) + + +def gather_endpoints( + polylines: List[NDArray[np.float64]], +) -> Tuple[List[Tuple[int, int]], NDArray[np.float64]]: + """Flatten polyline endpoints (start and end of each polyline) with + parallel ``(poly_idx, end_side)`` metadata. ``end_side`` is 0 for + the polyline's first point, 1 for its last. + + Polylines with fewer than 2 points are skipped (no meaningful + endpoint pair). + + Used by post-repair fusion and chromosome detection: after repair, + polylines that share a junction have exact (or near-exact) shared + endpoint coordinates, so clustering JUST the endpoints is both + sufficient and significantly cheaper than clustering all points. + """ + meta: List[Tuple[int, int]] = [] + coords_list: List[NDArray[np.float64]] = [] + for pi, poly in enumerate(polylines): + if len(poly) < 2: + continue + meta.append((pi, 0)) + coords_list.append(poly[0]) + meta.append((pi, 1)) + coords_list.append(poly[-1]) + if not coords_list: + return meta, np.empty((0, 2), dtype=float) + return meta, np.asarray(coords_list, dtype=float) + + +# ============================================================================ +# Spatial clustering +# ============================================================================ + + +def spatial_clusters( + coords: NDArray[np.float64], + tol: float, + point_poly_idx: List[int], +) -> List[List[int]]: + """Union-find clustering of 2D points with a cross-polyline rule. + + Two points are unioned IFF (a) they belong to DIFFERENT polylines + AND (b) they are within ``tol`` of each other. The cross-polyline + restriction is exactly what makes the result correspond to + JUNCTIONS: a long polyline that folds back close to itself is not + unioned to itself, so it doesn't spuriously generate a junction at + every fold. + + Uses a hashed grid of cell size ``tol`` so each point only checks + points in its own 3x3 neighborhood of grid cells, keeping the cost + near-linear in the number of input points. + + Returns a list of clusters, each a list of indices into ``coords``. + Singleton clusters (one point) ARE returned; callers that want only + multi-polyline clusters must filter themselves. + """ + n = len(coords) + if n == 0: + return [] + cell_size = max(tol, 1e-9) + grid: Dict[Tuple[int, int], List[int]] = {} + for idx in range(n): + c = (int(coords[idx, 0] // cell_size), int(coords[idx, 1] // cell_size)) + grid.setdefault(c, []).append(idx) + parent = list(range(n)) + + def find(x: int) -> int: + while parent[x] != x: + parent[x] = parent[parent[x]] + x = parent[x] + return x + + def union(a: int, b: int) -> None: + ra, rb = find(a), find(b) + if ra != rb: + parent[ra] = rb + + tol_sq = tol * tol + for (cx, cy), idxs in list(grid.items()): + candidates: List[int] = [] + for dx in (-1, 0, 1): + for dy in (-1, 0, 1): + key = (cx + dx, cy + dy) + if key in grid: + candidates.extend(grid[key]) + for i in idxs: + xi, yi = coords[i, 0], coords[i, 1] + pi_i = point_poly_idx[i] + for j in candidates: + if j <= i: + continue + if point_poly_idx[j] == pi_i: + continue + dx_ = xi - coords[j, 0] + dy_ = yi - coords[j, 1] + if dx_ * dx_ + dy_ * dy_ < tol_sq: + union(i, j) + groups: Dict[int, List[int]] = {} + for idx in range(n): + groups.setdefault(find(idx), []).append(idx) + return list(groups.values()) + + +# ============================================================================ +# Stable endpoint tangent (post-repair-aware) +# ============================================================================ + + +def stable_endpoint_tangent_and_curvature( + polyline: NDArray[np.float64], + end_side: int, + skip: int, + sample: int, +) -> Tuple[NDArray[np.float64], float]: + """Outward unit tangent and unsigned local curvature at one of a + polyline's endpoints, estimated from a window of stable interior + points past any post-repair "bridge". + + Background: ``repair_junctions`` replaces each junction-affected + polyline region with a single LS-solved junction point, then + interpolates 1-pixel-spaced points between that single point and + the polyline's stable interior. Those bridge points are nearly + colinear and lie BETWEEN the polyline's natural geometry and the + LS junction location - so PCA over a window that's mostly bridge + points returns a tangent pointing at the OLD junction location, + not in the polyline's true direction. For post-repair fusion and + chromosome detection we need a tangent estimate that ignores the + bridge. + + Strategy: skip the first ``skip`` points from the endpoint, then + PCA-fit the next ``sample`` points. With ``skip ~ 5`` (typical + bridge length is 4-6 points for the repair interp_max_spacing of + 1.0 px) and ``sample ~ 10``, the window is genuinely stable. + + Outward orientation: the returned tangent points TOWARD the chosen + endpoint (i.e. in the direction the stroke "leaves" the polyline + body, away from the interior), which is the convention fusion uses + for scoring anti-parallel matches. + + Curvature via the small-arc sagitta-to-curvature relation: + kappa ~ 8 * s / L^2 where s is the maximum perpendicular deviation + of the sample window from its PCA principal axis and L is the + chord length of the window. Returns 0 for windows too short or + degenerate to support a meaningful curvature estimate. + """ + n = len(polyline) + if n < 2: + return np.array([1.0, 0.0]), 0.0 + + if end_side == 0: + lo = min(skip, max(0, n - 2)) + hi = min(lo + sample, n) + else: + hi = max(n - skip, 2) + lo = max(hi - sample, 0) + + pts = polyline[lo:hi] + if len(pts) < 2: + # Fallback: full-polyline chord. + if end_side == 0: + d = polyline[-1] - polyline[0] + else: + d = polyline[0] - polyline[-1] + norm = float(np.linalg.norm(d)) + return (d / norm if norm > 1e-9 else np.array([1.0, 0.0])), 0.0 + + centroid = pts.mean(axis=0) + centered = pts - centroid + _, _, Vt = np.linalg.svd(centered, full_matrices=False) + tangent = Vt[0] + + endpoint = polyline[0] if end_side == 0 else polyline[-1] + ref = endpoint - centroid + if float(np.dot(tangent, ref)) < 0: + tangent = -tangent + + if len(pts) < 3: + return tangent, 0.0 + perp = Vt[1] + perp_dists = np.abs(centered @ perp) + max_perp = float(perp_dists.max()) + chord = float(np.linalg.norm(pts[-1] - pts[0])) + if chord < 1.0: + return tangent, 0.0 + curvature = 8.0 * max_perp / (chord * chord) + return tangent, curvature diff --git a/arc_line_vectorization_suede/segment/fusion.py b/arc_line_vectorization_suede/segment/fusion.py new file mode 100644 index 0000000..a896ebc --- /dev/null +++ b/arc_line_vectorization_suede/segment/fusion.py @@ -0,0 +1,622 @@ +"""Fuse adjacent skeleton segments into longer continuous strokes. + +After skeletonization and segmentation we have a list of polylines, each +running between two node pixels. At a branch where multiple segments meet, +the original drawing was likely a single continuous pen stroke that crossed +through; the skeletonizer broke it because the local topology is a 'star'. + +This module re-fuses such segments. The matching rule has two parts: + + (1) Reachability in the *original binary* (not just the skeleton): two + endpoints are candidates only if a short 8-connected path through + ink connects them. This handles "noisy" junctions where the two + segments have endpoints separated by a thick blob of original ink + that the skeletonizer chewed up. + + (2) Tangent alignment at the endpoints: of all candidates, the best + match is the one whose outward tangent is most anti-parallel to + the target's outward tangent, i.e. the one that looks like the + smoothest continuation of the same stroke. + +A greedy maximum-weight matching picks pairs starting from the highest +score, so at a 4-way junction the two best-aligned segments fuse first +and the remaining two fuse second (or stay separate). + +This module also exposes ``fuse_post_repair``, a second-pass fusion that +runs AFTER junction repair (and after chromosome resolution). It drops +the A*-through-binary path-finding because by then the polylines that +should join already share a clean endpoint coordinate; the only question +is which pairs at each shared location are the best-aligned continuation. +See its docstring for details. +""" + +from __future__ import annotations +import heapq +import math +from dataclasses import dataclass +from typing import List, Optional, Set, Tuple, TypedDict + +import numpy as np +from numpy.typing import NDArray + +from ._helpers import ( + gather_endpoints, + spatial_clusters, + stable_endpoint_tangent_and_curvature, +) + +# ============================================================================ +# Types +# ============================================================================ + + +@dataclass +class FusionCandidate: + """A potential fusion between two segment endpoints. + + Attributes: + seg_a, end_a: index of the segment and which side (0 = polyline + start, 1 = polyline end) of its two endpoints is involved. + seg_b, end_b: same for the other segment. + connecting_path: (K, 2) (x, y) pixels running from endpoint A to + endpoint B through the original binary ink. Inclusive of both + endpoints. 8-connected pixel-by-pixel. EMPTY for post-repair + fusion candidates (which join already-co-located endpoints). + path_length: number of pixels in `connecting_path`. 0 for + post-repair candidates. + tangent_score: in [-1, 1]. +1 = perfectly smooth continuation + (outward tangents anti-parallel); 0 = perpendicular; + -1 = the two segments fold back the same way at the junction. + curvature_a, curvature_b: local curvature magnitudes at each + endpoint (1/pixel units). Used by the score to penalize + matches between segments of very different curvature. + score: combined sort key for the greedy matcher. + """ + + seg_a: int + end_a: int + seg_b: int + end_b: int + connecting_path: NDArray[np.float64] + path_length: float + tangent_score: float + curvature_a: float + curvature_b: float + score: float + + +# ============================================================================ +# Helpers +# ============================================================================ + +NEIGHBOR_OFFSETS = [ + (-1, -1), + (-1, 0), + (-1, 1), + (0, -1), + (0, 1), + (1, -1), + (1, 0), + (1, 1), +] + + +def _shortest_path_through_binary( + binary: NDArray[np.bool_], + start: Tuple[int, int], + goal: Tuple[int, int], + max_cost: float, +) -> Optional[List[Tuple[int, int]]]: + """Optimal 8-connected shortest path through True pixels. + + Uses A* with Euclidean step costs (1.0 for orthogonal moves, sqrt(2) + for diagonals) and the octile-distance heuristic — the tightest + admissible heuristic for this cost structure, which guarantees A* + returns a true shortest-Euclidean-distance path while pruning the + search aggressively. + + Returns the path as a list of (y, x) tuples (inclusive of both + endpoints), or None if `goal` is not reachable from `start` through + True pixels within `max_cost` total Euclidean distance. + """ + H, W = binary.shape + if not (0 <= start[0] < H and 0 <= start[1] < W and binary[start]): + return None + if not (0 <= goal[0] < H and 0 <= goal[1] < W and binary[goal]): + return None + if start == goal: + return [start] + + SQRT2 = math.sqrt(2.0) + # (dy, dx, step_cost) + moves = ( + (-1, -1, SQRT2), + (-1, 0, 1.0), + (-1, 1, SQRT2), + (0, -1, 1.0), + (0, 1, 1.0), + (1, -1, SQRT2), + (1, 0, 1.0), + (1, 1, SQRT2), + ) + + def octile(p: Tuple[int, int]) -> float: + dy = abs(p[0] - goal[0]) + dx = abs(p[1] - goal[1]) + # octile distance = max + (sqrt(2) - 1) * min + if dy > dx: + return dy + (SQRT2 - 1.0) * dx + return dx + (SQRT2 - 1.0) * dy + + # Priority queue items: (f_score, tiebreak_counter, node) + open_pq: List[Tuple[float, int, Tuple[int, int]]] = [] + counter = 0 + g_score: dict = {start: 0.0} + parent: dict = {start: None} + heapq.heappush(open_pq, (octile(start), counter, start)) + + while open_pq: + _, _, cur = heapq.heappop(open_pq) + if cur == goal: + path: List[Tuple[int, int]] = [goal] + while parent[path[-1]] is not None: + path.append(parent[path[-1]]) + path.reverse() + return path + cur_g = g_score[cur] + # Stale entry (re-inserted with a worse g)? skip. + # (We don't need an explicit closed set because we only relax when + # we strictly improve g_score below.) + cy, cx = cur + for dy, dx, step_cost in moves: + ny, nx = cy + dy, cx + dx + if not (0 <= ny < H and 0 <= nx < W): + continue + if not binary[ny, nx]: + continue + tentative_g = cur_g + step_cost + if tentative_g > max_cost: + continue + nbr = (ny, nx) + if nbr in g_score and tentative_g >= g_score[nbr]: + continue + g_score[nbr] = tentative_g + parent[nbr] = cur + f_new = tentative_g + octile(nbr) + counter += 1 + heapq.heappush(open_pq, (f_new, counter, nbr)) + return None + + +def _local_tangent_and_curvature( + polyline: NDArray[np.float64], + at_end: bool, + lookback: int, +) -> Tuple[NDArray[np.float64], float]: + """Estimate outward tangent and unsigned local curvature near an endpoint. + + Uses PCA on the last (or first) `lookback` pixels rather than a simple + two-point difference, which makes the tangent much more robust to + single-pixel "jog" artifacts that skeletonization sometimes introduces + near junctions. The principal axis dominates as long as the bulk of the + pixels follow the segment's true direction, even if the very last one + or two pixels drift sideways. + + Curvature is approximated by the small-arc sagitta-to-curvature + relation: for a circular arc of chord length L and maximum sagitta s, + kappa = 1/R ~= 8 s / L^2. We measure s as the largest perpendicular + deviation of the lookback pixels from their PCA principal axis. + Returns 0 when the polyline is too short to support a meaningful + estimate. + """ + n = min(lookback, len(polyline)) + if n < 3: + if at_end: + d = polyline[-1] - polyline[0] + else: + d = polyline[0] - polyline[-1] + norm = float(np.linalg.norm(d)) + return (d / norm if norm > 1e-9 else np.zeros(2)), 0.0 + + pts = polyline[-n:] if at_end else polyline[:n] + + # PCA on the lookback window. The principal axis is the tangent + # direction (unsigned); the perpendicular axis is what we measure + # sagitta against. + centroid = pts.mean(axis=0) + centered = pts - centroid + _, _, Vt = np.linalg.svd(centered, full_matrices=False) + tangent = Vt[0] + perp = Vt[1] + + # Orient the tangent outward (i.e. pointing away from the segment's + # interior, toward the chosen endpoint). + ref = (pts[-1] - pts[0]) if at_end else (pts[0] - pts[-1]) + if float(np.dot(tangent, ref)) < 0: + tangent = -tangent + + # Curvature magnitude proxy + perp_dists = np.abs(centered @ perp) + max_perp = float(perp_dists.max()) + chord = float(np.linalg.norm(pts[-1] - pts[0])) + if chord < 1.0: + return tangent, 0.0 + curvature = 8.0 * max_perp / (chord * chord) + return tangent, curvature + + +# ============================================================================ +# Candidate enumeration & matching (FIRST PASS) +# ============================================================================ + + +class FuseConfig(TypedDict): + max_path_length: int + lookback: int + gap_penalty: float + min_tangent_score: float + curvature_penalty: float + + +def find_fusion_candidates( + segments: List[NDArray[np.float64]], + binary: NDArray[np.bool_], + max_path_length: int, + lookback: int, + gap_penalty: float, + curvature_penalty: float, +) -> List[FusionCandidate]: + """Enumerate every (endpoint, endpoint) pair that is: + (1) within 2 * max_path_length Euclidean distance, AND + (2) reachable through ink in <= max_path_length BFS steps. + Each pair becomes a FusionCandidate scored by tangent alignment minus + a small penalty for gap length and a curvature-mismatch penalty. + + Tangents and curvatures are estimated by PCA over the last `lookback` + pixels at each endpoint (more robust to small skeletonization jogs + than a two-point difference). The curvature term penalizes matches + between segments of very different local curvature — e.g. a straight + line meeting a tight curve. + + score = tangent_score + - gap_penalty * path_length + - curvature_penalty * |curvature_a - curvature_b| + """ + candidates: List[FusionCandidate] = [] + endpoints: List[Tuple[int, int, NDArray]] = [] + for i, seg in enumerate(segments): + if len(seg) < 2: + continue + endpoints.append((i, 0, seg[0])) + endpoints.append((i, 1, seg[-1])) + + for ai in range(len(endpoints)): + seg_a, end_a, pt_a = endpoints[ai] + for bi in range(ai + 1, len(endpoints)): + seg_b, end_b, pt_b = endpoints[bi] + if seg_a == seg_b: + continue + d_eucl = float(np.linalg.norm(pt_a - pt_b)) + if d_eucl > 2 * max_path_length: + continue + ya, xa = int(round(pt_a[1])), int(round(pt_a[0])) + yb, xb = int(round(pt_b[1])), int(round(pt_b[0])) + path_yx = _shortest_path_through_binary( + binary, (ya, xa), (yb, xb), max_path_length + ) + if path_yx is None: + continue + path_xy = np.array([(x, y) for y, x in path_yx], dtype=float) + path_length = float(len(path_yx)) + + tan_a, curv_a = _local_tangent_and_curvature( + segments[seg_a], at_end=(end_a == 1), lookback=lookback + ) + tan_b, curv_b = _local_tangent_and_curvature( + segments[seg_b], at_end=(end_b == 1), lookback=lookback + ) + tangent_score = float(-np.dot(tan_a, tan_b)) + curvature_diff = abs(curv_a - curv_b) + score = ( + tangent_score + - gap_penalty * path_length + - curvature_penalty * curvature_diff + ) + candidates.append( + FusionCandidate( + seg_a=seg_a, + end_a=end_a, + seg_b=seg_b, + end_b=end_b, + connecting_path=path_xy, + path_length=path_length, + tangent_score=tangent_score, + curvature_a=curv_a, + curvature_b=curv_b, + score=score, + ) + ) + return candidates + + +def fuse_segments( + segments: List[NDArray[np.float64]], + binary: NDArray[np.bool_], + config: FuseConfig, +) -> Tuple[List[NDArray[np.float64]], List[FusionCandidate], List[FusionCandidate]]: + """Fuse segments by greedy maximum-weight matching of endpoints. + + Returns: + (fused_segments, accepted, all_candidates) + fused_segments: list of (M, 2) polylines after fusion + accepted: candidates used in the final matching + all_candidates: every candidate considered, accepted or not + """ + all_cands = find_fusion_candidates( + segments, + binary, + max_path_length=config["max_path_length"], + lookback=config["lookback"], + gap_penalty=config["gap_penalty"], + curvature_penalty=config["curvature_penalty"], + ) + cands = [c for c in all_cands if c.tangent_score >= config["min_tangent_score"]] + cands.sort(key=lambda c: c.score, reverse=True) + + n = len(segments) + used: set = set() + accepted: List[FusionCandidate] = [] + # link[seg][end] = (other_seg, other_end, connecting_path) or None + link: List[List[Optional[Tuple[int, int, NDArray]]]] = [ + [None, None] for _ in range(n) + ] + + for c in cands: + ka = (c.seg_a, c.end_a) + kb = (c.seg_b, c.end_b) + if ka in used or kb in used: + continue + used.add(ka) + used.add(kb) + link[c.seg_a][c.end_a] = (c.seg_b, c.end_b, c.connecting_path) + link[c.seg_b][c.end_b] = (c.seg_a, c.end_a, c.connecting_path[::-1]) + accepted.append(c) + + visited = [False] * n + fused: List[NDArray[np.float64]] = [] + + def walk(start_seg: int, start_end: int) -> NDArray[np.float64]: + """Walk the chain starting at segments[start_seg] from endpoint + `start_end`, emitting a single concatenated polyline.""" + chain: List[NDArray] = [] + cur_seg, cur_end = start_seg, start_end + while not visited[cur_seg]: + visited[cur_seg] = True + poly = segments[cur_seg] + pts = poly if cur_end == 0 else poly[::-1] + if chain and np.allclose(chain[-1], pts[0], atol=0.5): + chain.extend(pts[1:]) + else: + chain.extend(pts) + exit_end = 1 - cur_end + link_info = link[cur_seg][exit_end] + if link_info is None: + break + next_seg, next_end, conn_path = link_info + if visited[next_seg]: + break # cycle + # Insert the *interior* of the connecting path; the two endpoints + # are already in `chain` (last) and will be the first pixel of + # the next segment. + if len(conn_path) > 2: + interior = conn_path[1:-1] + chain.extend(interior) + cur_seg, cur_end = next_seg, next_end + return np.array(chain, dtype=float) + + # Walk from free (unlinked) endpoints first. + for i in range(n): + if visited[i]: + continue + start_end = None + if link[i][0] is None: + start_end = 0 + elif link[i][1] is None: + start_end = 1 + if start_end is not None: + fused.append(walk(i, start_end)) + # Whatever remains is part of a fully-linked cycle. + for i in range(n): + if visited[i]: + continue + fused.append(walk(i, 0)) + + return fused, accepted, all_cands + + +# ============================================================================ +# Post-repair fusion (SECOND PASS) +# ============================================================================ + + +class PostRepairFuseConfig(TypedDict): + junction_tol: float + """Max distance between endpoints (from different polylines) to be + considered as meeting at the same point. Should match (or be a + touch larger than) the value used in repair.""" + + tangent_skip: int + """Number of points to skip from each endpoint before sampling for + tangent estimation. Skipping past repair's bridge points avoids a + PCA tangent biased toward the OLD junction location. ~5 is + reasonable when ``repair`` ran with ``interp_max_spacing=1.0``.""" + + tangent_sample: int + """Number of points to PCA-fit for tangent estimation, starting + ``tangent_skip`` points in from the endpoint. ~10 is reasonable.""" + + min_tangent_score: float + """Threshold on ``-dot(tan_a, tan_b)`` below which a pair is not + considered a fusion candidate at all. 0.6-0.7 is reasonable; the + first pass accepts down to about 0.6 in typical configs and the + second pass can afford to be a touch stricter since tangents are + cleaner.""" + + curvature_penalty: float + """Penalty per unit local-curvature difference between the two + endpoints being joined. Discourages joining a tight curve to a + straight section. Keep small (~1.0); curvature estimates from a + short window are noisy and shouldn't dominate the score.""" + + +def fuse_post_repair( + polylines: List[NDArray[np.float64]], + config: PostRepairFuseConfig, +) -> Tuple[List[NDArray[np.float64]], List[FusionCandidate]]: + """Second-pass fusion for polylines that share endpoint coordinates + after junction repair (and optionally chromosome resolution). + + Unlike ``fuse_segments``, this doesn't do A*-through-binary path + finding. Post-repair, polylines that should be one stroke already + share an endpoint coordinate (the LS-cleaned junction point, or the + chromosome midpoint), so the only question is which pairs at each + shared location are the smoothest continuation. + + Tangents at endpoints are estimated by ``stable_endpoint_tangent_ + and_curvature``, which samples past repair's bridge points. Without + that, the PCA tangent would be biased by the bridge's near-colinear + points and could point at the OLD junction location instead of in + the polyline's true direction. + + Matching is greedy max-weight, same as the first pass. Walking the + chain is simpler because there are no connecting paths to splice + in: the two ends of every accepted link are already at the same + coordinate, so we just dedupe the seam. + + Returns: + ``(joined_polylines, accepted)``: + + - ``joined_polylines``: the new polyline list with accepted + pairs concatenated. + + - ``accepted``: the candidates that were used. Their + ``connecting_path`` arrays are empty since there's no path + to bridge; ``path_length`` is 0. + """ + polylines = [np.asarray(p, dtype=float) for p in polylines] + n = len(polylines) + if n == 0: + return [], [] + + meta, coords = gather_endpoints(polylines) + if len(coords) == 0: + return list(polylines), [] + point_poly_idx = [m[0] for m in meta] + raw_clusters = spatial_clusters(coords, config["junction_tol"], point_poly_idx) + + # Build candidates. + candidates: List[FusionCandidate] = [] + skip = config["tangent_skip"] + sample = config["tangent_sample"] + min_score = config["min_tangent_score"] + curv_pen = config["curvature_penalty"] + for cl in raw_clusters: + members = [meta[i] for i in cl] + if len({pi for pi, _ in members}) < 2: + continue + # Cache tangents within the cluster to avoid recomputing per pair. + member_tans: List[Tuple[int, int, NDArray, float]] = [] + for pi, end in members: + tan, curv = stable_endpoint_tangent_and_curvature( + polylines[pi], end, skip, sample + ) + member_tans.append((pi, end, tan, curv)) + m = len(member_tans) + for i in range(m): + pi_a, end_a, tan_a, curv_a = member_tans[i] + for j in range(i + 1, m): + pi_b, end_b, tan_b, curv_b = member_tans[j] + if pi_a == pi_b: + continue + tan_score = float(-np.dot(tan_a, tan_b)) + if tan_score < min_score: + continue + score = tan_score - curv_pen * abs(curv_a - curv_b) + candidates.append( + FusionCandidate( + seg_a=pi_a, + end_a=end_a, + seg_b=pi_b, + end_b=end_b, + connecting_path=np.empty((0, 2), dtype=float), + path_length=0.0, + tangent_score=tan_score, + curvature_a=curv_a, + curvature_b=curv_b, + score=score, + ) + ) + + # Greedy max-weight matching. + candidates.sort(key=lambda c: c.score, reverse=True) + used: Set[Tuple[int, int]] = set() + accepted: List[FusionCandidate] = [] + link: List[List[Optional[Tuple[int, int]]]] = [[None, None] for _ in range(n)] + for c in candidates: + ka = (c.seg_a, c.end_a) + kb = (c.seg_b, c.end_b) + if ka in used or kb in used: + continue + used.add(ka) + used.add(kb) + link[c.seg_a][c.end_a] = (c.seg_b, c.end_b) + link[c.seg_b][c.end_b] = (c.seg_a, c.end_a) + accepted.append(c) + + # Walk chains and concatenate. + visited = [False] * n + fused: List[NDArray[np.float64]] = [] + + def walk(start_seg: int, start_end: int) -> NDArray[np.float64]: + """Walk the chain starting at polylines[start_seg] from endpoint + side `start_end`. The seam between successive polylines is at + the same coordinate by construction, so the duplicate point at + the join is dropped by an ``allclose`` check.""" + chain: List[NDArray[np.float64]] = [] + cur_seg, cur_end = start_seg, start_end + while not visited[cur_seg]: + visited[cur_seg] = True + poly = polylines[cur_seg] + pts = poly if cur_end == 0 else poly[::-1] + if chain and np.allclose(chain[-1], pts[0], atol=0.5): + chain.extend(pts[1:]) + else: + chain.extend(pts) + exit_end = 1 - cur_end + link_info = link[cur_seg][exit_end] + if link_info is None: + break + next_seg, next_end = link_info + if visited[next_seg]: + break # cycle + cur_seg, cur_end = next_seg, next_end + return np.array(chain, dtype=float) + + # Walk from free (unlinked) endpoints first. + for i in range(n): + if visited[i]: + continue + start_end: Optional[int] = None + if link[i][0] is None: + start_end = 0 + elif link[i][1] is None: + start_end = 1 + if start_end is not None: + fused.append(walk(i, start_end)) + # Whatever remains is in a fully-linked cycle. + for i in range(n): + if visited[i]: + continue + fused.append(walk(i, 0)) + + return fused, accepted diff --git a/arc_line_vectorization_suede/segment/repair.py b/arc_line_vectorization_suede/segment/repair.py new file mode 100644 index 0000000..2abd7f9 --- /dev/null +++ b/arc_line_vectorization_suede/segment/repair.py @@ -0,0 +1,558 @@ +"""Local junction repair for skeleton-derived polylines. + +After ``fuse_segments``, polylines still pass through (or terminate at) +skeleton-level junctions. The handful of pixels nearest a Y, T, or X +junction is typically distorted: skeletonization commits each pixel to +exactly one stroke, so at a crossing it has to make local compromises +that look like small jogs, multi-pixel-thick "horizontal blobs", or +kinks in the polyline path. These don't match the actual stroke +geometry on either side of the junction, and they later confuse the +line/arc fitter. + +The fix: trust each segment's behaviour OUTSIDE the junction (its +tangent in the stable region a few pixels away), not the distorted +local geometry. For every junction, solve for the single 2D point +that all stable approaches meet at (a 2x2 least-squares problem), then +replace each polyline's junction-affected indices with that single +point. Polylines that pass through a junction now do so without the +jog; polylines that terminate at one now end at the same shared point +as their neighbours. + +Junction validity and region extent are decided GEOMETRICALLY, from +the polyline's smoothed second-derivative magnitude |r''|: + + (a) A cluster of close points across polylines is only treated as a + real junction if at least one polyline has a |r''| spike in a + padded window around the cluster. A spike means the polyline + bends sharply somewhere near the cluster -- exactly what the + skeletonizer's local "decide where to put pixels" compromises + produce at a real X/Y/T crossing. Two smooth polylines that just + happen to run within tol of each other (parallel coincidence, + tangent kiss) have no such spike and are correctly skipped. + Searching a few pixels OUTSIDE the strict cluster matters for + terminating-polyline junctions: there the cluster captures only + the polyline's first or last few points, and the actual kink + sits 1-5 pixels in from the endpoint. + + (b) For each polyline in a real junction, the region replaced by + the junction point is grown outward from the strict spatial + cluster while |r''| stays elevated above the polyline's local + baseline. Without this extension, skeleton noise that spills a + few pixels past the close-points zone survives into the output + as a small "shelf" or jog next to the cleaned junction. + + (c) Tangents for the LS solve are taken from samples just OUTSIDE + the extended region, so the meeting point is constrained by + genuinely smooth polyline geometry rather than by pixels that + are still inside the noisy region. + +Algorithm: + +1. Build a junction graph over polyline points. Two points are linked + if they lie within ``junction_tol`` of each other AND belong to + different polylines. Connected components are junction clusters. + +2. For each cluster spanning >= 2 polylines: + - Reject if any polyline contributes more than + ``max_junction_region_length`` consecutive points to the + cluster (a sanity net against pathologically long runs -- + typically closed-loop polylines whose start and end coincide). + - Reject unless at least one polyline shows a |r''| spike in a + padded window around its cluster indices (the kink test). + - Extend each polyline's region by walking outward through + elevated-|r''| pixels. + - Build approaches from samples beyond each extended region. + - Reject if the approach tangents are all nearly parallel + (tangent kiss). + - Solve the 2x2 LS problem for the junction point. + +3. For each polyline involved in the junction, splice out the + junction-affected indices (extended) and replace them with the + single junction point. Adjacent emitted segments are bridged with + 1-pixel-spaced interpolated points so the polyline stays visually + continuous and downstream fitters don't see a gap. + +References (close in spirit, none are exactly this algorithm): + - Hilaire & Tombre 2006, "Robust and Accurate Vectorization of Line + Drawings". + - Favreau, Lafarge, Bousseau 2016, "Fidelity vs. Simplicity". + - Bessmeltsev & Solomon 2019, "Vectorization of Line Drawings via + Polyvector Fields". + - Bao & Fu 2023, "Joint Curve Network Optimization for Drawing + Vectorization" (does this jointly with global primitive alignment). +""" + +from __future__ import annotations +from typing import List, Tuple, TypedDict + +import numpy as np +from numpy.typing import NDArray + +from ._helpers import gather_points, spatial_clusters + + +def _build_approaches(polyline, junction_indices, stable_skip, stable_sample): + n = len(polyline) + if not junction_indices: + return [] + region_start = min(junction_indices) + region_end = max(junction_indices) + arms = [] + if region_start > 0: + inner_idx = region_start - 1 + sample_near = inner_idx - stable_skip + sample_far = inner_idx - stable_skip - stable_sample + 1 + sample_near = max(0, sample_near) + sample_far = max(0, sample_far) + if sample_near < sample_far: + sample_near, sample_far = sample_far, sample_near + if sample_near - sample_far < 1: + sample_far = 0 + sample_near = inner_idx + if sample_near > sample_far: + samples = polyline[sample_far : sample_near + 1] + anchor = samples.mean(axis=0) + v = samples[0] - samples[-1] + norm = float(np.linalg.norm(v)) + if norm > 1e-9: + arms.append((inner_idx, -1, anchor, v / norm)) + if region_end < n - 1: + inner_idx = region_end + 1 + sample_near = inner_idx + stable_skip + sample_far = inner_idx + stable_skip + stable_sample - 1 + sample_near = min(n - 1, sample_near) + sample_far = min(n - 1, sample_far) + if sample_far < sample_near: + sample_near, sample_far = sample_far, sample_near + if sample_far - sample_near < 1: + sample_near = inner_idx + sample_far = n - 1 + if sample_far > sample_near: + samples = polyline[sample_near : sample_far + 1] + anchor = samples.mean(axis=0) + v = samples[-1] - samples[0] + norm = float(np.linalg.norm(v)) + if norm > 1e-9: + arms.append((inner_idx, +1, anchor, v / norm)) + return arms + + +def _tangent_spread_deg(tangents): + if len(tangents) < 2: + return 180.0 + angles = np.mod(np.array([np.arctan2(t[1], t[0]) for t in tangents]), np.pi) + sorted_a = np.sort(angles) + gaps = np.append(np.diff(sorted_a), np.pi - sorted_a[-1] + sorted_a[0]) + max_gap = float(np.max(gaps)) + return float(np.degrees(np.pi - max_gap)) + + +def _solve_junction_point(anchors, outward_tangents, fallback=None): + if not anchors: + if fallback is None: + raise ValueError("No approaches and no fallback") + return fallback.copy() + if len(anchors) == 1: + return anchors[0].copy() + A = np.zeros((2, 2)) + b = np.zeros(2) + for a, t in zip(anchors, outward_tangents): + M = np.eye(2) - np.outer(t, t) + A += M + b += M @ a + try: + return np.linalg.solve(A, b) + except np.linalg.LinAlgError: + return np.mean(anchors, axis=0) if fallback is None else fallback.copy() + + +def _merge_cascade_repairs( + repairs: List[Tuple[int, int, NDArray]], + cascade_gap: int = 3, + max_jp_distance: float = 3.0, +) -> List[Tuple[int, int, NDArray]]: + """Merge repairs whose junction regions are within ``cascade_gap`` + pixels on the same polyline AND whose junction points are within + ``max_jp_distance`` pixels of each other. + + The original case this targets is one physical junction that + spatial clustering split into two adjacent clusters with very + slightly different centroids (e.g., P2 meets P4 at one centroid + and P5 at a slightly different centroid two pixels away). Applied + independently, the two repairs would chew through the polyline in + series and remove a much larger chunk than either junction + actually warrants; merging them keeps the chunk small. + + The JP-distance guard prevents merging when a single polyline + threads through two PHYSICALLY DISTINCT junctions whose extended + regions happen to overlap. Without the guard, both junction + points get averaged into a single point in between, and the + polyline's path no longer reaches either junction's true meeting + point -- a clearly visible disconnect. + """ + if len(repairs) <= 1: + return list(repairs) + repairs = sorted(repairs, key=lambda r: r[0]) + merged: List[Tuple[int, int, NDArray]] = [repairs[0]] + for js, je, jp in repairs[1:]: + last_js, last_je, last_jp = merged[-1] + close_in_index = js - last_je <= cascade_gap + close_in_space = ( + float(np.linalg.norm(np.asarray(jp) - np.asarray(last_jp))) + <= max_jp_distance + ) + if close_in_index and close_in_space: + new_jp = ( + np.asarray(last_jp, dtype=float) + np.asarray(jp, dtype=float) + ) / 2.0 + merged[-1] = (last_js, max(last_je, je), new_jp) + else: + merged.append((js, je, jp)) + return merged + + +def _polyline_d2_mag(polyline, smoothing_window=3): + """Smoothed |2nd derivative| magnitude per polyline point. + + Used to detect skeleton-level kinks. For a polyline parameterized + at ~1px spacing, |r''| ~ 1/R for a circular section (radius R), so + a smooth curve has a stable low baseline, while a sharp local + distortion (e.g. the multi-pixel "horizontal blob" left by skeleton- + izing an X-crossing) spikes well above that baseline. + """ + n = len(polyline) + if n < 3: + return np.zeros(n) + dx = np.gradient(polyline[:, 0]) + dy = np.gradient(polyline[:, 1]) + ddx = np.gradient(dx) + ddy = np.gradient(dy) + mag = np.sqrt(ddx * ddx + ddy * ddy) + if smoothing_window > 1: + k = min(smoothing_window, n) + kernel = np.ones(k, dtype=float) / float(k) + mag = np.convolve(mag, kernel, mode="same") + return mag + + +def _kink_present(d2_mag, rs, re, context_window, spike_ratio, search_pad=4): + """Test if |r''| spikes near the cluster vs the baseline further out. + + Looks for the spike in ``[rs - search_pad, re + search_pad]`` rather + than just the strict cluster: at terminating-polyline junctions, the + cluster captures only the few pixels nearest the meeting point, but + the kink itself (the bend where the polyline turns to fit the + skeleton's junction blob) usually sits 1-5 pixels INSIDE from the + endpoint. Searching a slightly padded window catches that. Baseline + is measured from the pixels just OUTSIDE this padded window. + """ + n = len(d2_mag) + if n == 0 or rs > re: + return False + rs_c = max(0, rs) + re_c = min(n - 1, re) + search_start = max(0, rs_c - search_pad) + search_end = min(n - 1, re_c + search_pad) + in_search = d2_mag[search_start : search_end + 1] + if len(in_search) == 0: + return False + base_start = max(0, search_start - context_window) + base_end = min(n, search_end + 1 + context_window) + outside_pieces = [] + if base_start < search_start: + outside_pieces.append(d2_mag[base_start:search_start]) + if search_end + 1 < base_end: + outside_pieces.append(d2_mag[search_end + 1 : base_end]) + if not outside_pieces: + return True + outside = np.concatenate(outside_pieces) + if len(outside) < 2: + return True + baseline = max(float(np.median(outside)), 1e-3) + spike = float(np.max(in_search)) + return spike >= spike_ratio * baseline + + +def _extend_region_by_d2(d2_mag, rs, re, max_extend, baseline_factor=1.5): + """Extend [rs, re] outward as long as |r''| stays elevated. + + Skeleton-level distortion often extends a few pixels BEYOND the + strict spatial cluster: the polyline is no longer within tolerance + of the other polyline out there, but its local geometry is still + bent. Without this extension, those out-of-cluster noisy pixels + survive into the output as a small "shelf" or jog next to the + cleaned junction. We walk outward from the strict region on each + side and grow it while |r''| exceeds ``baseline_factor`` x the + baseline (median of values further outside). + """ + n = len(d2_mag) + if n == 0 or rs > re: + return rs, re + rs_c = max(0, rs) + re_c = min(n - 1, re) + # Baseline measured FAR from both the strict region and the + # potential extension zone, so the extension can't bias it. + far_left_end = max(0, rs_c - max_extend) + far_right_start = min(n, re_c + 1 + max_extend) + context = 8 + far_left_start = max(0, far_left_end - context) + far_right_end = min(n, far_right_start + context) + pieces = [] + if far_left_start < far_left_end: + pieces.append(d2_mag[far_left_start:far_left_end]) + if far_right_start < far_right_end: + pieces.append(d2_mag[far_right_start:far_right_end]) + if not pieces: + return rs_c, re_c + outside = np.concatenate(pieces) + if len(outside) < 2: + return rs_c, re_c + baseline = max(float(np.median(outside)), 1e-3) + threshold = baseline * baseline_factor + rs_ext = rs_c + for off in range(1, max_extend + 1): + idx = rs_c - off + if idx < 0: + break + if d2_mag[idx] > threshold: + rs_ext = idx + else: + break + re_ext = re_c + for off in range(1, max_extend + 1): + idx = re_c + off + if idx >= n: + break + if d2_mag[idx] > threshold: + re_ext = idx + else: + break + return rs_ext, re_ext + + +def _interpolate(start, end, max_spacing): + start = np.asarray(start, dtype=float) + end = np.asarray(end, dtype=float) + d = float(np.linalg.norm(end - start)) + if d <= max_spacing or max_spacing <= 0: + return np.empty((0, 2), dtype=float) + n_segments = int(np.ceil(d / max_spacing)) + t = np.linspace(0.0, 1.0, n_segments + 1)[1:-1] + return start[None, :] + t[:, None] * (end - start)[None, :] + + +def _apply_repairs(polyline, repairs, interp_max_spacing=1.0): + if not repairs: + return polyline.copy() + repairs = sorted(repairs, key=lambda r: r[0]) + pieces: List[NDArray] = [] + cursor = 0 + last_emitted: NDArray | None = None + + def push(arr): + nonlocal last_emitted + if len(arr) == 0: + return + if last_emitted is not None: + bridge = _interpolate(last_emitted, arr[0], interp_max_spacing) + if len(bridge) > 0: + pieces.append(bridge) + pieces.append(arr) + last_emitted = arr[-1] + + for js, je, jp in repairs: + if js > cursor: + push(polyline[cursor:js]) + push(np.asarray(jp, dtype=float).reshape(1, 2)) + cursor = je + 1 + if cursor < len(polyline): + push(polyline[cursor:]) + if not pieces: + return np.empty((0, 2), dtype=polyline.dtype) + return np.concatenate(pieces, axis=0).astype(polyline.dtype, copy=False) + + +class RepairConfig(TypedDict): + + junction_tol: float + """ + Max distance for two polylines' points to be considered part of + the same junction. ~2.5 px works for typical 8-connected skeletons. + """ + + stable_skip: int + """ + Pixels to skip BEYOND the extended junction region before sampling + for the tangent. Buffers against any residual distortion just + outside the noisy zone. ~2 is usually right. + """ + + stable_sample: int + """ + Number of pixels to sample for tangent estimation once past the + skip buffer. Larger = more robust but smooths away genuine + curvature in short approaches. + """ + + max_junction_region_length: int + """ + Reject junctions where any polyline contributes more than this + many consecutive points to the strict spatial cluster. Now a + sanity net rather than the primary discriminator: the kink test + below is what distinguishes real crossings from parallel + coincidences. Useful default ~20-30 to catch closed-loop polylines + whose start/end coincidence makes their cluster span the entire + polyline. + """ + + min_tangent_spread_deg: float + """ + Reject junctions where all approach tangents point in nearly the + same direction (mod 180 deg). Backstop against tangent kisses + that slip past the kink test. Default 15 deg accepts oblique X + crossings (down to ~15 deg between strokes) while rejecting + tangent kisses and parallel coincidences. + """ + + interp_max_spacing: float + """ + After replacing a junction region with the junction point, the + polyline is filled with interpolated points so that consecutive + points are at most this far apart. Default 1.0 (one pixel) matches + the skeleton spacing, so the polyline remains visually continuous + and the downstream fitter doesn't see a gap. + """ + + min_output_polyline_length: int + """ + After repair, drop polylines shorter than this (rare; happens if a + polyline was almost entirely consumed by junctions). + """ + + min_curvature_spike_ratio: float + """ + A cluster is treated as a real junction only if some polyline has + ``max(|r''|)`` in a padded window around the cluster that is at + least this many times the median |r''| baseline measured just + OUTSIDE that window. Default 2.0 cleanly separates skeleton-level + kinks (ratios of ~5-20x) from parallel coincidences (~1x). + """ + + curvature_context_window: int + """ + Half-width (in pixels) of the |r''| context window used both for + the kink test (where it sets the baseline measurement zone) and + for the region-extension walk (where it caps how far the + junction-affected region can grow on each side). Default 8. + """ + + +def repair_junctions(polylines, config: RepairConfig): + if not polylines: + return [] + polylines = [np.asarray(p, dtype=float) for p in polylines] + meta, coords = gather_points(polylines) + if len(coords) == 0: + return list(polylines) + point_poly_idx = [m[0] for m in meta] + raw_clusters = spatial_clusters(coords, config["junction_tol"], point_poly_idx) + clusters = [] + for c in raw_clusters: + polys = {meta[i][0] for i in c} + if len(polys) >= 2: + clusters.append(c) + per_poly_repairs = {pi: [] for pi in range(len(polylines))} + + # Pre-compute |r''| profile per polyline for the kink-presence test. + d2_profiles = [_polyline_d2_mag(p, smoothing_window=3) for p in polylines] + + for cluster in clusters: + poly_to_indices = {} + for global_idx in cluster: + pi, ii = meta[global_idx] + poly_to_indices.setdefault(pi, []).append(ii) + too_long = False + for indices in poly_to_indices.values(): + if max(indices) - min(indices) + 1 > config["max_junction_region_length"]: + too_long = True + break + if too_long: + continue + # KINK CHECK: distinguishes real crossings (sharp local turn in + # at least one polyline) from "parallel coincidences" where two + # smooth polylines happen to run within tol of each other. + # Without this, raising max_junction_region_length to cover + # legitimate X-crossings with multi-pixel-thick skeleton blobs + # would also accept parallel runs. + n_kinky = 0 + for pi, indices in poly_to_indices.items(): + rs_i, re_i = min(indices), max(indices) + if _kink_present( + d2_profiles[pi], + rs_i, + re_i, + context_window=config["curvature_context_window"], + spike_ratio=config["min_curvature_spike_ratio"], + ): + n_kinky += 1 + if n_kinky < 1: + continue + # Extend each polyline's junction-affected region by walking + # outward as long as |r''| stays elevated. The skeleton-level + # distortion at an X-crossing often spills a few pixels past + # the strict spatial cluster, and leaving them in the output + # produces a visible "shelf" or jog next to the cleaned point. + extended_poly_indices = {} + for pi, indices in poly_to_indices.items(): + rs_i, re_i = min(indices), max(indices) + rs_ext, re_ext = _extend_region_by_d2( + d2_profiles[pi], + rs_i, + re_i, + max_extend=config["curvature_context_window"], + baseline_factor=1.5, + ) + extended_poly_indices[pi] = (rs_ext, re_ext) + all_anchors = [] + all_tangents = [] + poly_to_region = {} + for pi, indices in poly_to_indices.items(): + rs_ext, re_ext = extended_poly_indices[pi] + # Build approaches from samples beyond the EXTENDED region, + # so the tangent estimate is taken from genuinely smooth + # polyline (not from pixels that are still inside the noise). + ext_indices = list(range(rs_ext, re_ext + 1)) + arms = _build_approaches( + polylines[pi], + ext_indices, + config["stable_skip"], + config["stable_sample"], + ) + if not arms: + continue + for _inner, _dir, anchor, tangent in arms: + all_anchors.append(anchor) + all_tangents.append(tangent) + poly_to_region[pi] = (rs_ext, re_ext) + if len(poly_to_region) < 2 or len(all_anchors) < 2: + continue + spread = _tangent_spread_deg(all_tangents) + if spread < config["min_tangent_spread_deg"]: + continue + cluster_centroid = coords[cluster].mean(axis=0) + jp = _solve_junction_point(all_anchors, all_tangents, fallback=cluster_centroid) + max_dev = 5.0 * config["junction_tol"] + if np.linalg.norm(jp - cluster_centroid) > max_dev: + jp = cluster_centroid + for pi, (rs, re) in poly_to_region.items(): + per_poly_repairs[pi].append((rs, re, jp)) + repaired = [] + for pi, poly in enumerate(polylines): + raw = per_poly_repairs[pi] + merged = _merge_cascade_repairs(raw, cascade_gap=3) if len(raw) >= 2 else raw + new_poly = _apply_repairs( + poly, merged, interp_max_spacing=config["interp_max_spacing"] + ) + if len(new_poly) >= config["min_output_polyline_length"]: + repaired.append(new_poly) + return repaired diff --git a/arc_line_vectorization_suede/segment/trace.py b/arc_line_vectorization_suede/segment/trace.py new file mode 100644 index 0000000..ba6c371 --- /dev/null +++ b/arc_line_vectorization_suede/segment/trace.py @@ -0,0 +1,156 @@ +import numpy as np +from scipy.ndimage import label as cc_label +from typing import List, cast +from numpy.typing import NDArray + +NEIGHBOR_OFFSETS = [ + (-1, -1), + (-1, 0), + (-1, 1), + (0, -1), + (0, 1), + (1, -1), + (1, 0), + (1, 1), +] + + +def _skeleton_neighbors(skel: np.ndarray, y: int, x: int): + H, W = skel.shape + for dy, dx in NEIGHBOR_OFFSETS: + ny, nx = y + dy, x + dx + if 0 <= ny < H and 0 <= nx < W and skel[ny, nx]: + yield ny, nx + + +def _compute_degree(skel: np.ndarray) -> np.ndarray: + """Per-pixel count of 8-connected ON neighbours.""" + deg = np.zeros_like(skel, dtype=np.int8) + ys, xs = np.where(skel) + for y, x in zip(ys, xs): + deg[y, x] = sum(1 for _ in _skeleton_neighbors(skel, y, x)) + return deg + + +def _label_node_components(skel: np.ndarray, deg: np.ndarray): + """Assign a unique non-zero id to every "node pixel" in the skeleton. + + A node pixel is either: + - an endpoint (degree == 1), or + - a member of a super-junction: a connected component (8-connected) of + pixels with degree >= 3. All pixels in one component share the same id. + + Returns (node_id, n_super_junctions). node_id is an int array with the same + shape as skel; 0 means "not a node" (degree-2 interior). + """ + + structure = np.ones((3, 3), dtype=bool) + junction_mask = deg >= 3 + sj_labels, n_sj = cast( + tuple[NDArray[np.int32], int], + cc_label(junction_mask, structure=structure), + ) + + node_id = np.zeros_like(skel, dtype=np.int32) + node_id[junction_mask] = sj_labels[junction_mask] # ids 1..n_sj + + endpoint_ys, endpoint_xs = np.where(deg == 1) + next_id = n_sj + 1 + for y, x in zip(endpoint_ys, endpoint_xs): + node_id[y, x] = next_id + next_id += 1 + + return node_id, n_sj + + +def trace_skeleton(skel: np.ndarray) -> List[NDArray[np.float64]]: + """Trace a skeleton into a list of polylines. + + Returns a list of (N, 2) float arrays of (`x, y) pixel coordinates. + + Endpoints and junction-clusters anchor the polylines. Multiple adjacent + junction pixels are treated as a single node (super-junction), so the + walk does not spawn spurious 2-pixel polylines inside thick junctions. + + Closed loops with no endpoints/junctions are broken at an arbitrary pixel. + """ + skel = skel.astype(bool) + deg = _compute_degree(skel) + node_id, _ = _label_node_components(skel, deg) + + polylines: List[NDArray[np.float64]] = [] + visited_starts = set() # (start_pixel, first_step_pixel) + + def trace(start_pixel, start_id, first_step): + """Walk from start_pixel via first_step until reaching a different + node. Returns the path or None if first_step is in the same node.""" + if node_id[first_step] == start_id: + return None # hopping inside the same super-junction + path = [start_pixel, first_step] + prev = start_pixel + cur = first_step + while node_id[cur] == 0: + next_p = None + for nny, nnx in _skeleton_neighbors(skel, cur[0], cur[1]): + cand = (nny, nnx) + if cand == prev: + continue + # Prefer non-node neighbours; otherwise accept the first node. + if node_id[cand] == 0: + next_p = cand + break + if next_p is None: + next_p = cand + if next_p is None: + break + prev = cur + cur = next_p + path.append(cur) + return path + + # All node pixels (members of super-junctions + endpoints) + node_ys, node_xs = np.where(node_id != 0) + for sy, sx in zip(node_ys, node_xs): + start_pixel = (sy, sx) + start_id = int(node_id[sy, sx]) + for ny, nx in _skeleton_neighbors(skel, sy, sx): + first_step = (ny, nx) + key = (start_pixel, first_step) + if key in visited_starts: + continue + path = trace(start_pixel, start_id, first_step) + if path is None: + continue + visited_starts.add(key) + # Also mark the reverse direction so we don't re-trace the same edge + if len(path) >= 2: + visited_starts.add((path[-1], path[-2])) + polylines.append(np.array([(x, y) for y, x in path], dtype=float)) + + # Closed loops: connected components of skel pixels that contain no node. + seen = np.zeros_like(skel, dtype=bool) + for poly in polylines: + for x, y in poly: + seen[int(round(y)), int(round(x))] = True + for y, x in zip(*np.where(skel & (node_id == 0) & ~seen)): + path = [(y, x)] + seen[y, x] = True + prev = None + cur = (y, x) + while True: + next_p = None + for ny, nx in _skeleton_neighbors(skel, cur[0], cur[1]): + if (ny, nx) != prev and not seen[ny, nx]: + next_p = (ny, nx) + break + if next_p is None: + break + seen[next_p] = True + path.append(next_p) + prev = cur + cur = next_p + if len(path) >= 3: + path.append(path[0]) # close the loop + polylines.append(np.array([(x, y) for y, x in path], dtype=float)) + + return polylines diff --git a/arc_line_vectorization_suede/skeletonize/__init__.py b/arc_line_vectorization_suede/skeletonize/__init__.py new file mode 100644 index 0000000..d6ee64c --- /dev/null +++ b/arc_line_vectorization_suede/skeletonize/__init__.py @@ -0,0 +1,137 @@ +import numpy as np +from numpy.typing import NDArray +from skimage import io +from skimage.color import rgb2gray, rgba2rgb +from skimage.morphology import skeletonize as _skeletonize + +from .cleanup import collapse_small_holes, CollapseConfig +from .crossings import ( + DetectConfig, + DetectResult, + detect_crossings, + resolve_crossings, +) + +from typing import Literal, TypedDict, NamedTuple, Union + + +class BinarizeConfig(TypedDict): + threshold: float # 0.0 to 1.0 + + +ImageSource = Union[str, np.ndarray] + + +def _to_grayscale_float(img: np.ndarray) -> np.ndarray: + """Normalize an arbitrary image array to a float grayscale in [0, 1].""" + if img.ndim == 3: + if img.shape[-1] == 4: + img = rgba2rgb(img) # also yields float in [0, 1] + if img.shape[-1] == 3: + return rgb2gray(img) # float in [0, 1] + raise ValueError( + f"Unsupported channel count {img.shape[-1]} for 3D image input; " + "expected 3 (RGB) or 4 (RGBA)." + ) + if img.ndim != 2: + raise ValueError( + f"Unsupported image ndim {img.ndim}; expected 2 (grayscale) or 3." + ) + if img.dtype == np.bool_: + return img.astype(np.float64) + if np.issubdtype(img.dtype, np.integer): + return img / np.float64(np.iinfo(img.dtype).max) + return img.astype(np.float64, copy=False) + + +def to_binary(source: ImageSource, config: BinarizeConfig) -> NDArray[np.bool_]: + if isinstance(source, np.ndarray): + img = _to_grayscale_float(source) + else: + img = io.imread(source, as_gray=True) + if img.dtype != np.float64 and img.dtype != np.float32: + img = img / 255.0 + return img < config["threshold"] + + +class SkeletonizeConfig(TypedDict): + method: Literal["lee", "zhang"] + + +def skeletonize( + mask: NDArray[np.bool_], config: SkeletonizeConfig +) -> NDArray[np.bool_]: + return _skeletonize(mask, method=config["method"]) + + +class Skeletonize: + """End-to-end skeletonization pipeline. + + Stages: + 1. binarize -> self.binary + 2. skeletonize -> self.skeletonized + 3. collapse_small_holes -> self.collapsed + (fixes thin double-pixel artifacts from Lee thinning) + 4. detect_crossings -> self.detection + (identifies ribbon-collapse regions in the binary; uses + the cleaned skeleton from stage 3 so arm endpoints land + on actual cleaned-skeleton pixels for downstream resolution) + 5. resolve_crossings -> self.uncrossed + (rewrites each detected ribbon collapse: erases the merged + skeleton segment and replaces it with two straight lines + between paired arm endpoints, restoring two non-intersecting + paths through the crossing) + + The final output for downstream consumption is `self.uncrossed`. + Earlier stages are kept on the instance so visualization / debug + can compare them. + """ + + class Config: + class Binarize(BinarizeConfig): + pass + + class Skeletonize(SkeletonizeConfig): + pass + + class Collapse(CollapseConfig): + pass + + class Detect(DetectConfig): + pass + + class Output(NamedTuple): + binary: NDArray[np.bool_] + skeletonized: NDArray[np.bool_] + collapsed: NDArray[np.bool_] + detection: DetectResult + uncrossed: NDArray[np.bool_] + + def __init__( + self, + source: ImageSource, + binarize_config: Config.Binarize, + skeletonize_config: Config.Skeletonize, + collapse_config: Config.Collapse, + detect_config: Config.Detect, + ): + self.binary = to_binary(source, binarize_config) + self.skeletonized = skeletonize(self.binary, skeletonize_config) + self.collapsed = collapse_small_holes(self.skeletonized, collapse_config) + # Detection uses the BINARY for its distance-transform analysis + # but is given the CLEANED skeleton so arm endpoints land on + # pixels that will survive the resolver's edits. + self.detection = detect_crossings( + self.binary, detect_config, skel=self.collapsed + ) + # Resolution rewrites the cleaned skeleton in place at every + # detected crossing. + self.uncrossed = resolve_crossings(self.collapsed, self.detection) + + self.output = self.Output( + binary=self.binary, + skeletonized=self.skeletonized, + collapsed=self.collapsed, + detection=self.detection, + uncrossed=self.uncrossed, + ) diff --git a/arc_line_vectorization_suede/skeletonize/cleanup.py b/arc_line_vectorization_suede/skeletonize/cleanup.py new file mode 100644 index 0000000..28a0289 --- /dev/null +++ b/arc_line_vectorization_suede/skeletonize/cleanup.py @@ -0,0 +1,162 @@ +"""Post-process a skeleton to collapse small enclosed regions. + +Some skeletonization implementations — notably `method="lee"` on certain +line shapes — leave behind tiny closed loops where the line failed to +collapse to 1-pixel width. Two common artifact shapes: + + 1. "Double diagonal": two parallel diagonal runs of pixels share corner- + only contact, sandwiching 1-pixel-wide enclosed regions between them. + 2. "Long thin bulge": a stretch where the line briefly thickens to 2 + pixels, producing a long but narrow enclosed region (high area, but + never more than a couple of pixels thick anywhere). + +This module provides a function that finds those enclosed regions and +fills them, then thins the result so the line ends up 1 pixel wide. + +A region is collapsed if EITHER criterion is met: + (a) its pixel area is below `max_hole_area` (catches tiny holes + regardless of shape), OR + (b) its maximum local thickness is at most `max_thin_thickness` + (catches long-and-thin artifacts that pass the area test). + +Local thickness is measured via the Euclidean distance transform: +`distance_transform_edt(region)` gives, for each interior pixel, the +distance to the nearest non-region pixel. The maximum of that field is +the radius of the largest disk that fits inside the region; we treat +that as half the local thickness. For a region that is k pixels thick +everywhere, the max distance-transform value is approximately k/2 (so +max_dt <= 1.0 catches up to ~2-pixel-thick regions). +""" + +from __future__ import annotations + +import numpy as np +from numpy.typing import NDArray +from scipy.ndimage import distance_transform_edt +from scipy.ndimage import label as cc_label +from skimage.morphology import skeletonize + +from typing import Literal, TypedDict + +# 4-connectivity: only orthogonal neighbours count as connected. Crucially, +# this is the right choice for hole detection on a skeleton — a pair of +# skeleton pixels that meet only at a corner are 8-connected to each other, +# but the background pixels they sandwich are 4-isolated from the outside. +_CROSS_4 = np.array([[0, 1, 0], [1, 1, 1], [0, 1, 0]], dtype=bool) + + +def _enclosed_region_masks(skel: NDArray[np.bool_]) -> list[NDArray[np.bool_]]: + """Return a list of boolean masks, one per enclosed background region. + + A region is "enclosed" when it is a 4-connected component of background + pixels that does not touch the image border. + """ + H, W = skel.shape + if H == 0 or W == 0: + return [] + background = ~skel + labeled, n = cc_label(background, structure=_CROSS_4) + if n == 0: + return [] + border_labels: set[int] = set() + border_labels.update(np.unique(labeled[0, :]).tolist()) + border_labels.update(np.unique(labeled[-1, :]).tolist()) + border_labels.update(np.unique(labeled[:, 0]).tolist()) + border_labels.update(np.unique(labeled[:, -1]).tolist()) + border_labels.discard(0) # 0 is the foreground; not a real component + return [labeled == lid for lid in range(1, n + 1) if lid not in border_labels] + + +def _max_thickness(mask: NDArray[np.bool_]) -> float: + """Return the maximum local thickness of a region (in pixel units). + + Uses the Euclidean distance transform: for each interior pixel, the + distance to the nearest non-region pixel; the maximum of that field + is half the local thickness at the deepest interior point. + """ + if not mask.any(): + return 0.0 + dt = distance_transform_edt(mask) + return float(dt.max()) + + +def find_enclosed_regions( + skel: NDArray[np.bool_], + max_area: int | None = None, + max_thickness: float | None = None, +) -> list[NDArray[np.bool_]]: + """Return masks of enclosed background regions matching the filters. + + A region passes if EITHER of the supplied filters matches: + - `max_area`: area of the region in pixels is <= max_area + - `max_thickness`: the region's maximum local thickness (via distance + transform) is <= max_thickness + + If both are None, returns every enclosed region. If both are given, + a region is included if either criterion matches (logical OR). + """ + skel = skel.astype(bool) + masks = _enclosed_region_masks(skel) + if max_area is None and max_thickness is None: + return masks + + out = [] + for m in masks: + passes = False + if max_area is not None and int(m.sum()) <= max_area: + passes = True + if not passes and max_thickness is not None: + if _max_thickness(m) <= max_thickness: + passes = True + if passes: + out.append(m) + return out + + +class CollapseConfig(TypedDict): + max_hole_area: int | None + max_thin_thickness: float | None + reskeletonize: bool + skeletonize_method: Literal["lee", "zhang"] + + +def collapse_small_holes( + skel: NDArray[np.bool_], + config: CollapseConfig, +) -> NDArray[np.bool_]: + """Fill collapsible enclosed background regions and (optionally) re-thin. + + A region is filled if EITHER of the following holds: + (a) its area in pixels is <= `max_hole_area`, OR + (b) its maximum local thickness is <= `max_thin_thickness`. + + Defaults catch the common Lee-method artifacts: + - `max_hole_area=4`: handles "double diagonal" 1-pixel holes and + small elbow artifacts. + - `max_thin_thickness=1.0`: handles long-and-thin bulges where the + line briefly doubles up; this catches anything up to ~2 pixels + thick everywhere, regardless of length. + + The skeletonizer's topology preservation means strokes that ran + through a filled region remain connected, so the line continues + correctly through the patch. + + Set either parameter to 0 / `None` to disable that criterion. + """ + skel = skel.astype(bool) + if skel.size == 0: + return skel.copy() + + masks = find_enclosed_regions( + skel, + max_area=config.get("max_hole_area"), + max_thickness=config.get("max_thin_thickness"), + ) + filled = skel.copy() + for mask in masks: + filled[mask] = True + + if config.get("reskeletonize") and masks: + filled = skeletonize(filled, method=config.get("skeletonize_method")) + + return filled diff --git a/arc_line_vectorization_suede/skeletonize/crossings.py b/arc_line_vectorization_suede/skeletonize/crossings.py new file mode 100644 index 0000000..7864244 --- /dev/null +++ b/arc_line_vectorization_suede/skeletonize/crossings.py @@ -0,0 +1,505 @@ +"""Detect ribbon-collapse artifacts in a binarized line drawing. + +Two strokes that briefly share ink — by truly crossing over, by being +tangent, or by aligning for a short segment — produce a characteristic +skeletonization artifact: the merged region of ink is wide enough that +the thinning algorithm is forced to draw the skeleton as a single line +through the middle, even though the original drawing has two distinct +strokes there. After the merged region ends, the skeleton splits back +into two lines via a Y-junction on each side. The whole pattern looks +like ``)―(`` in the skeleton: a single connecting segment between two +Y-junctions. The literature calls this a "ribbon collapse" (Bessmeltsev +et al., PolyVector Fields, SIGGRAPH 2019), a "stroke merge" (Favreau et +al., SIGGRAPH 2016), or informally a "chromosome". + +This module detects those artifacts so downstream code can resolve them +(pair the four arms, rewrite the skeleton as two distinct strokes). + +Why this approach instead of the obvious ones +--------------------------------------------- +Global stroke-half-width thresholding does not work when a drawing +contains both thin and thick strokes: the threshold is set by the +thicker strokes and the thinner-stroke ribbon collapses are too narrow +to qualify. The fix is a per-pixel LOCAL stroke half-width estimated +from nearby skeleton pixels. + +Counting how many non-fat ink components touch a fat region also does +not work: the rest of the drawing's ink is one globally connected +blob, so the count is always 1. The fix is to count skeleton-arm +exits via 8-connected components of the skeleton in a ring around the +fat region, since the skeleton's 1-pixel arms are naturally separated +by background. + +A degree-4 fat region is necessary but not sufficient: four strokes +meeting at a single point look identical at the topology level. Two +further checks distinguish actual ribbon collapses from point +junctions: + +- **Stroke pairing**: the four arm tangents must pair into two + anti-parallel pairs (two strokes continuing through). Wheel-spoke + junctions fail this because the arms radiate in 4 directions. +- **Merged segment length**: the skeleton trapped inside the fat + region must be a meaningful segment, not just a junction vertex. + This is the most direct test for the ``)―(`` pattern: a ribbon + collapse has a long single-line segment inside the fat region; a + 4-way junction has only a junction node and its immediate + neighbors there. + +Pipeline +-------- +1. Compute `dt` (distance transform) and `skel` (skeleton). +2. Compute LOCAL tau per pixel: mean DT over skeleton pixels in a + window. Adapts to varying stroke widths. +3. Threshold fat pixels: `dt > fat_ratio * local_tau` AND ink. +4. GROUP nearby fat fragments by dilation-then-cc-label, so a single + ribbon collapse whose fat region has split into two disconnected + peaks counts as one logical region. +5. For each group, count SKELETON-EXIT degree (8-connected components + of skeleton in a ring outside the dilated group). +6. For degree-4 groups, compute the best STROKE PAIRING score. +7. Apply chromosome filters: pairing score above threshold AND + skeleton-length-inside-fat above threshold. +""" + +from __future__ import annotations + +from typing import List, NamedTuple, Optional, Tuple, TypedDict, cast + +import numpy as np +from numpy.typing import NDArray +from scipy.ndimage import ( + binary_dilation, + distance_transform_edt, + label as _ndi_label, + uniform_filter, +) + +_CONNECTIVITY_8 = np.ones((3, 3), dtype=bool) + + +# ============================================================================ +# Config & output types +# ============================================================================ + + +class DetectConfig(TypedDict): + local_tau_radius: int + """Window half-size for local tau. ~40 px works on 1024-px sketches.""" + + fat_ratio: float + """Pixel is fat iff dt > fat_ratio * local_tau. ~1.3 is permissive.""" + + min_fat_area: int + """Drop fat groups below this combined pixel count.""" + + group_dilate: int + """Dilation iterations applied to fat_mask before grouping. Fat + fragments whose dilated extents touch are merged into one logical + region. Should be at least the maximum distance you expect between + fragments of one crossing. ~12-20 px works.""" + + skel_ring_dilate: int + """Additional dilation past `group_dilate` when looking for + skeleton arms in the exit ring.""" + + pairing_tangent_steps: int + """How far along each arm to look for its outward tangent. Larger + is more robust to noise but less local. 6-10 works.""" + + pairing_threshold: float + """Minimum allowed sum of (-dot products) over the best pairing + of 4 arms into 2 pairs. Both pairs perfectly anti-parallel sum + to 2.0; 1.0 means an average pair angle of 120 deg. ~1.2 is a + reasonable cutoff for 'the arms continue rather than meet'.""" + + min_chromosome_skel_length: int + """Minimum number of skeleton pixels that must sit INSIDE the + fat region for the detection to be reported as a ribbon-collapse + crossing. A true ribbon collapse forces the skeleton through the + merged region as a single line whose length is proportional to + the chromosome's extent; values of 15+ pixels on a 1024-px sketch + cleanly separate real ribbon collapses from 4-way point junctions + (where only the junction vertex itself sits in the fat region).""" + + +class Crossing(NamedTuple): + fat_pixels: NDArray[np.int32] + """(N, 2) integer (y, x) coordinates of the fat pixels in this group.""" + + group_id: int + """The group's label in ``DetectResult.group_labels``. Lets the + resolver retrieve the dilated group region (the area inside which + the merged skeleton segment lives) without recomputing the + dilation.""" + + degree: int + """Number of distinct skeleton arms exiting the group.""" + + peak_dt: float + """Maximum DT in the fat pixels.""" + + centroid: Tuple[float, float] + """(y, x) centroid.""" + + pairing_score: Optional[float] + """Best pairing score; only computed for degree-4 groups. None + otherwise.""" + + chromosome_skel_length: int + """Number of skeleton pixels that fall inside the fat region. + A measure of the merged-segment length. Tiny for 4-way junctions + (just the junction vertex), large for ribbon collapses.""" + + arm_endpoints: NDArray[np.int32] + """(degree, 2) array of (y, x) coordinates: each arm's proximal + end (the skeleton pixel closest to the dilated group region). + These pixels lie OUTSIDE the dilated group, so they survive when + the resolver erases the group's interior, and new replacement + lines can be drawn between paired endpoints to restore the two + crossing strokes.""" + + arm_tangents: NDArray[np.float64] + """(degree, 2) array of outward unit tangents, one per arm, + pointing from the group's centroid toward the arm.""" + + arm_pairing: Optional[Tuple[Tuple[int, int], Tuple[int, int]]] + """The best pairing of the 4 arms into 2 anti-parallel pairs, + given as a pair of (arm_index, arm_index) tuples (e.g. + ``((0, 2), (1, 3))``). The resolver draws one replacement line + per pair, between the two arms' endpoints. Only set for degree-4 + groups; ``None`` otherwise.""" + + +class DetectResult(NamedTuple): + local_tau: NDArray[np.float64] + dt: NDArray[np.float64] + skel: NDArray[np.bool_] + fat_mask: NDArray[np.bool_] + group_labels: NDArray[np.intp] + all_regions: List[Crossing] + crossings: List[Crossing] + """Degree-4 groups that pass the stroke-pairing check.""" + + +# ============================================================================ +# Internals +# ============================================================================ + + +def _label_tuple( + inp: NDArray, + structure: Optional[NDArray] = None, +) -> Tuple[NDArray[np.intp], int]: + result = _ndi_label(inp, structure=structure, output=None) + labels, n = result # type: ignore[misc] + return np.asarray(labels, dtype=np.intp), int(n) + + +def _local_tau( + skel: NDArray[np.bool_], + dt: NDArray[np.float64], + radius: int, +) -> NDArray[np.float64]: + """Per-pixel mean DT over skeleton pixels in a (2r+1) window. + + Implemented as the ratio of two `uniform_filter` outputs: one over + skel*dt (numerator), one over skel (denominator). Where the window + contains no skeleton pixels, falls back to the global median. + """ + skel_f = skel.astype(np.float64) + window = 2 * radius + 1 + weighted = uniform_filter(skel_f * dt, size=window, mode="constant", cval=0.0) + count = uniform_filter(skel_f, size=window, mode="constant", cval=0.0) + fallback = float(np.median(dt[skel])) if skel.any() else 1.0 + return np.where(count > 0, weighted / np.maximum(count, 1e-12), fallback) + + +def _arm_tangents_and_endpoints( + arm_labels: NDArray[np.intp], + n_arms: int, + centroid: Tuple[float, float], + steps: int, +) -> Tuple[NDArray[np.float64], NDArray[np.int32]]: + """For each skeleton arm, compute the outward tangent (unit vector) + AND the proximal endpoint (the arm's skeleton pixel closest to the + group centroid). + + Returns: + tangents: ``(n_arms, 2)`` float array of unit vectors pointing + from centroid toward the arm. + endpoints: ``(n_arms, 2)`` integer array of (y, x) coordinates + of each arm's proximal end. These pixels lie just outside + the dilated group region, on the original skeleton, and + serve as connection points for resolver-drawn lines. + + The tangent for each arm is the direction from the group centroid + to the centroid of the arm's first ``steps`` skeleton pixels + (ordered by distance to the group centroid). The endpoint is the + single closest pixel. + """ + cy, cx = centroid + tangents = np.zeros((n_arms, 2), dtype=np.float64) + endpoints = np.zeros((n_arms, 2), dtype=np.int32) + for arm_id in range(1, n_arms + 1): + ys, xs = np.where(arm_labels == arm_id) + idx = arm_id - 1 + if len(ys) == 0: + continue + d2 = (ys - cy) ** 2 + (xs - cx) ** 2 + order = np.argsort(d2) + head = order[: max(1, min(steps, len(order)))] + head_y = float(ys[head].mean()) + head_x = float(xs[head].mean()) + vec = np.array([head_y - cy, head_x - cx], dtype=np.float64) + n = float(np.linalg.norm(vec)) + if n >= 1e-9: + tangents[idx] = vec / n + # The proximal endpoint is the single closest pixel to the + # centroid. It is guaranteed to be in `arm_skel`, i.e. on the + # original skeleton and outside the dilated group region. + endpoints[idx, 0] = int(ys[order[0]]) + endpoints[idx, 1] = int(xs[order[0]]) + return tangents, endpoints + + +def _best_pairing( + tangents: NDArray[np.float64], +) -> Tuple[float, Tuple[Tuple[int, int], Tuple[int, int]]]: + """Best score AND the winning pairing over the 3 ways to pair 4 + unit vectors into 2 pairs. + + Each pair contributes ``-dot(t_i, t_j)``: 1.0 if anti-parallel + (the arms continue across the group as one stroke), -1.0 if + parallel (same direction; bad), 0.0 if perpendicular. Score is in + ``[-2, 2]``; ``>= 1.5`` ≈ both pairs at least 135 deg apart; + ``>= 1.2`` ≈ both pairs at least 120 deg. + """ + assert tangents.shape == (4, 2) + matchings = [ + ((0, 1), (2, 3)), + ((0, 2), (1, 3)), + ((0, 3), (1, 2)), + ] + best_score = -np.inf + best_match = matchings[0] + for match in matchings: + (a, b), (c, d) = match + s = -float(np.dot(tangents[a], tangents[b])) - float( + np.dot(tangents[c], tangents[d]) + ) + if s > best_score: + best_score = s + best_match = match + return float(best_score), best_match + + +# ============================================================================ +# Detection +# ============================================================================ + + +def detect_crossings( + binary: NDArray[np.bool_], + config: DetectConfig, + skel: NDArray[np.bool_], +) -> DetectResult: + """Detect ribbon-collapse crossings in ``binary``. + + See module docstring for the algorithm. + + Arguments: + binary: the binarized image (True where ink is present). + config: detection parameters; see ``DetectConfig`` for fields. + skel: pre-computed skeleton to use instead of running + ``skimage.morphology.skeletonize(binary)`` internally. Pass + this when you have a CLEANED skeleton (e.g. after + ``collapse_small_holes``) so that the arm endpoints land + on actual cleaned-skeleton pixels — important if the + resolver will then run on the cleaned skeleton. + """ + binary_b = binary.astype(bool) + dt = cast(NDArray[np.float64], distance_transform_edt(binary_b)) + + local_tau = _local_tau(skel, dt, config["local_tau_radius"]) + fat_mask = binary_b & (dt > config["fat_ratio"] * local_tau) + + # GROUPING: dilate fat_mask, then connected components on the + # dilation. Nearby fragments collapse into one group. + group_dilate = int(config["group_dilate"]) + dilated_fat = binary_dilation(fat_mask, iterations=group_dilate) + group_labels, n_groups = _label_tuple(dilated_fat, structure=_CONNECTIVITY_8) + + min_area = int(config["min_fat_area"]) + skel_ring_dilate = int(config["skel_ring_dilate"]) + ring_dilate = group_dilate + skel_ring_dilate + + pairing_steps = int(config["pairing_tangent_steps"]) + pairing_threshold = float(config["pairing_threshold"]) + min_chrom_length = int(config["min_chromosome_skel_length"]) + + all_regions: List[Crossing] = [] + crossings: List[Crossing] = [] + + for gid in range(1, n_groups + 1): + group_dilated_mask = group_labels == gid + fat_in_group = fat_mask & group_dilated_mask + area = int(fat_in_group.sum()) + if area < min_area: + continue + + # Skeleton arms in the exit ring. The ring is a band outside + # the dilated group region. Skel pixels inside the dilated + # region are excluded (they're either fat or within the fat's + # blur halo). Result: each clean arm appears as one 8-connected + # component, since arms are 1px wide and separated by background. + arm_region = binary_dilation(fat_in_group, iterations=ring_dilate) + arm_skel = arm_region & ~group_dilated_mask & skel + arm_labels, n_arms = _label_tuple(arm_skel, structure=_CONNECTIVITY_8) + + ys, xs = np.where(fat_in_group) + pixels = np.stack([ys, xs], axis=1).astype(np.int32) + peak_dt = float(dt[fat_in_group].max()) + centroid = (float(ys.mean()), float(xs.mean())) + + # Length of the merged segment: skeleton pixels INSIDE the fat + # region. Distinguishes ribbon collapses (long single-line + # segment forced through merged ink) from 4-way point junctions + # (just the junction vertex inside fat). + chrom_skel_len = int((fat_in_group & skel).sum()) + + # Compute arm tangents and endpoints for any group with >=1 + # detected arm; needed for visualization and (when degree==4) + # for the resolver to draw replacement line segments. + if n_arms > 0: + arm_tangents, arm_endpoints = _arm_tangents_and_endpoints( + arm_labels, + n_arms, + centroid, + pairing_steps, + ) + else: + arm_tangents = np.zeros((0, 2), dtype=np.float64) + arm_endpoints = np.zeros((0, 2), dtype=np.int32) + + pairing_score: Optional[float] = None + arm_pairing: Optional[Tuple[Tuple[int, int], Tuple[int, int]]] = None + passed_pairing = False + if n_arms == 4: + pairing_score, arm_pairing = _best_pairing(arm_tangents) + passed_pairing = pairing_score >= pairing_threshold + + region = Crossing( + fat_pixels=pixels, + group_id=gid, + degree=n_arms, + peak_dt=peak_dt, + centroid=centroid, + pairing_score=pairing_score, + chromosome_skel_length=chrom_skel_len, + arm_endpoints=arm_endpoints, + arm_tangents=arm_tangents, + arm_pairing=arm_pairing, + ) + all_regions.append(region) + if passed_pairing and chrom_skel_len >= min_chrom_length: + crossings.append(region) + + return DetectResult( + local_tau=local_tau, + dt=dt, + skel=skel, + fat_mask=fat_mask, + group_labels=group_labels, + all_regions=all_regions, + crossings=crossings, + ) + + +# ============================================================================ +# Resolution +# ============================================================================ + + +def resolve_crossings( + skel: NDArray[np.bool_], + detection: DetectResult, +) -> NDArray[np.bool_]: + """Rewrite the skeleton at each detected crossing to undo the + ribbon-collapse artifact. + + For each detected crossing, the merged skeleton segment that the + thinning algorithm wrote through the fat region is erased and + replaced by two new straight-line segments — one per paired pair of + arms. The new lines start from the original arm endpoints (which + lie just OUTSIDE the dilated group region and so are preserved + when the group interior is erased), so the surviving arm skeleton + connects to the replacement lines seamlessly. + + The result has two non-intersecting topological paths through each + crossing region, recovering the original two-stroke topology that + skeletonization had collapsed. Downstream segment tracing will now + follow each stroke through the crossing as a single curve. + + Arguments: + skel: the cleaned skeleton (one-pixel-wide ink). Modified + crossings will be written into a copy; ``skel`` is not + mutated. + detection: the output of ``detect_crossings``. The + ``crossings`` list provides the fat pixels, arm endpoints, + and pairings used by the resolver. + + Returns: + A new boolean array of the same shape as ``skel`` with each + detected crossing rewritten. Crossings with no pairing + (``arm_pairing is None``) are left alone. + + Notes: + - Straight Bresenham lines are used between paired endpoints. + For most crossings this is geometrically faithful (the two + strokes either cross at the centroid or run nearly parallel + through it). For strongly curved strokes a Bezier or + spline-based connection would be more accurate, but the + downstream pipeline typically re-fits curves to the resolved + skeleton, so the straight-line approximation is sufficient + here. + - Erasing the dilated group region (rather than just the fat + pixels) is important: the merged skeleton segment can extend + slightly beyond the fat-thresholded pixels into the + "shoulders" where the DT dips just below the fat threshold + but the strokes are still merged. The dilated group region + captures that shoulder zone. + """ + # Local imports keep the heavy skimage.draw module optional for + # users who only need detection. + from skimage.draw import line as _bresenham + + result = skel.copy() + group_labels = detection.group_labels + + for crossing in detection.crossings: + if crossing.arm_pairing is None: + continue + + # Erase the merged skeleton segment(s). The dilated group + # region (from group_labels) covers every pixel within + # group_dilate of any fat pixel in this crossing — which is + # exactly the area inside which the thinning algorithm wrote + # the merged centerline. + group_region = group_labels == crossing.group_id + result[group_region] = False + + # Draw the two replacement segments, one per paired pair. + H, W = result.shape + for i, j in crossing.arm_pairing: + y0, x0 = int(crossing.arm_endpoints[i, 0]), int( + crossing.arm_endpoints[i, 1] + ) + y1, x1 = int(crossing.arm_endpoints[j, 0]), int( + crossing.arm_endpoints[j, 1] + ) + rr, cc = _bresenham(y0, x0, y1, x1) + # Defensive: clip out anything that would land out of bounds. + valid = (rr >= 0) & (rr < H) & (cc >= 0) & (cc < W) + result[rr[valid], cc[valid]] = True + + return result diff --git a/arc_line_vectorization_suede/vectorize/__init__.py b/arc_line_vectorization_suede/vectorize/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/arc_line_vectorization_suede/vectorize/high_geometry/__init__.py b/arc_line_vectorization_suede/vectorize/high_geometry/__init__.py new file mode 100644 index 0000000..63724f5 --- /dev/null +++ b/arc_line_vectorization_suede/vectorize/high_geometry/__init__.py @@ -0,0 +1,604 @@ +"""PNG -> drawing-robot commands pipeline. + +Pipeline stages (matching the recipe): + + 1. skeletonize the binary image (skimage) + 2. trace the skeleton into polylines using a graph walk + 3. resample, smooth, and segment each polyline at curvature peaks (corners) + 4. classify each segment as line or arc (LS line + algebraic circle fit) + 5. order strokes greedily to minimize pen-up travel + 6. emit a list of DrawingCommand dicts matching the TS type +""" + +from __future__ import annotations +import math +from typing import List, Tuple, TypedDict + +import numpy as np +from numpy.typing import NDArray +from scipy.ndimage import gaussian_filter1d +from scipy.signal import find_peaks + +from ...commands import ( + ArcCommand, + ArcPrimitive, + DrawingCommand, + LineCommand, + LinePrimitive, + Primitive, + SpinCommand, + Stroke, +) + +from .consolidate import consolidate_commands, ConsolidateConfig + +# ============================================================================ +# Stage 3: curvature-based segmentation +# ============================================================================ + + +def _resample_polyline(poly: np.ndarray, n: int) -> Tuple[np.ndarray, np.ndarray]: + """Resample a polyline to `n` uniformly spaced points by arc length. + + Returns (resampled_points, original_arc_length_at_each_resample). + """ + diffs = np.diff(poly, axis=0) + seg_lens = np.linalg.norm(diffs, axis=1) + s = np.concatenate([[0.0], np.cumsum(seg_lens)]) + total = s[-1] + if total < 1e-9: + return poly.copy(), s + s_new = np.linspace(0.0, total, n) + x_new = np.interp(s_new, s, poly[:, 0]) + y_new = np.interp(s_new, s, poly[:, 1]) + return np.column_stack([x_new, y_new]), s_new + + +def _curvature(xy: np.ndarray, sigma: float) -> np.ndarray: + """Discrete signed curvature along a uniformly-spaced polyline.""" + xs = gaussian_filter1d(xy[:, 0], sigma) + ys = gaussian_filter1d(xy[:, 1], sigma) + dx = np.gradient(xs) + dy = np.gradient(ys) + ddx = np.gradient(dx) + ddy = np.gradient(dy) + denom = (dx * dx + dy * dy) ** 1.5 + denom = np.where(denom < 1e-10, 1e-10, denom) + return (dx * ddy - dy * ddx) / denom + + +def segment_at_corners( + poly: np.ndarray, + sigma: float = 2.0, + corner_threshold: float = 0.25, + min_segment_length: int = 4, +) -> List[np.ndarray]: + """Cut a polyline at peaks of |curvature|. + + Args: + poly: ordered (N, 2) points. + sigma: Gaussian smoothing scale (in resampled units). + corner_threshold: minimum |curvature| to be considered a corner. + min_segment_length: drop segments shorter than this many points. + """ + if len(poly) < 2 * min_segment_length: + return [poly] + + diffs = np.diff(poly, axis=0) + total_len = float(np.sum(np.linalg.norm(diffs, axis=1))) + if total_len < 4.0: + return [poly] + + n_samples = max(int(total_len), 16) + resampled, s_resampled = _resample_polyline(poly, n_samples) + kappa = _curvature(resampled, sigma) + abs_k = np.abs(kappa) + + peaks, _ = find_peaks( + abs_k, + height=corner_threshold, + distance=max(int(sigma * 3), 3), + ) + + # Map each peak (in resampled-index space) back to an index in `poly`. + if len(peaks) == 0: + return [poly] + + # Cumulative arc length along the original polyline: + s_orig = np.concatenate([[0.0], np.cumsum(np.linalg.norm(diffs, axis=1))]) + cut_points = [] + for p in peaks: + s_at_peak = s_resampled[p] + idx = int(np.searchsorted(s_orig, s_at_peak)) + idx = max(1, min(len(poly) - 2, idx)) + cut_points.append(idx) + cut_points = sorted(set(cut_points)) + + # Build segments, dropping those that are too short. + segments: List[np.ndarray] = [] + prev = 0 + for ci in cut_points: + if ci - prev >= min_segment_length: + # +1 to include the cut point as the segment endpoint + segments.append(poly[prev : ci + 1].copy()) + prev = ci + if len(poly) - prev >= min_segment_length: + segments.append(poly[prev:].copy()) + + return segments if segments else [poly] + + +# ============================================================================ +# Stage 4: line vs arc fitting +# ============================================================================ + + +def fit_line(pts: np.ndarray) -> Tuple[np.ndarray, np.ndarray, float]: + """Total-least-squares line fit; returns (start, end, max_perp_residual). + + `start` and `end` are the projections of the first/last points onto the + fitted axis. + """ + centroid = pts.mean(axis=0) + centered = pts - centroid + _, _, Vt = np.linalg.svd(centered, full_matrices=False) + direction = Vt[0] + perp = Vt[1] + projections = centered @ direction + perp_dists = np.abs(centered @ perp) + # Use the projection of the first/last polyline point so endpoints + # match the original ordering. + t_first = projections[0] + t_last = projections[-1] + start = centroid + t_first * direction + end = centroid + t_last * direction + return start, end, float(perp_dists.max()) + + +def fit_circle_algebraic(pts: np.ndarray): + """Algebraic circle fit (a refined version of Kåsa's method). + + Returns (cx, cy, r) or None if the system is degenerate. + """ + xs = pts[:, 0] + ys = pts[:, 1] + x_m = xs.mean() + y_m = ys.mean() + u = xs - x_m + v = ys - y_m + Suu = float((u * u).sum()) + Svv = float((v * v).sum()) + Suv = float((u * v).sum()) + Suuu = float((u * u * u).sum()) + Svvv = float((v * v * v).sum()) + Suvv = float((u * v * v).sum()) + Svuu = float((v * u * u).sum()) + A = np.array([[Suu, Suv], [Suv, Svv]]) + b = 0.5 * np.array([Suuu + Suvv, Svvv + Svuu]) + try: + uc, vc = np.linalg.solve(A, b) + except np.linalg.LinAlgError: + return None + cx = uc + x_m + cy = vc + y_m + r = math.sqrt(uc * uc + vc * vc + (Suu + Svv) / len(pts)) + return cx, cy, r + + +def _arc_traversal_ccw(pts: np.ndarray, cx: float, cy: float) -> bool: + """True if the polyline traverses around (cx, cy) counter-clockwise.""" + angles = np.arctan2(pts[:, 1] - cy, pts[:, 0] - cx) + unwrapped = np.unwrap(angles) + return float(unwrapped[-1] - unwrapped[0]) > 0.0 + + +class Arc(TypedDict): + center: np.ndarray + radius: float + residual: float + ccw: bool + + +def fit_arc(pts: np.ndarray) -> Arc | None: + """Fit a circular arc; returns dict with center/radius/residual or None.""" + if len(pts) < 3: + return None + fit = fit_circle_algebraic(pts) + if fit is None: + return None + cx, cy, r = fit + if not math.isfinite(r) or r < 0.5: + return None + dists = np.sqrt((pts[:, 0] - cx) ** 2 + (pts[:, 1] - cy) ** 2) + residual = float(np.abs(dists - r).max()) + return { + "center": np.array([cx, cy]), + "radius": r, + "ccw": _arc_traversal_ccw(pts, cx, cy), + "residual": residual, + } + + +def _segment_scale(pts: np.ndarray) -> float: + """A characteristic length for a segment, robust to closed paths. + + For an open polyline, this is close to the chord length. For a closed + loop (where start == end and the chord length is zero), this returns + the bounding-box diagonal, which is a meaningful spatial scale. + """ + bbox = pts.max(axis=0) - pts.min(axis=0) + diag = float(np.linalg.norm(bbox)) + return max(diag, 1.0) + + +def classify_segment( + pts: np.ndarray, + arc_advantage: float = 1.5, + max_radius_factor: float = 8.0, + min_arc_radius: float = 1.5, +) -> Primitive: + """Choose between a line and an arc primitive for a single segment. + + The arc is preferred only if its residual is meaningfully better than the + line residual AND its radius is reasonable (extremely large radii are + indistinguishable from lines and should be drawn as lines). + """ + line_start, line_end, line_res = fit_line(pts) + scale = _segment_scale(pts) + arc = fit_arc(pts) + + use_arc = ( + arc is not None + and arc["radius"] >= min_arc_radius + and arc["radius"] <= max_radius_factor * scale + and arc["residual"] * arc_advantage < max(line_res, 1e-6) + ) + + if use_arc and arc is not None: + return ArcPrimitive( + center=arc["center"], + radius=arc["radius"], + # Snap to the actual polyline endpoints so adjacent primitives + # remain connected. Reproject onto the fitted circle for sanity. + start=_project_to_circle(pts[0], arc["center"], arc["radius"]), + end=_project_to_circle(pts[-1], arc["center"], arc["radius"]), + ccw=arc["ccw"], + ) + + return LinePrimitive(start=pts[0].copy(), end=pts[-1].copy()) + + +def _project_to_circle(p: np.ndarray, center: np.ndarray, r: float) -> np.ndarray: + d = p - center + n = np.linalg.norm(d) + if n < 1e-9: + return p.copy() + return center + (r / n) * d + + +def _best_fit_with_residual( + pts: np.ndarray, + arc_advantage: float = 1.5, + max_radius_factor: float = 8.0, + min_arc_radius: float = 1.5, +) -> Tuple[Primitive, float]: + """Same logic as `classify_segment` but also returns the absolute + residual (in pixels) of the chosen fit. Used by the residual-based + recursive subdivision in `polyline_to_stroke`. + """ + line_start, line_end, line_res = fit_line(pts) + scale = _segment_scale(pts) + arc = fit_arc(pts) + + use_arc = ( + arc is not None + and arc["radius"] >= min_arc_radius + and arc["radius"] <= max_radius_factor * scale + and arc["residual"] * arc_advantage < max(line_res, 1e-6) + ) + if use_arc and arc is not None: + prim = ArcPrimitive( + center=arc["center"], + radius=arc["radius"], + start=_project_to_circle(pts[0], arc["center"], arc["radius"]), + end=_project_to_circle(pts[-1], arc["center"], arc["radius"]), + ccw=arc["ccw"], + ) + return prim, float(arc["residual"]) + return LinePrimitive(start=pts[0].copy(), end=pts[-1].copy()), float(line_res) + + +# ============================================================================ +# Stage 5: build strokes and order them +# ============================================================================ + + +def polyline_to_stroke( + poly: np.ndarray, + sigma: float = 2.0, + corner_threshold: float = 0.25, + max_fit_residual: float | None = None, + min_split_length: int = 4, +) -> Stroke: + """Segment a polyline at corners and fit a primitive per segment. + + If `max_fit_residual` is set (in absolute pixel units), every segment + whose best line/arc fit has residual greater than that threshold is + recursively split at its midpoint until each piece fits within + tolerance or becomes shorter than `min_split_length` points. This is + the right knob for raising fidelity: corner-based segmentation alone + doesn't help when a polyline has *gradually* varying curvature (like + a heart or S-curve) — there's no curvature peak to cut at, and a + single line/arc cannot capture the shape. Residual-based splitting + forces extra cuts wherever a curve isn't well approximated by one + primitive, regardless of whether there's a "corner" there. + + Set `max_fit_residual=None` (the default) for original behaviour: + only corner-based segmentation, no residual cap. + """ + segments = segment_at_corners(poly, sigma=sigma, corner_threshold=corner_threshold) + if max_fit_residual is None: + return [classify_segment(s) for s in segments if len(s) >= 2] + + primitives: Stroke = [] + # DFS via stack: segments processed left-to-right after splits. + stack = list(reversed([s for s in segments if len(s) >= 2])) + while stack: + seg = stack.pop() + prim, residual = _best_fit_with_residual(seg) + if residual <= max_fit_residual or len(seg) < 2 * min_split_length: + primitives.append(prim) + continue + mid = len(seg) // 2 + # Both halves share the midpoint so adjacent primitives connect. + left = seg[: mid + 1] + right = seg[mid:] + # Push right first so left is processed next. + stack.append(right) + stack.append(left) + return primitives + + +def _stroke_endpoints(stroke: Stroke) -> Tuple[np.ndarray, np.ndarray]: + return stroke[0].start, stroke[-1].end + + +def _reverse_primitive(p: Primitive) -> Primitive: + if isinstance(p, LinePrimitive): + return LinePrimitive(start=p.end.copy(), end=p.start.copy()) + return ArcPrimitive( + center=p.center.copy(), + radius=p.radius, + start=p.end.copy(), + end=p.start.copy(), + ccw=not p.ccw, + ) + + +def _reverse_stroke(stroke: Stroke) -> Stroke: + return [_reverse_primitive(p) for p in reversed(stroke)] + + +def order_strokes(strokes: List[Stroke], origin: np.ndarray) -> List[Stroke]: + """Greedy nearest-neighbour ordering; reverses strokes when their far end + is closer to the cursor.""" + remaining = list(range(len(strokes))) + cursor = origin.astype(float).copy() + ordered: List[Stroke] = [] + while remaining: + best_i = remaining[0] + best_dist = math.inf + best_reverse = False + for i in remaining: + s, e = _stroke_endpoints(strokes[i]) + ds = float(np.linalg.norm(cursor - s)) + de = float(np.linalg.norm(cursor - e)) + if ds < best_dist: + best_dist, best_i, best_reverse = ds, i, False + if de < best_dist: + best_dist, best_i, best_reverse = de, i, True + chosen = strokes[best_i] + if best_reverse: + chosen = _reverse_stroke(chosen) + ordered.append(chosen) + cursor = chosen[-1].end.copy() + remaining.remove(best_i) + return ordered + + +# ============================================================================ +# Stage 6: ordered strokes -> DrawingCommand list +# ============================================================================ + + +def _wrap_pi(a: float) -> float: + """Wrap an angle to (-pi, pi].""" + return ((a + math.pi) % (2 * math.pi)) - math.pi + + +def _emit_spin_to(target_heading: float, heading: float, eps: float = 1e-4): + """Yield a SpinCommand if needed to align `heading` with `target_heading`. + + Returns (cmd_or_None, new_heading). + """ + delta = _wrap_pi(target_heading - heading) + if abs(delta) < eps: + return None, heading + return SpinCommand(kind="spin", degrees=math.degrees(delta)), target_heading + + +def _emit_primitive(prim: Primitive, pos: np.ndarray, heading: float): + """Convert a single primitive into commands. + + Returns (cmds_list, new_pos, new_heading). + """ + cmds: List[DrawingCommand] = [] + + if isinstance(prim, LinePrimitive): + delta = prim.end - pos + dist = float(np.linalg.norm(delta)) + if dist < 1e-6: + return cmds, pos, heading + target_heading = math.atan2(delta[1], delta[0]) + spin_cmd, heading = _emit_spin_to(target_heading, heading) + if spin_cmd is not None: + cmds.append(spin_cmd) + cmds.append(LineCommand(kind="line", distance=dist, penDown=True)) + return cmds, prim.end.copy(), heading + + # ArcPrimitive + radial = prim.start - prim.center + # Tangent: 90deg CCW from radial if traversing CCW, else 90deg CW. + if prim.ccw: + tangent = np.array([-radial[1], radial[0]]) + else: + tangent = np.array([radial[1], -radial[0]]) + target_heading = math.atan2(tangent[1], tangent[0]) + spin_cmd, heading = _emit_spin_to(target_heading, heading) + if spin_cmd is not None: + cmds.append(spin_cmd) + + start_a = math.atan2(prim.start[1] - prim.center[1], prim.start[0] - prim.center[0]) + end_a = math.atan2(prim.end[1] - prim.center[1], prim.end[0] - prim.center[0]) + if prim.ccw: + sweep = end_a - start_a + if sweep <= 0: + sweep += 2 * math.pi + else: + sweep = end_a - start_a + if sweep >= 0: + sweep -= 2 * math.pi + + cmds.append( + ArcCommand( + kind="arc", + radius=float(prim.radius), + degrees=math.degrees(sweep), + ) + ) + new_heading = _wrap_pi(heading + sweep) + return cmds, prim.end.copy(), new_heading + + +def strokes_to_commands( + strokes: List[Stroke], + start_pos: np.ndarray | None = None, + start_heading: float = 0.0, +) -> List[DrawingCommand]: + """Walk an ordered list of strokes and produce a flat command sequence. + + Between strokes, emits (spin, line penDown=False) to traverse pen-up. + Within a stroke, emits a spin to align the heading with each primitive's + starting tangent, then the primitive itself. + """ + if start_pos is None: + start_pos = np.zeros(2) + pos = start_pos.astype(float).copy() + heading = float(start_heading) + out: List[DrawingCommand] = [] + + for stroke in strokes: + if not stroke: + continue + + # Pen-up traversal to the start of this stroke. + target = stroke[0].start + delta = target - pos + dist = float(np.linalg.norm(delta)) + if dist > 1e-6: + target_heading = math.atan2(delta[1], delta[0]) + spin_cmd, heading = _emit_spin_to(target_heading, heading) + if spin_cmd is not None: + out.append(spin_cmd) + out.append(LineCommand(kind="line", distance=dist, penDown=False)) + pos = target.copy() + + # Draw the stroke. + for prim in stroke: + cmds, pos, heading = _emit_primitive(prim, pos, heading) + out.extend(cmds) + + return out + + +# ============================================================================ +# Top-level driver +# ============================================================================ + + +def polylines_to_commands( + polylines: List[np.ndarray], + sigma: float = 2.0, + corner_threshold: float = 0.25, + max_fit_residual: float | None = None, + start_pos: np.ndarray | None = None, + start_heading: float = 0.0, +) -> List[DrawingCommand]: + """Run the polyline -> commands half of the pipeline. + + Each polyline is segmented at curvature peaks, every segment is + classified as a line or an arc, the resulting strokes are ordered + greedily for minimum pen-up travel, and the whole thing is emitted + as a flat command list. This is the part of the pipeline that does + not care where the polylines came from — raw skeleton tracing, the + fusion stage, or anything else that produces (N, 2) pixel polylines + is fine input. + + If `max_fit_residual` is set, segments are recursively subdivided + until each one fits a single line or arc with residual at or below + that pixel threshold. Useful when corner-based segmentation alone + does not produce enough cuts on smoothly-varying curves. + """ + strokes = [ + polyline_to_stroke( + p, + sigma=sigma, + corner_threshold=corner_threshold, + max_fit_residual=max_fit_residual, + ) + for p in polylines + ] + strokes = [s for s in strokes if s] + if start_pos is None: + start_pos = np.zeros(2) + ordered = order_strokes(strokes, origin=np.asarray(start_pos, dtype=float)) + return strokes_to_commands( + ordered, + start_pos=start_pos, + start_heading=start_heading, + ) + + +class Vectorize: + class Config: + class ToCommands(TypedDict): + sigma: float + corner_threshold: float + max_fit_residual: float | None + + class Consolidate(ConsolidateConfig): + pass + + def __init__( + self, + polylines: List[NDArray[np.float64]], + start_pos: NDArray[np.float64], + start_heading: float, + commands: Config.ToCommands, + consolidate: Config.Consolidate, + ): + self.commands = polylines_to_commands( + polylines, + sigma=commands["sigma"], + corner_threshold=commands["corner_threshold"], + max_fit_residual=commands["max_fit_residual"], + start_pos=start_pos, + start_heading=start_heading, + ) + + self.consolidated, self.report = consolidate_commands( + self.commands, + start_pos=start_pos, + start_heading=start_heading, + config=consolidate, + ) diff --git a/arc_line_vectorization_suede/vectorize/high_geometry/consolidate.py b/arc_line_vectorization_suede/vectorize/high_geometry/consolidate.py new file mode 100644 index 0000000..74d0f70 --- /dev/null +++ b/arc_line_vectorization_suede/vectorize/high_geometry/consolidate.py @@ -0,0 +1,1223 @@ +"""Post-pass over a finished DrawingCommand list. + +The vectorizer fits each segment as an independent line or arc, so a +single underlying circle in the input (a wheel, an eye, a head outline) +that gets fragmented at junctions ends up as several arcs with slightly +different (center, radius) parameters. Visually the result looks broken +even though each piece is a good fit on its own. + +This module does a pure post-pass on the emitted commands: + + 1. Reverse-simulate the command sequence to recover geometric + primitives (each with start, end, and arc params if applicable). + 2. Cluster pen-down arcs whose (center, radius) are similar. + 3. For every cluster of size >= 2, compute a consensus circle + (sweep-weighted average) and snap every arc in the cluster onto + that circle: its center and radius are replaced, and its endpoints + are projected radially onto the new circle. + 4. Snapping moves endpoints, which creates disjoints with neighbouring + primitives. Walk the primitive list and stitch every gap so that + prim[i+1].start == prim[i].end exactly. For neighbouring arcs that + are not in any cluster, refit center and radius so the moved + endpoint stays on a circle (preserving the original bend + direction). + 5. Re-emit a fresh command sequence from the modified primitives. + +The output is self-consistent: every arc primitive satisfies +|start - center| = |end - center| = radius, and adjacent primitives +share their connecting endpoint exactly. That is what guarantees the +robot doesn't drift -- the emitter computes spin/distance/sweep from +the primitive's geometry, and any inconsistency would manifest as +accumulating offset along the chain. +""" + +from __future__ import annotations +import math +from dataclasses import dataclass, field +from typing import Any, Dict, List, Optional, Sequence, Tuple, TypedDict + +import numpy as np +from numpy.typing import NDArray + +from .junction_graph import JunctionGraph +from ...commands import DrawingCommand + +# ============================================================================= +# Internal geometric representation +# ============================================================================= + + +@dataclass +class _Primitive: + """One geometric primitive between two robot states. + + `pen_down` is False only for inter-stroke pen-up traversals (which + have free endpoints and therefore aren't constrained by the stitcher + on their END side -- they just go wherever the next primitive starts). + Arc fields are populated only when kind == "arc". + """ + + kind: str # "line" | "arc" + start: NDArray[np.float64] + end: NDArray[np.float64] + pen_down: bool = True + # arc-only: + center: Optional[NDArray[np.float64]] = None + radius: float = 0.0 + ccw: bool = True + + +def _wrap_pi(a: float) -> float: + return ((a + math.pi) % (2 * math.pi)) - math.pi + + +# ============================================================================= +# Stage 1: simulate commands -> primitives +# ============================================================================= + + +def _commands_to_primitives( + commands: Sequence[DrawingCommand], + start_pos: NDArray[np.float64], + start_heading: float, +) -> List[_Primitive]: + """Walk the command list while tracking (pos, heading); emit one + primitive per line/arc command (spins are absorbed into the heading). + """ + pos = np.asarray(start_pos, dtype=float).copy() + heading = float(start_heading) + prims: List[_Primitive] = [] + for c in commands: + kind = c["kind"] + if kind == "spin": + heading = _wrap_pi(heading + math.radians(c["degrees"])) + elif kind == "line": + dist = float(c["distance"]) + new_pos = pos + dist * np.array([math.cos(heading), math.sin(heading)]) + prims.append( + _Primitive( + kind="line", + start=pos.copy(), + end=new_pos.copy(), + pen_down=bool(c.get("penDown", True)), + ) + ) + pos = new_pos + elif kind == "arc": + r = float(c["radius"]) + sweep = math.radians(float(c["degrees"])) + ccw = sweep > 0 + # Center is perpendicular to heading at distance r; left for + # CCW, right for CW. + if ccw: + perp = np.array([-math.sin(heading), math.cos(heading)]) + else: + perp = np.array([math.sin(heading), -math.cos(heading)]) + center = pos + r * perp + # Rotate the radial vector around the center by `sweep`. + radial = pos - center + cs, sn = math.cos(sweep), math.sin(sweep) + new_radial = np.array( + [ + radial[0] * cs - radial[1] * sn, + radial[0] * sn + radial[1] * cs, + ] + ) + new_pos = center + new_radial + prims.append( + _Primitive( + kind="arc", + start=pos.copy(), + end=new_pos.copy(), + pen_down=True, + center=center, + radius=r, + ccw=ccw, + ) + ) + pos = new_pos + heading = _wrap_pi(heading + sweep) + else: + raise ValueError(f"unknown command kind: {kind!r}") + return prims + + +# ============================================================================= +# Stage 2: cluster arcs by shared underlying circle +# ============================================================================= + + +def _cluster_arcs( + prims: List[_Primitive], + center_tol_rel: float, + radius_tol_rel: float, + center_tol_abs: float, + radius_tol_abs: float, +) -> List[List[int]]: + """Union-find arcs with similar (center, radius). Tolerances are + 'either / or': two arcs are linked if their centers AND radii are + close, where close = max(rel * R, abs). + Returns clusters as lists of primitive indices (only clusters of + size >= 2 are returned).""" + arc_idxs = [i for i, p in enumerate(prims) if p.kind == "arc" and p.pen_down] + n = len(arc_idxs) + parent = list(range(n)) + + def find(x: int) -> int: + while parent[x] != x: + parent[x] = parent[parent[x]] + x = parent[x] + return x + + def union(a: int, b: int) -> None: + ra, rb = find(a), find(b) + if ra != rb: + parent[ra] = rb + + for ii in range(n): + pi = prims[arc_idxs[ii]] + for jj in range(ii + 1, n): + pj = prims[arc_idxs[jj]] + r_max = max(pi.radius, pj.radius) + ctol = max(center_tol_rel * r_max, center_tol_abs) + rtol = max(radius_tol_rel * r_max, radius_tol_abs) + cdist = float(np.linalg.norm(pi.center - pj.center)) + rdiff = abs(pi.radius - pj.radius) + if cdist < ctol and rdiff < rtol: + union(ii, jj) + + bucket: Dict[int, List[int]] = {} + for ii in range(n): + bucket.setdefault(find(ii), []).append(arc_idxs[ii]) + return [g for g in bucket.values() if len(g) >= 2] + + +def _arc_signed_sweep(p: _Primitive) -> float: + """Sweep angle in radians (positive for CCW, negative for CW), as + you would compute from the primitive's start, end and center.""" + a_s = math.atan2(p.start[1] - p.center[1], p.start[0] - p.center[0]) + a_e = math.atan2(p.end[1] - p.center[1], p.end[0] - p.center[0]) + sweep = a_e - a_s + if p.ccw and sweep <= 0: + sweep += 2 * math.pi + elif not p.ccw and sweep >= 0: + sweep -= 2 * math.pi + return sweep + + +def _consensus_circle( + prims: List[_Primitive], cluster: List[int] +) -> Tuple[NDArray, float]: + """Sweep-weighted average of (center, radius). Longer arcs vote more + because they carry more information about the underlying circle.""" + weights, centers, radii = [], [], [] + for idx in cluster: + p = prims[idx] + # Minimum weight floors the contribution of vanishingly short arcs + # without dropping them entirely. + w = max(abs(_arc_signed_sweep(p)), 0.05) + weights.append(w) + centers.append(p.center) + radii.append(p.radius) + weights = np.asarray(weights) + weights = weights / weights.sum() + cc = np.zeros(2) + for w, c in zip(weights, centers): + cc = cc + w * c + cr = float(sum(w * r for w, r in zip(weights, radii))) + return cc, cr + + +# ============================================================================= +# Stage 3: snap each clustered arc to its consensus circle +# ============================================================================= + + +def _project_onto_circle(p: NDArray, center: NDArray, radius: float) -> NDArray: + v = p - center + d = float(np.linalg.norm(v)) + if d < 1e-9: + return center + np.array([radius, 0.0]) + return center + (radius / d) * v + + +def _move_endpoint( + prims: List[_Primitive], + idx: int, + attr: str, + new_pos: NDArray, + jg: Optional[JunctionGraph], + skip_refit: Optional[set] = None, +) -> List[Tuple[int, str]]: + """Move a single endpoint, optionally propagating to its junction-mates. + + If ``jg`` is None, only ``prims[idx].`` is updated. If ``jg`` is + provided, every endpoint sharing the same junction is moved to + ``new_pos``, and each propagated-to arc whose index is NOT in + ``skip_refit`` is refit through its (now-changed) endpoints. The + ``skip_refit`` set typically contains all arc cluster members -- + those will have their (center, radius) controlled by their own + cluster's consensus and shouldn't be refit by accident. + + Returns the propagation list (excluding the requested endpoint). + """ + if jg is None: + p = prims[idx] + np_pos = np.asarray(new_pos, dtype=float).copy() + if attr == "start": + p.start = np_pos + else: + p.end = np_pos + return [] + moved_others = jg.move(prims, idx, attr, new_pos) + if skip_refit is None: + skip_refit = set() + for oi, _ in moved_others: + if prims[oi].kind == "arc" and oi not in skip_refit: + _refit_arc_through_endpoints(prims[oi]) + return moved_others + + +def _snap_arc_to_consensus( + p: _Primitive, + idx: int, + cc: NDArray, + cr: float, + prims: List[_Primitive], + jg: Optional[JunctionGraph] = None, + skip_refit: Optional[set] = None, +) -> None: + """Replace the arc's center and radius with the cluster consensus, + and project both endpoints radially onto the new circle. The arc's + ccw direction is preserved. If ``jg`` is provided, the endpoint + moves are propagated to all junction-mates (so that other strokes + that touched this arc at its endpoints follow along).""" + p.center = cc.copy() + p.radius = cr + new_start = _project_onto_circle(p.start, cc, cr) + new_end = _project_onto_circle(p.end, cc, cr) + _move_endpoint(prims, idx, "start", new_start, jg, skip_refit) + _move_endpoint(prims, idx, "end", new_end, jg, skip_refit) + + +# ============================================================================= +# Stage 2b/3b: cluster collinear lines and snap them to a shared consensus +# ============================================================================= + + +def _line_geometry(p: _Primitive) -> Tuple[NDArray, float]: + """Return the line's unit direction and length. Caller is expected to + have filtered out zero-length lines already.""" + v = p.end - p.start + L = float(np.linalg.norm(v)) + return v / L, L + + +def _cluster_lines( + prims: List[_Primitive], + angle_tol_rad: float, + offset_tol_abs: float, + min_length: float, +) -> List[List[int]]: + """Group pen-down lines that are collinear: same direction within + `angle_tol_rad` (sign-ambiguous, since a line has no preferred + orientation), and same perpendicular offset within `offset_tol_abs` + pixels. Lines shorter than `min_length` are excluded because their + direction estimate is too noisy.""" + line_idxs: List[int] = [] + dirs: List[NDArray] = [] + for i, p in enumerate(prims): + if p.kind != "line" or not p.pen_down: + continue + L = float(np.linalg.norm(p.end - p.start)) + if L < min_length: + continue + line_idxs.append(i) + dirs.append((p.end - p.start) / L) + + n = len(line_idxs) + parent = list(range(n)) + + def find(x: int) -> int: + while parent[x] != x: + parent[x] = parent[parent[x]] + x = parent[x] + return x + + def union(a: int, b: int) -> None: + ra, rb = find(a), find(b) + if ra != rb: + parent[ra] = rb + + cos_tol = math.cos(angle_tol_rad) + for ii in range(n): + pi = prims[line_idxs[ii]] + di = dirs[ii] + for jj in range(ii + 1, n): + pj = prims[line_idxs[jj]] + dj = dirs[jj] + # Direction parallel? Use |dot| because a line has no orientation. + if abs(float(np.dot(di, dj))) < cos_tol: + continue + # Perpendicular offset: distance from one line's start to the + # other's infinite line. + perp_j = np.array([-dj[1], dj[0]]) + offset = abs(float(np.dot(perp_j, pi.start - pj.start))) + if offset < offset_tol_abs: + union(ii, jj) + + bucket: Dict[int, List[int]] = {} + for ii in range(n): + bucket.setdefault(find(ii), []).append(line_idxs[ii]) + return [g for g in bucket.values() if len(g) >= 2] + + +def _consensus_line( + prims: List[_Primitive], + cluster: List[int], +) -> Tuple[NDArray, NDArray]: + """Total-least-squares line through all endpoints of the clustered + lines. Returns (centroid, unit_direction); the infinite consensus + line is the set of points centroid + t * direction for any t in R.""" + pts = [] + for idx in cluster: + p = prims[idx] + pts.append(p.start) + pts.append(p.end) + arr = np.asarray(pts) + centroid = arr.mean(axis=0) + centered = arr - centroid + _, _, Vt = np.linalg.svd(centered, full_matrices=False) + direction = Vt[0] + return centroid, direction + + +def _project_onto_line(p: NDArray, centroid: NDArray, direction: NDArray) -> NDArray: + return centroid + float(np.dot(p - centroid, direction)) * direction + + +def _snap_line_to_consensus( + p: _Primitive, + idx: int, + centroid: NDArray, + direction: NDArray, + prims: List[_Primitive], + jg: Optional[JunctionGraph] = None, + skip_refit: Optional[set] = None, +) -> None: + """Project both endpoints of a line primitive onto the consensus + infinite line (perpendicular projection). If ``jg`` is provided, + the moves propagate to all junction-mates.""" + new_start = _project_onto_line(p.start, centroid, direction) + new_end = _project_onto_line(p.end, centroid, direction) + _move_endpoint(prims, idx, "start", new_start, jg, skip_refit) + _move_endpoint(prims, idx, "end", new_end, jg, skip_refit) + + +def _max_line_endpoint_snap_distance( + prims: List[_Primitive], + cluster: List[int], + centroid: NDArray, + direction: NDArray, +) -> float: + """Max perpendicular distance any endpoint must travel to reach the + consensus line. Large values indicate the cluster is bogus.""" + perp = np.array([-direction[1], direction[0]]) + max_d = 0.0 + for idx in cluster: + p = prims[idx] + for endpoint in (p.start, p.end): + d = abs(float(np.dot(perp, endpoint - centroid))) + if d > max_d: + max_d = d + return max_d + + +# ============================================================================= +# Stage 4: stitching (refit non-clustered arcs whose endpoints moved) +# ============================================================================= + + +def _refit_arc_through_endpoints(p: _Primitive) -> None: + """The arc's start and end no longer satisfy + |x - center| = radius. Find the unique arc that PASSES THROUGH the + current start and end with center on the perpendicular bisector, + closest to the original center. Preserves the original bending side + (and therefore the ccw flag).""" + chord = p.end - p.start + chord_len = float(np.linalg.norm(chord)) + if chord_len < 1e-9: + return # degenerate + midpoint = 0.5 * (p.start + p.end) + # perpendicular to chord (unit) + perp = np.array([-chord[1], chord[0]]) / chord_len + # Match the original bending side + if float(np.dot(p.center - midpoint, perp)) < 0: + perp = -perp + # Project the original center onto the perpendicular line + t = float(np.dot(p.center - midpoint, perp)) + new_center = midpoint + t * perp + new_radius = math.sqrt(t * t + (chord_len / 2.0) ** 2) + p.center = new_center + p.radius = new_radius + # ccw is preserved + + +def _stitch( + prims: List[_Primitive], + consensus_arcs: Dict[int, Tuple[NDArray, float]], + consensus_lines: Dict[int, Tuple[NDArray, NDArray]], + jg: Optional[JunctionGraph] = None, +) -> None: + """Walk the primitive list in order. For each adjacent pen-down pair + whose connecting endpoints disagree, reconcile while preserving each + side's consensus geometry where possible. Endpoint mutations go + through the junction graph if provided, so any fix here also + propagates to other strokes that touched the same junction. + + Resolution rules: + - Pen-up "line" primitives are unconstrained on both ends; they + just snap to whatever the neighbours dictate. The symmetric + case (prev is pen-up) is also handled, so a stroke whose + terminal endpoint moved earlier (e.g. by a proximity snap) + doesn't get pulled back by an unconstrained pen-up. + - If exactly one of (prev, cur) is on a consensus geometry, the + other adapts: the moving endpoint matches the clustered side; + non-clustered arc neighbours get refit through new endpoints. + - If neither side is clustered, take the midpoint and refit. + - If both sides are clustered (rare, e.g. a circle meeting a + line at a junction), take the midpoint, project each onto its + own consensus, average once more. Sub-pixel residual accepted. + """ + EPS = 1e-3 + cluster_member_set = set(consensus_arcs) | set(consensus_lines) + + def set_endpoint(idx: int, attr: str, new_pos: NDArray) -> None: + _move_endpoint(prims, idx, attr, new_pos, jg, cluster_member_set) + + def project_to_own(idx: int, attr: str) -> None: + """Project the (start|end) of prims[idx] onto its consensus + geometry, if any. Otherwise no-op.""" + p = prims[idx] + pt = getattr(p, attr) + if idx in consensus_arcs: + cc, cr = consensus_arcs[idx] + v = pt - cc + d = float(np.linalg.norm(v)) + if d > 1e-9: + set_endpoint(idx, attr, cc + (cr / d) * v) + elif idx in consensus_lines: + centroid, direction = consensus_lines[idx] + set_endpoint(idx, attr, _project_onto_line(pt, centroid, direction)) + + def is_clustered(idx: int) -> bool: + return idx in consensus_arcs or idx in consensus_lines + + def is_pen_up(p: _Primitive) -> bool: + return p.kind == "line" and not p.pen_down + + for i in range(1, len(prims)): + prev = prims[i - 1] + cur = prims[i] + + # Pen-up traversals adapt to whichever neighbour is constrained. + if is_pen_up(cur): + set_endpoint(i, "start", prev.end) + continue + if is_pen_up(prev): + set_endpoint(i - 1, "end", cur.start) + continue + + gap = float(np.linalg.norm(cur.start - prev.end)) + if gap < EPS: + set_endpoint(i, "start", prev.end) + continue + + prev_clustered = is_clustered(i - 1) + cur_clustered = is_clustered(i) + + if prev_clustered and not cur_clustered: + set_endpoint(i, "start", prev.end) + if cur.kind == "arc": + _refit_arc_through_endpoints(cur) + elif cur_clustered and not prev_clustered: + set_endpoint(i - 1, "end", cur.start) + if prev.kind == "arc": + _refit_arc_through_endpoints(prev) + elif prev_clustered and cur_clustered: + mid = 0.5 * (prev.end + cur.start) + set_endpoint(i - 1, "end", mid) + set_endpoint(i, "start", mid) + project_to_own(i - 1, "end") + project_to_own(i, "start") + mid = 0.5 * (prev.end + cur.start) + set_endpoint(i - 1, "end", mid) + set_endpoint(i, "start", mid) + else: + mid = 0.5 * (prev.end + cur.start) + set_endpoint(i - 1, "end", mid) + set_endpoint(i, "start", mid) + if prev.kind == "arc": + _refit_arc_through_endpoints(prev) + if cur.kind == "arc": + _refit_arc_through_endpoints(cur) + + +# ============================================================================= +# Stage 5: primitives -> commands (re-emit) +# ============================================================================= + + +def _primitives_to_commands( + prims: List[_Primitive], + start_pos: NDArray, + start_heading: float, +) -> List[DrawingCommand]: + """Walk the (now self-consistent) primitive list and emit + spin / line / arc commands. Mirrors the existing `strokes_to_commands` + behaviour but drives off the primitive list rather than strokes.""" + pos = np.asarray(start_pos, dtype=float).copy() + heading = float(start_heading) + out: List[DrawingCommand] = [] + EPS_HEADING = 1e-4 + EPS_DIST = 1e-6 + + for p in prims: + if p.kind == "line": + delta = p.end - pos + dist = float(np.linalg.norm(delta)) + if dist < EPS_DIST: + continue + target_h = math.atan2(delta[1], delta[0]) + spin = _wrap_pi(target_h - heading) + if abs(spin) > EPS_HEADING: + out.append({"kind": "spin", "degrees": math.degrees(spin)}) + heading = target_h + out.append({"kind": "line", "distance": dist, "penDown": p.pen_down}) + pos = p.end.copy() + continue + + # arc + radial = p.start - p.center + if p.ccw: + tangent = np.array([-radial[1], radial[0]]) + else: + tangent = np.array([radial[1], -radial[0]]) + target_h = math.atan2(tangent[1], tangent[0]) + spin = _wrap_pi(target_h - heading) + if abs(spin) > EPS_HEADING: + out.append({"kind": "spin", "degrees": math.degrees(spin)}) + heading = target_h + sweep = _arc_signed_sweep(p) + out.append( + {"kind": "arc", "radius": float(p.radius), "degrees": math.degrees(sweep)} + ) + pos = p.end.copy() + heading = _wrap_pi(heading + sweep) + + return out + + +# ============================================================================= +# Public API +# ============================================================================= + + +def max_drift( + commands: Sequence[Dict[str, Any]], + start_pos: Sequence[float] = (0.0, 0.0), + start_heading: float = 0.0, +) -> float: + """Maximum within-stroke disjoint, in pixels, in a command list. + + A correctly-emitted DrawingCommand list should have zero drift: every + pen-down primitive's start should equal the previous primitive's end. + Useful as an assertion after any post-processing. + """ + sp = np.asarray(start_pos, dtype=float) + prims = _commands_to_primitives(commands, sp, start_heading) + max_d = 0.0 + for i in range(1, len(prims)): + prev, cur = prims[i - 1], prims[i] + if cur.kind == "line" and not cur.pen_down: + continue + max_d = max(max_d, float(np.linalg.norm(cur.start - prev.end))) + return max_d + + +@dataclass +class ConsolidationReport: + """Diagnostics returned alongside the consolidated commands.""" + + n_arc_clusters: int + arc_cluster_sizes: List[int] + n_arc_clusters_rejected: int = 0 + n_arcs_added_by_proximity: int = 0 + n_line_clusters: int = 0 + line_cluster_sizes: List[int] = field(default_factory=list) + n_line_clusters_rejected: int = 0 + n_lines_added_by_proximity: int = 0 + n_disjoints_fixed: int = 0 + max_disjoint_gap: float = 0.0 + n_proximate_endpoints_snapped: int = 0 + # Junction graph statistics + n_junctions: int = 0 + n_shared_junctions: int = 0 + shared_junction_sizes: List[int] = field(default_factory=list) + + +def _max_endpoint_snap_distance( + prims: List[_Primitive], + cluster: List[int], + cc: NDArray, + cr: float, +) -> float: + """Maximum radial distance any endpoint would have to travel to reach + the consensus circle. Used as a safety check: if even the best + consensus circle requires moving an endpoint by a lot, the arcs in + the cluster aren't really part of one circle.""" + max_d = 0.0 + for idx in cluster: + p = prims[idx] + for endpoint in (p.start, p.end): + d = float(np.linalg.norm(endpoint - cc)) + max_d = max(max_d, abs(d - cr)) + return max_d + + +def _extend_arc_clusters_with_proximate( + prims: List[_Primitive], + accepted_clusters: List[List[int]], + consensus_circles: List[Tuple[NDArray, float]], + max_endpoint_snap_rel: float, + max_endpoint_snap_abs: float, + min_radius_ratio: float = 0.4, +) -> List[List[int]]: + """Add arcs to clusters by ENDPOINT proximity, regardless of how + well their own fitted (center, radius) match. + + Rationale: an arc that is geometrically a piece of a wheel can have + a noisy individual fit (badly off radius/center) yet still have both + of its endpoints sitting close to the underlying circle. The primary + union-find pass clusters by similarity of fitted parameters and + misses these. This pass looks at every pen-down arc not yet in any + cluster and, if both its endpoints are within the snap-safety + threshold of some cluster's consensus circle, adds it to that + cluster. The "best" cluster (closest endpoints) wins ties. + + Critically, the consensus circle is NOT recomputed after extension: + the original cluster members determined a good circle, and the + newcomer's noisy fit wouldn't improve it. The newcomer just snaps + to the existing consensus. + + Guard against tiny detail arcs (small eyes, nose tips, whisker + crossings) whose endpoints happen to sit near a big circle by + requiring the candidate's own radius to be within a factor of + ``min_radius_ratio`` of the consensus radius. Without this guard, + a tiny arc forced onto a big circle ends up with two snapped + endpoints far apart on the big circle, and its preserved ``ccw`` + direction can sweep the long way around -- producing a 350° + phantom arc on the consensus circle. + """ + cluster_membership: Dict[int, int] = {} + for ci, cluster in enumerate(accepted_clusters): + for idx in cluster: + cluster_membership[idx] = ci + + for i, p in enumerate(prims): + if p.kind != "arc" or not p.pen_down or i in cluster_membership: + continue + best_ci = None + best_max_d = float("inf") + for ci, (cc, cr) in enumerate(consensus_circles): + r_ratio = min(p.radius, cr) / max(p.radius, cr) + if r_ratio < min_radius_ratio: + continue + d_start = abs(float(np.linalg.norm(p.start - cc)) - cr) + d_end = abs(float(np.linalg.norm(p.end - cc)) - cr) + max_d = max(d_start, d_end) + snap_allowed = max(max_endpoint_snap_abs, max_endpoint_snap_rel * cr) + if max_d > snap_allowed: + continue + if max_d < best_max_d: + best_max_d = max_d + best_ci = ci + if best_ci is not None: + accepted_clusters[best_ci].append(i) + cluster_membership[i] = best_ci + return accepted_clusters + + +def _snap_proximate_endpoints_to_consensus_circles( + prims: List[_Primitive], + consensus_circles: List[Tuple[NDArray, float]], + cluster_members: set, + max_snap_rel: float, + max_snap_abs: float, + jg: JunctionGraph, + min_radius_ratio: float = 0.4, +) -> int: + """For every pen-down primitive that is NOT a cluster member, look + at each endpoint and check whether it lies within snap threshold of + any consensus circle. If so, project it onto the circle and + propagate through the junction graph so any junction-mates follow. + + The junction graph already handles cases where a non-cluster + endpoint exactly coincided with a cluster-member endpoint -- those + moved during cluster snapping. This pass picks up the OTHER cases: + an endpoint that's geometrically near the consensus circle but + didn't share a junction with any cluster arc. The bike's spokes + are the canonical example: each spoke ends at the rim, but the + rim's arc-arc junctions and the rim's spoke-touching points are at + different angular positions, so they don't share junctions in the + graph. + + Same tiny-arc guard as the extension pass: for arc primitives, the + candidate's own radius must be within ``min_radius_ratio`` of the + candidate consensus radius. Lines have no inherent curvature, so + the gate is skipped for them. + + Only TERMINAL endpoints are eligible -- an endpoint is terminal + when it sits at the start or end of the primitive list, or when + the adjacent primitive is a pen-up traversal. Interior endpoints + (mid-stroke) are skipped. Without this restriction, a short + interior line segment with BOTH endpoints near a consensus circle + gets snapped on both ends, turning it into a chord of the circle + and producing a visible bump along the consolidated outline. The + bike's spokes don't have this problem because their non-rim + endpoint is at the hub, far from the consensus circle. + + Returns the count of endpoints that were snapped. + """ + n = len(prims) + n_snapped = 0 + + def is_pen_up_line(p: _Primitive) -> bool: + return p.kind == "line" and not p.pen_down + + def is_terminal(i: int, attr: str) -> bool: + if attr == "start": + return i == 0 or is_pen_up_line(prims[i - 1]) + return i == n - 1 or is_pen_up_line(prims[i + 1]) + + def maybe_snap_target(pt: NDArray, p: _Primitive) -> Optional[NDArray]: + best_ci = None + best_d = float("inf") + for ci, (cc, cr) in enumerate(consensus_circles): + # Arc-only: skip consensus circles whose curvature is + # incompatible with this primitive's own radius. + if p.kind == "arc": + r_ratio = min(p.radius, cr) / max(p.radius, cr) + if r_ratio < min_radius_ratio: + continue + d = abs(float(np.linalg.norm(pt - cc)) - cr) + snap_allowed = max(max_snap_abs, max_snap_rel * cr) + if d < snap_allowed and d < best_d: + best_d = d + best_ci = ci + if best_ci is None: + return None + cc, cr = consensus_circles[best_ci] + v = pt - cc + norm = float(np.linalg.norm(v)) + if norm < 1e-9: + return None + return cc + (cr / norm) * v + + for i, p in enumerate(prims): + if not p.pen_down or i in cluster_members: + continue + endpoint_changed = False + for attr in ("start", "end"): + if not is_terminal(i, attr): + continue + new = maybe_snap_target(getattr(p, attr), p) + if new is not None: + _move_endpoint(prims, i, attr, new, jg, cluster_members) + endpoint_changed = True + n_snapped += 1 + # Refit our own arc if its endpoints moved off-circle + if endpoint_changed and p.kind == "arc": + d_start = float(np.linalg.norm(p.start - p.center)) + d_end = float(np.linalg.norm(p.end - p.center)) + if abs(d_start - p.radius) > 0.5 or abs(d_end - p.radius) > 0.5: + _refit_arc_through_endpoints(p) + return n_snapped + + +def _extend_line_clusters_with_proximate( + prims: List[_Primitive], + accepted_clusters: List[List[int]], + consensus_lines: List[Tuple[NDArray, NDArray]], + angle_tol_rad: float, + max_line_endpoint_snap_abs: float, + min_length: float, +) -> List[List[int]]: + """Same idea as `_extend_arc_clusters_with_proximate` but for lines: + add a line to a cluster if its direction is parallel to the + cluster's consensus line (within `angle_tol_rad`) AND both endpoints + sit within `max_line_endpoint_snap_abs` perpendicular pixels of that + consensus line. + """ + cluster_membership: Dict[int, int] = {} + for ci, cluster in enumerate(accepted_clusters): + for idx in cluster: + cluster_membership[idx] = ci + + cos_tol = math.cos(angle_tol_rad) + for i, p in enumerate(prims): + if p.kind != "line" or not p.pen_down or i in cluster_membership: + continue + L = float(np.linalg.norm(p.end - p.start)) + if L < min_length: + continue + u = (p.end - p.start) / L + best_ci = None + best_max_d = float("inf") + for ci, (centroid, direction) in enumerate(consensus_lines): + if abs(float(np.dot(u, direction))) < cos_tol: + continue + perp = np.array([-direction[1], direction[0]]) + d_start = abs(float(np.dot(perp, p.start - centroid))) + d_end = abs(float(np.dot(perp, p.end - centroid))) + max_d = max(d_start, d_end) + if max_d > max_line_endpoint_snap_abs: + continue + if max_d < best_max_d: + best_max_d = max_d + best_ci = ci + if best_ci is not None: + accepted_clusters[best_ci].append(i) + cluster_membership[i] = best_ci + return accepted_clusters + + +class ConsolidateConfig(TypedDict): + center_tol_rel: float + """ + cluster two arcs if their centers are within max(rel * R_max, abs) pixels. + """ + + radius_tol_rel: float + """ + cluster two arcs if their radii differ by less than max(rel * R_max, abs). + """ + + center_tol_abs: float + """ + absolute fallback for center clustering tolerance in pixels. + """ + + radius_tol_abs: float + """ + absolute fallback for radius clustering tolerance in pixels. + """ + + max_endpoint_snap_rel: float + """ + reject an arc cluster if any endpoint would need to move more than + max(rel * R, abs) pixels to match the consensus circle. + """ + + max_endpoint_snap_abs: float + """ + absolute fallback for endpoint snap threshold in pixels when validating + arc clusters. + """ + + proximity_min_radius_ratio: float + + line_angle_tol_deg: float + """ + cluster two lines only if their direction vectors (up to sign) agree + within this many degrees. + """ + + line_offset_tol_abs: float + """ + cluster two lines only if their perpendicular offsets differ by no more + than this many pixels. + """ + + min_line_length: float + """ + lines shorter than this are excluded from line clustering because their + direction estimate is too noisy. + """ + + max_line_endpoint_snap_abs: float + """ + reject a line cluster if any endpoint would need to move more than this + many pixels to reach the consensus line. + """ + + junction_epsilon: float + """ + when building the junction graph, consider two endpoints to be linked if + they are within this many pixels of each other in the initial primitive set. + """ + + merge_arcs: bool + """ + set False to disable the arc-cluster pass. + """ + + merge_lines: bool + """ + set False to disable the line-cluster pass. + """ + + return_report: bool + """ + if True, return (commands, ConsolidationReport) instead of only commands. + """ + + +def consolidate_commands( + commands: Sequence[DrawingCommand], + start_pos: NDArray[np.float64], + start_heading: float, + config: ConsolidateConfig, +): + """Merge fragmented arcs that probably belong to one circle, merge + fragmented lines that probably belong to one infinite line, and + stitch all disjoints introduced by either operation. The output is + self-consistent (zero drift on the robot). + + Args: + commands: the DrawingCommand list to post-process. + start_pos, start_heading: the robot's initial state, used to + simulate the input commands back into geometric primitives. + + Arc-cluster tolerances (used when merge_arcs=True): + center_tol_rel/_abs: cluster two arcs if their centers are + within max(rel * R_max, abs) pixels. + radius_tol_rel/_abs: same for radius difference. + max_endpoint_snap_rel/_abs: a candidate cluster is REJECTED if + its consensus circle would require any endpoint to move + more than max(rel * R, abs) pixels. Catches the + "two unrelated arcs that just happen to look similar" case. + + Line-cluster tolerances (used when merge_lines=True): + line_angle_tol_deg: two lines cluster only if their direction + unit vectors agree (up to sign) within this many degrees. + line_offset_tol_abs: ...AND their perpendicular offset agrees + within this many pixels. + min_line_length: lines shorter than this are excluded from + clustering (their direction estimate is too noisy). + max_line_endpoint_snap_abs: a candidate line cluster is REJECTED + if any endpoint would have to move more than this many + pixels to reach the consensus line. + + Switches: + merge_arcs: set False to disable the arc-cluster pass. + merge_lines: set False to disable the line-cluster pass. + return_report: also return a ConsolidationReport. + + Returns: + consolidated_commands -- a fresh DrawingCommand list, or + (consolidated_commands, ConsolidationReport) if return_report. + """ + sp = np.asarray(start_pos, dtype=float) + prims = _commands_to_primitives(commands, sp, start_heading) + + # Build the junction graph from the *initial* primitive positions. + # All subsequent endpoint moves go through this graph so endpoints + # that originally coincided in 2D travel together. This is what + # keeps a spoke aligned with the wheel rim when the rim snaps onto + # a consensus circle. Membership is fixed at construction; later + # mutations don't change which endpoints are considered linked. + jg = JunctionGraph(prims, epsilon=config["junction_epsilon"]) + + # ---------------- arc clustering ---------------- + accepted_arc_clusters: List[List[int]] = [] + rejected_arc_clusters: List[List[int]] = [] + arc_consensus_circles: List[Tuple[NDArray, float]] = [] + consensus_arcs: Dict[int, Tuple[NDArray, float]] = {} + if config["merge_arcs"]: + candidate_arc_clusters = _cluster_arcs( + prims, + config["center_tol_rel"], + config["radius_tol_rel"], + config["center_tol_abs"], + config["radius_tol_abs"], + ) + for cluster in candidate_arc_clusters: + cc, cr = _consensus_circle(prims, cluster) + snap_d = _max_endpoint_snap_distance(prims, cluster, cc, cr) + snap_allowed = max( + config["max_endpoint_snap_abs"], config["max_endpoint_snap_rel"] * cr + ) + if snap_d <= snap_allowed: + accepted_arc_clusters.append(cluster) + arc_consensus_circles.append((cc, cr)) + else: + rejected_arc_clusters.append(cluster) + + # Extension pass: include any arc whose endpoints are already + # close to one of the accepted consensus circles, regardless of + # how noisy its own fit was. + n_arc_added_by_proximity = sum(len(c) for c in accepted_arc_clusters) + _extend_arc_clusters_with_proximate( + prims, + accepted_arc_clusters, + arc_consensus_circles, + config["max_endpoint_snap_rel"], + config["max_endpoint_snap_abs"], + min_radius_ratio=config["proximity_min_radius_ratio"], + ) + n_arc_added_by_proximity = ( + sum(len(c) for c in accepted_arc_clusters) - n_arc_added_by_proximity + ) + else: + n_arc_added_by_proximity = 0 + + # ---------------- line clustering ---------------- + accepted_line_clusters: List[List[int]] = [] + rejected_line_clusters: List[List[int]] = [] + line_consensus_geoms: List[Tuple[NDArray, NDArray]] = [] + consensus_lines: Dict[int, Tuple[NDArray, NDArray]] = {} + if config["merge_lines"]: + line_angle_tol_rad = math.radians(config["line_angle_tol_deg"]) + candidate_line_clusters = _cluster_lines( + prims, + angle_tol_rad=line_angle_tol_rad, + offset_tol_abs=config["line_offset_tol_abs"], + min_length=config["min_line_length"], + ) + for cluster in candidate_line_clusters: + centroid, direction = _consensus_line(prims, cluster) + snap_d = _max_line_endpoint_snap_distance( + prims, cluster, centroid, direction + ) + if snap_d <= config["max_line_endpoint_snap_abs"]: + accepted_line_clusters.append(cluster) + line_consensus_geoms.append((centroid, direction)) + else: + rejected_line_clusters.append(cluster) + + # Extension pass for lines. + n_line_added_by_proximity = sum(len(c) for c in accepted_line_clusters) + _extend_line_clusters_with_proximate( + prims, + accepted_line_clusters, + line_consensus_geoms, + angle_tol_rad=line_angle_tol_rad, + max_line_endpoint_snap_abs=config["max_line_endpoint_snap_abs"], + min_length=config["min_line_length"], + ) + n_line_added_by_proximity = ( + sum(len(c) for c in accepted_line_clusters) - n_line_added_by_proximity + ) + else: + n_line_added_by_proximity = 0 + + # Pre-compute the set of all cluster members across both kinds. + # Propagation must not refit cluster member arcs (that would + # overwrite their freshly-set consensus center/radius). + all_cluster_members: set = set() + for c in accepted_arc_clusters: + all_cluster_members.update(c) + for c in accepted_line_clusters: + all_cluster_members.update(c) + + # Now do the actual snapping, with junction-graph propagation. + if config["merge_arcs"]: + for cluster, (cc, cr) in zip(accepted_arc_clusters, arc_consensus_circles): + for idx in cluster: + _snap_arc_to_consensus( + prims[idx], + idx, + cc, + cr, + prims, + jg, + all_cluster_members, + ) + consensus_arcs[idx] = (cc, cr) + if config["merge_lines"]: + for cluster, (centroid, direction) in zip( + accepted_line_clusters, + line_consensus_geoms, + ): + for idx in cluster: + _snap_line_to_consensus( + prims[idx], + idx, + centroid, + direction, + prims, + jg, + all_cluster_members, + ) + consensus_lines[idx] = (centroid, direction) + + # Proximate-endpoint pass: catches non-cluster endpoints (typically + # belonging to spokes or other strokes meeting a consolidated curve) + # that are near a consensus circle but didn't share a junction with + # any cluster member. Without this, those endpoints stay at their + # original positions while the rim moves, producing visible + # spoke-overshoots-rim artifacts. + n_endpoints_snapped = 0 + if config["merge_arcs"] and arc_consensus_circles: + n_endpoints_snapped = _snap_proximate_endpoints_to_consensus_circles( + prims, + arc_consensus_circles, + all_cluster_members, + config["max_endpoint_snap_rel"], + config["max_endpoint_snap_abs"], + jg, + min_radius_ratio=config["proximity_min_radius_ratio"], + ) + + # ---------------- pre-stitch diagnostics ---------------- + max_gap = 0.0 + n_disjoints = 0 + for i in range(1, len(prims)): + prev, cur = prims[i - 1], prims[i] + if cur.kind == "line" and not cur.pen_down: + continue + g = float(np.linalg.norm(cur.start - prev.end)) + if g > 1e-3: + n_disjoints += 1 + if g > max_gap: + max_gap = g + + _stitch(prims, consensus_arcs, consensus_lines, jg) + + # Final reprojection pass: cross-cluster propagation through the + # junction graph could have pulled a clustered arc's endpoint + # slightly off its consensus circle. Re-snap clustered arc/line + # endpoints back onto their own consensus geometry. Don't propagate + # this final fix (we don't want to set off another wave of moves); + # just clamp each cluster member onto its consensus. + for idx, (cc, cr) in consensus_arcs.items(): + p = prims[idx] + p.start = _project_onto_circle(p.start, cc, cr) + p.end = _project_onto_circle(p.end, cc, cr) + for idx, (centroid, direction) in consensus_lines.items(): + p = prims[idx] + p.start = _project_onto_line(p.start, centroid, direction) + p.end = _project_onto_line(p.end, centroid, direction) + + out = _primitives_to_commands(prims, sp, start_heading) + + if config["return_report"]: + report = ConsolidationReport( + n_arc_clusters=len(accepted_arc_clusters), + arc_cluster_sizes=[len(c) for c in accepted_arc_clusters], + n_arc_clusters_rejected=len(rejected_arc_clusters), + n_arcs_added_by_proximity=n_arc_added_by_proximity, + n_line_clusters=len(accepted_line_clusters), + line_cluster_sizes=[len(c) for c in accepted_line_clusters], + n_line_clusters_rejected=len(rejected_line_clusters), + n_lines_added_by_proximity=n_line_added_by_proximity, + n_disjoints_fixed=n_disjoints, + max_disjoint_gap=max_gap, + n_proximate_endpoints_snapped=n_endpoints_snapped, + n_junctions=jg.n_junctions(), + n_shared_junctions=jg.n_shared_junctions(), + shared_junction_sizes=jg.shared_junction_sizes(), + ) + return out, report + return out, None diff --git a/arc_line_vectorization_suede/vectorize/high_geometry/consolidate_v1.py b/arc_line_vectorization_suede/vectorize/high_geometry/consolidate_v1.py new file mode 100644 index 0000000..6ca29de --- /dev/null +++ b/arc_line_vectorization_suede/vectorize/high_geometry/consolidate_v1.py @@ -0,0 +1,783 @@ +"""Post-pass over a finished DrawingCommand list. + +The vectorizer fits each segment as an independent line or arc, so a +single underlying circle in the input (a wheel, an eye, a head outline) +that gets fragmented at junctions ends up as several arcs with slightly +different (center, radius) parameters. Visually the result looks broken +even though each piece is a good fit on its own. + +This module does a pure post-pass on the emitted commands: + + 1. Reverse-simulate the command sequence to recover geometric + primitives (each with start, end, and arc params if applicable). + 2. Cluster pen-down arcs whose (center, radius) are similar. + 3. For every cluster of size >= 2, compute a consensus circle + (sweep-weighted average) and snap every arc in the cluster onto + that circle: its center and radius are replaced, and its endpoints + are projected radially onto the new circle. + 4. Snapping moves endpoints, which creates disjoints with neighbouring + primitives. Walk the primitive list and stitch every gap so that + prim[i+1].start == prim[i].end exactly. For neighbouring arcs that + are not in any cluster, refit center and radius so the moved + endpoint stays on a circle (preserving the original bend + direction). + 5. Re-emit a fresh command sequence from the modified primitives. + +The output is self-consistent: every arc primitive satisfies +|start - center| = |end - center| = radius, and adjacent primitives +share their connecting endpoint exactly. That is what guarantees the +robot doesn't drift -- the emitter computes spin/distance/sweep from +the primitive's geometry, and any inconsistency would manifest as +accumulating offset along the chain. +""" + +from __future__ import annotations +import math +from dataclasses import dataclass, field +from typing import Any, Dict, List, Optional, Sequence, Tuple, TypedDict, cast + +import numpy as np +from numpy.typing import NDArray + +from commands import DrawingCommand, SpinCommand, LineCommand, ArcCommand + +# ============================================================================= +# Internal geometric representation +# ============================================================================= + + +@dataclass +class _Primitive: + """One geometric primitive between two robot states. + + `pen_down` is False only for inter-stroke pen-up traversals (which + have free endpoints and therefore aren't constrained by the stitcher + on their END side -- they just go wherever the next primitive starts). + Arc fields are populated only when kind == "arc". + """ + + kind: str # "line" | "arc" + start: NDArray[np.float64] + end: NDArray[np.float64] + pen_down: bool = True + # arc-only: + center: Optional[NDArray[np.float64]] = None + radius: float = 0.0 + ccw: bool = True + + +def _wrap_pi(a: float) -> float: + return ((a + math.pi) % (2 * math.pi)) - math.pi + + +# ============================================================================= +# Stage 1: simulate commands -> primitives +# ============================================================================= + + +def _commands_to_primitives( + commands: Sequence[DrawingCommand], + start_pos: NDArray[np.float64], + start_heading: float, +) -> List[_Primitive]: + """Walk the command list while tracking (pos, heading); emit one + primitive per line/arc command (spins are absorbed into the heading). + """ + pos = np.asarray(start_pos, dtype=float).copy() + heading = float(start_heading) + prims: List[_Primitive] = [] + for c in commands: + kind = c["kind"] + if kind == "spin": + spin = cast(SpinCommand, c) + heading = _wrap_pi(heading + math.radians(spin["degrees"])) + elif kind == "line": + line = cast(LineCommand, c) + dist = float(line["distance"]) + new_pos = pos + dist * np.array([math.cos(heading), math.sin(heading)]) + prims.append( + _Primitive( + kind="line", + start=pos.copy(), + end=new_pos.copy(), + pen_down=bool(line.get("penDown", True)), + ) + ) + pos = new_pos + elif kind == "arc": + arc = cast(ArcCommand, c) + r = float(arc["radius"]) + sweep = math.radians(float(arc["degrees"])) + ccw = sweep > 0 + # Center is perpendicular to heading at distance r; left for + # CCW, right for CW. + if ccw: + perp = np.array([-math.sin(heading), math.cos(heading)]) + else: + perp = np.array([math.sin(heading), -math.cos(heading)]) + center = pos + r * perp + # Rotate the radial vector around the center by `sweep`. + radial = pos - center + cs, sn = math.cos(sweep), math.sin(sweep) + new_radial = np.array( + [ + radial[0] * cs - radial[1] * sn, + radial[0] * sn + radial[1] * cs, + ] + ) + new_pos = center + new_radial + prims.append( + _Primitive( + kind="arc", + start=pos.copy(), + end=new_pos.copy(), + pen_down=True, + center=center, + radius=r, + ccw=ccw, + ) + ) + pos = new_pos + heading = _wrap_pi(heading + sweep) + else: + raise ValueError(f"unknown command kind: {kind!r}") + return prims + + +# ============================================================================= +# Stage 2: cluster arcs by shared underlying circle +# ============================================================================= + + +def _cluster_arcs( + prims: List[_Primitive], + center_tol_rel: float, + radius_tol_rel: float, + center_tol_abs: float, + radius_tol_abs: float, +) -> List[List[int]]: + """Union-find arcs with similar (center, radius). Tolerances are + 'either / or': two arcs are linked if their centers AND radii are + close, where close = max(rel * R, abs). + Returns clusters as lists of primitive indices (only clusters of + size >= 2 are returned).""" + arc_idxs = [i for i, p in enumerate(prims) if p.kind == "arc" and p.pen_down] + n = len(arc_idxs) + parent = list(range(n)) + + def find(x: int) -> int: + while parent[x] != x: + parent[x] = parent[parent[x]] + x = parent[x] + return x + + def union(a: int, b: int) -> None: + ra, rb = find(a), find(b) + if ra != rb: + parent[ra] = rb + + for ii in range(n): + pi = prims[arc_idxs[ii]] + if pi.center is None: + continue + for jj in range(ii + 1, n): + pj = prims[arc_idxs[jj]] + if pj.center is None: + continue + r_max = max(pi.radius, pj.radius) + ctol = max(center_tol_rel * r_max, center_tol_abs) + rtol = max(radius_tol_rel * r_max, radius_tol_abs) + cdist = float(np.linalg.norm(pi.center - pj.center)) + rdiff = abs(pi.radius - pj.radius) + if cdist < ctol and rdiff < rtol: + union(ii, jj) + + bucket: Dict[int, List[int]] = {} + for ii in range(n): + bucket.setdefault(find(ii), []).append(arc_idxs[ii]) + return [g for g in bucket.values() if len(g) >= 2] + + +def _arc_signed_sweep(p: _Primitive) -> float: + """Sweep angle in radians (positive for CCW, negative for CW), as + you would compute from the primitive's start, end and center.""" + if p.center is None: + return 0.0 + a_s = math.atan2(p.start[1] - p.center[1], p.start[0] - p.center[0]) + a_e = math.atan2(p.end[1] - p.center[1], p.end[0] - p.center[0]) + sweep = a_e - a_s + if p.ccw and sweep <= 0: + sweep += 2 * math.pi + elif not p.ccw and sweep >= 0: + sweep -= 2 * math.pi + return sweep + + +def _consensus_circle( + prims: List[_Primitive], cluster: List[int] +) -> Tuple[NDArray, float]: + """Sweep-weighted average of (center, radius). Longer arcs vote more + because they carry more information about the underlying circle.""" + weights, centers, radii = [], [], [] + for idx in cluster: + p = prims[idx] + # Minimum weight floors the contribution of vanishingly short arcs + # without dropping them entirely. + w = max(abs(_arc_signed_sweep(p)), 0.05) + weights.append(w) + centers.append(p.center) + radii.append(p.radius) + weights = np.asarray(weights) + weights = weights / weights.sum() + cc = np.zeros(2) + for w, c in zip(weights, centers): + cc = cc + w * c + cr = float(sum(w * r for w, r in zip(weights, radii))) + return cc, cr + + +# ============================================================================= +# Stage 3: snap each clustered arc to its consensus circle +# ============================================================================= + + +def _project_onto_circle(p: NDArray, center: NDArray, radius: float) -> NDArray: + v = p - center + d = float(np.linalg.norm(v)) + if d < 1e-9: + return center + np.array([radius, 0.0]) + return center + (radius / d) * v + + +def _snap_arc_to_consensus(p: _Primitive, cc: NDArray, cr: float) -> None: + """Replace the arc's center and radius with the cluster consensus, + and project both endpoints radially onto the new circle. The arc's + ccw direction is preserved.""" + p.center = cc.copy() + p.radius = cr + p.start = _project_onto_circle(p.start, cc, cr) + p.end = _project_onto_circle(p.end, cc, cr) + + +# ============================================================================= +# Stage 2b/3b: cluster collinear lines and snap them to a shared consensus +# ============================================================================= + + +def _line_geometry(p: _Primitive) -> Tuple[NDArray, float]: + """Return the line's unit direction and length. Caller is expected to + have filtered out zero-length lines already.""" + v = p.end - p.start + L = float(np.linalg.norm(v)) + return v / L, L + + +def _cluster_lines( + prims: List[_Primitive], + angle_tol_rad: float, + offset_tol_abs: float, + min_length: float, +) -> List[List[int]]: + """Group pen-down lines that are collinear: same direction within + `angle_tol_rad` (sign-ambiguous, since a line has no preferred + orientation), and same perpendicular offset within `offset_tol_abs` + pixels. Lines shorter than `min_length` are excluded because their + direction estimate is too noisy.""" + line_idxs: List[int] = [] + dirs: List[NDArray] = [] + for i, p in enumerate(prims): + if p.kind != "line" or not p.pen_down: + continue + L = float(np.linalg.norm(p.end - p.start)) + if L < min_length: + continue + line_idxs.append(i) + dirs.append((p.end - p.start) / L) + + n = len(line_idxs) + parent = list(range(n)) + + def find(x: int) -> int: + while parent[x] != x: + parent[x] = parent[parent[x]] + x = parent[x] + return x + + def union(a: int, b: int) -> None: + ra, rb = find(a), find(b) + if ra != rb: + parent[ra] = rb + + cos_tol = math.cos(angle_tol_rad) + for ii in range(n): + pi = prims[line_idxs[ii]] + di = dirs[ii] + for jj in range(ii + 1, n): + pj = prims[line_idxs[jj]] + dj = dirs[jj] + # Direction parallel? Use |dot| because a line has no orientation. + if abs(float(np.dot(di, dj))) < cos_tol: + continue + # Perpendicular offset: distance from one line's start to the + # other's infinite line. + perp_j = np.array([-dj[1], dj[0]]) + offset = abs(float(np.dot(perp_j, pi.start - pj.start))) + if offset < offset_tol_abs: + union(ii, jj) + + bucket: Dict[int, List[int]] = {} + for ii in range(n): + bucket.setdefault(find(ii), []).append(line_idxs[ii]) + return [g for g in bucket.values() if len(g) >= 2] + + +def _consensus_line( + prims: List[_Primitive], + cluster: List[int], +) -> Tuple[NDArray, NDArray]: + """Total-least-squares line through all endpoints of the clustered + lines. Returns (centroid, unit_direction); the infinite consensus + line is the set of points centroid + t * direction for any t in R.""" + pts = [] + for idx in cluster: + p = prims[idx] + pts.append(p.start) + pts.append(p.end) + arr = np.asarray(pts) + centroid = arr.mean(axis=0) + centered = arr - centroid + _, _, Vt = np.linalg.svd(centered, full_matrices=False) + direction = Vt[0] + return centroid, direction + + +def _project_onto_line(p: NDArray, centroid: NDArray, direction: NDArray) -> NDArray: + return centroid + float(np.dot(p - centroid, direction)) * direction + + +def _snap_line_to_consensus( + p: _Primitive, + centroid: NDArray, + direction: NDArray, +) -> None: + """Project both endpoints of a line primitive onto the consensus + infinite line (perpendicular projection).""" + p.start = _project_onto_line(p.start, centroid, direction) + p.end = _project_onto_line(p.end, centroid, direction) + + +def _max_line_endpoint_snap_distance( + prims: List[_Primitive], + cluster: List[int], + centroid: NDArray, + direction: NDArray, +) -> float: + """Max perpendicular distance any endpoint must travel to reach the + consensus line. Large values indicate the cluster is bogus.""" + perp = np.array([-direction[1], direction[0]]) + max_d = 0.0 + for idx in cluster: + p = prims[idx] + for endpoint in (p.start, p.end): + d = abs(float(np.dot(perp, endpoint - centroid))) + if d > max_d: + max_d = d + return max_d + + +# ============================================================================= +# Stage 4: stitching (refit non-clustered arcs whose endpoints moved) +# ============================================================================= + + +def _refit_arc_through_endpoints(p: _Primitive) -> None: + """The arc's start and end no longer satisfy + |x - center| = radius. Find the unique arc that PASSES THROUGH the + current start and end with center on the perpendicular bisector, + closest to the original center. Preserves the original bending side + (and therefore the ccw flag).""" + if p.center is None: + return + chord = p.end - p.start + chord_len = float(np.linalg.norm(chord)) + if chord_len < 1e-9: + return # degenerate + midpoint = 0.5 * (p.start + p.end) + # perpendicular to chord (unit) + perp = np.array([-chord[1], chord[0]]) / chord_len + # Match the original bending side + if float(np.dot(p.center - midpoint, perp)) < 0: + perp = -perp + # Project the original center onto the perpendicular line + t = float(np.dot(p.center - midpoint, perp)) + new_center = midpoint + t * perp + new_radius = math.sqrt(t * t + (chord_len / 2.0) ** 2) + p.center = new_center + p.radius = new_radius + # ccw is preserved + + +def _stitch( + prims: List[_Primitive], + consensus_arcs: Dict[int, Tuple[NDArray, float]], + consensus_lines: Dict[int, Tuple[NDArray, NDArray]], +) -> None: + """Walk the primitive list in order. For each adjacent pen-down pair + whose connecting endpoints disagree, reconcile while preserving each + side's consensus geometry where possible. + + Resolution rules: + - Pen-up "line" primitives are unconstrained on both ends; they + just snap to whatever the neighbours dictate. + - If exactly one of (prev, cur) is on a consensus geometry (arc or + line), the other adapts: the moving endpoint is set to match the + clustered side; non-clustered arc neighbours get refit through + their new endpoints, non-clustered line neighbours just store + the new endpoint. + - If neither side is clustered (only happens if a non-clustered arc + was already disturbed indirectly), take the midpoint and refit. + - If both sides are clustered (e.g. a circle meeting a line at + a junction), take the midpoint, project each side onto its own + consensus, average once more. Sub-pixel residual is accepted. + """ + EPS = 1e-3 + + def project_to_own(idx: int, attr: str) -> None: + """Project the (start|end) of prims[idx] onto its consensus + geometry, if any. Otherwise no-op.""" + p = prims[idx] + pt = getattr(p, attr) + if idx in consensus_arcs: + cc, cr = consensus_arcs[idx] + v = pt - cc + d = float(np.linalg.norm(v)) + if d > 1e-9: + setattr(p, attr, cc + (cr / d) * v) + elif idx in consensus_lines: + centroid, direction = consensus_lines[idx] + setattr(p, attr, _project_onto_line(pt, centroid, direction)) + + def is_clustered(idx: int) -> bool: + return idx in consensus_arcs or idx in consensus_lines + + for i in range(1, len(prims)): + prev = prims[i - 1] + cur = prims[i] + + if cur.kind == "line" and not cur.pen_down: + cur.start = prev.end.copy() + continue + + gap = float(np.linalg.norm(cur.start - prev.end)) + if gap < EPS: + cur.start = prev.end.copy() + continue + + prev_clustered = is_clustered(i - 1) + cur_clustered = is_clustered(i) + + if prev_clustered and not cur_clustered: + cur.start = prev.end.copy() + if cur.kind == "arc": + _refit_arc_through_endpoints(cur) + elif cur_clustered and not prev_clustered: + prev.end = cur.start.copy() + if prev.kind == "arc": + _refit_arc_through_endpoints(prev) + elif prev_clustered and cur_clustered: + mid = 0.5 * (prev.end + cur.start) + prev.end = mid.copy() + cur.start = mid.copy() + project_to_own(i - 1, "end") + project_to_own(i, "start") + mid = 0.5 * (prev.end + cur.start) + prev.end = mid.copy() + cur.start = mid.copy() + else: + mid = 0.5 * (prev.end + cur.start) + prev.end = mid.copy() + cur.start = mid.copy() + if prev.kind == "arc": + _refit_arc_through_endpoints(prev) + if cur.kind == "arc": + _refit_arc_through_endpoints(cur) + + +# ============================================================================= +# Stage 5: primitives -> commands (re-emit) +# ============================================================================= + + +def _primitives_to_commands( + prims: List[_Primitive], + start_pos: NDArray, + start_heading: float, +) -> List[DrawingCommand]: + """Walk the (now self-consistent) primitive list and emit + spin / line / arc commands. Mirrors the existing `strokes_to_commands` + behaviour but drives off the primitive list rather than strokes.""" + pos = np.asarray(start_pos, dtype=float).copy() + heading = float(start_heading) + out: List[DrawingCommand] = [] + EPS_HEADING = 1e-4 + EPS_DIST = 1e-6 + + for p in prims: + if p.kind == "line": + delta = p.end - pos + dist = float(np.linalg.norm(delta)) + if dist < EPS_DIST: + continue + target_h = math.atan2(delta[1], delta[0]) + spin = _wrap_pi(target_h - heading) + if abs(spin) > EPS_HEADING: + out.append({"kind": "spin", "degrees": math.degrees(spin)}) + heading = target_h + out.append({"kind": "line", "distance": dist, "penDown": p.pen_down}) + pos = p.end.copy() + continue + + if p.center is None: + continue + # arc + radial = p.start - p.center + if p.ccw: + tangent = np.array([-radial[1], radial[0]]) + else: + tangent = np.array([radial[1], -radial[0]]) + target_h = math.atan2(tangent[1], tangent[0]) + spin = _wrap_pi(target_h - heading) + if abs(spin) > EPS_HEADING: + out.append({"kind": "spin", "degrees": math.degrees(spin)}) + heading = target_h + sweep = _arc_signed_sweep(p) + out.append( + {"kind": "arc", "radius": float(p.radius), "degrees": math.degrees(sweep)} + ) + pos = p.end.copy() + heading = _wrap_pi(heading + sweep) + + return out + + +# ============================================================================= +# Public API +# ============================================================================= + + +def max_drift( + commands: Sequence[DrawingCommand], + start_pos: Sequence[float] = (0.0, 0.0), + start_heading: float = 0.0, +) -> float: + """Maximum within-stroke disjoint, in pixels, in a command list. + + A correctly-emitted DrawingCommand list should have zero drift: every + pen-down primitive's start should equal the previous primitive's end. + Useful as an assertion after any post-processing. + """ + sp = np.asarray(start_pos, dtype=float) + prims = _commands_to_primitives(commands, sp, start_heading) + max_d = 0.0 + for i in range(1, len(prims)): + prev, cur = prims[i - 1], prims[i] + if cur.kind == "line" and not cur.pen_down: + continue + max_d = max(max_d, float(np.linalg.norm(cur.start - prev.end))) + return max_d + + +@dataclass +class ConsolidationReport: + """Diagnostics returned alongside the consolidated commands.""" + + n_arc_clusters: int + arc_cluster_sizes: List[int] + n_arc_clusters_rejected: int = 0 + n_line_clusters: int = 0 + line_cluster_sizes: List[int] = field(default_factory=list) + n_line_clusters_rejected: int = 0 + n_disjoints_fixed: int = 0 + max_disjoint_gap: float = 0.0 + + +def _max_endpoint_snap_distance( + prims: List[_Primitive], + cluster: List[int], + cc: NDArray, + cr: float, +) -> float: + """Maximum radial distance any endpoint would have to travel to reach + the consensus circle. Used as a safety check: if even the best + consensus circle requires moving an endpoint by a lot, the arcs in + the cluster aren't really part of one circle.""" + max_d = 0.0 + for idx in cluster: + p = prims[idx] + for endpoint in (p.start, p.end): + d = float(np.linalg.norm(endpoint - cc)) + max_d = max(max_d, abs(d - cr)) + return max_d + + # start_pos: Sequence[float] = (0.0, 0.0), + # start_heading: float = 0.0, + # # arc-cluster knobs + # center_tol_rel: float = 0.25, + # radius_tol_rel: float = 0.25, + # center_tol_abs: float = 3.0, + # radius_tol_abs: float = 3.0, + # max_endpoint_snap_rel: float = 0.15, + # max_endpoint_snap_abs: float = 6.0, + # # line-cluster knobs + # line_angle_tol_deg: float = 6.0, + # line_offset_tol_abs: float = 5.0, + # min_line_length: float = 5.0, + # max_line_endpoint_snap_abs: float = 5.0, + # # behaviour switches + # merge_arcs: bool = True, + # merge_lines: bool = True, + # return_report: bool = False, + + +class ConsolidateConfig(TypedDict): + center_tol_rel: float + radius_tol_rel: float + center_tol_abs: float + radius_tol_abs: float + max_endpoint_snap_rel: float + max_endpoint_snap_abs: float + line_angle_tol_deg: float + line_offset_tol_abs: float + min_line_length: float + max_line_endpoint_snap_abs: float + merge_arcs: bool + merge_lines: bool + return_report: bool + + +def consolidate_commands( + commands: Sequence[DrawingCommand], + start_pos: NDArray[np.float64], + start_heading: float, + config: ConsolidateConfig, +): + """Merge fragmented arcs that probably belong to one circle, merge + fragmented lines that probably belong to one infinite line, and + stitch all disjoints introduced by either operation. The output is + self-consistent (zero drift on the robot). + + Args: + commands: the DrawingCommand list to post-process. + start_pos, start_heading: the robot's initial state, used to + simulate the input commands back into geometric primitives. + + Arc-cluster tolerances (used when merge_arcs=True): + center_tol_rel/_abs: cluster two arcs if their centers are + within max(rel * R_max, abs) pixels. + radius_tol_rel/_abs: same for radius difference. + max_endpoint_snap_rel/_abs: a candidate cluster is REJECTED if + its consensus circle would require any endpoint to move + more than max(rel * R, abs) pixels. Catches the + "two unrelated arcs that just happen to look similar" case. + + Line-cluster tolerances (used when merge_lines=True): + line_angle_tol_deg: two lines cluster only if their direction + unit vectors agree (up to sign) within this many degrees. + line_offset_tol_abs: ...AND their perpendicular offset agrees + within this many pixels. + min_line_length: lines shorter than this are excluded from + clustering (their direction estimate is too noisy). + max_line_endpoint_snap_abs: a candidate line cluster is REJECTED + if any endpoint would have to move more than this many + pixels to reach the consensus line. + + Switches: + merge_arcs: set False to disable the arc-cluster pass. + merge_lines: set False to disable the line-cluster pass. + return_report: also return a ConsolidationReport. + + Returns: + consolidated_commands -- a fresh DrawingCommand list, or + (consolidated_commands, ConsolidationReport) if return_report. + """ + prims = _commands_to_primitives(commands, start_pos, start_heading) + + # ---------------- arc clustering ---------------- + accepted_arc_clusters: List[List[int]] = [] + rejected_arc_clusters: List[List[int]] = [] + consensus_arcs: Dict[int, Tuple[NDArray, float]] = {} + if config["merge_arcs"]: + candidate_arc_clusters = _cluster_arcs( + prims, + config["center_tol_rel"], + config["radius_tol_rel"], + config["center_tol_abs"], + config["radius_tol_abs"], + ) + for cluster in candidate_arc_clusters: + cc, cr = _consensus_circle(prims, cluster) + snap_d = _max_endpoint_snap_distance(prims, cluster, cc, cr) + snap_allowed = max( + config["max_endpoint_snap_abs"], config["max_endpoint_snap_rel"] * cr + ) + if snap_d <= snap_allowed: + accepted_arc_clusters.append(cluster) + for idx in cluster: + _snap_arc_to_consensus(prims[idx], cc, cr) + consensus_arcs[idx] = (cc, cr) + else: + rejected_arc_clusters.append(cluster) + + # ---------------- line clustering ---------------- + accepted_line_clusters: List[List[int]] = [] + rejected_line_clusters: List[List[int]] = [] + consensus_lines: Dict[int, Tuple[NDArray, NDArray]] = {} + if config["merge_lines"]: + candidate_line_clusters = _cluster_lines( + prims, + angle_tol_rad=math.radians(config["line_angle_tol_deg"]), + offset_tol_abs=config["line_offset_tol_abs"], + min_length=config["min_line_length"], + ) + for cluster in candidate_line_clusters: + centroid, direction = _consensus_line(prims, cluster) + snap_d = _max_line_endpoint_snap_distance( + prims, cluster, centroid, direction + ) + if snap_d <= config["max_line_endpoint_snap_abs"]: + accepted_line_clusters.append(cluster) + for idx in cluster: + _snap_line_to_consensus(prims[idx], centroid, direction) + consensus_lines[idx] = (centroid, direction) + else: + rejected_line_clusters.append(cluster) + + # ---------------- pre-stitch diagnostics ---------------- + max_gap = 0.0 + n_disjoints = 0 + for i in range(1, len(prims)): + prev, cur = prims[i - 1], prims[i] + if cur.kind == "line" and not cur.pen_down: + continue + g = float(np.linalg.norm(cur.start - prev.end)) + if g > 1e-3: + n_disjoints += 1 + if g > max_gap: + max_gap = g + + _stitch(prims, consensus_arcs, consensus_lines) + out = _primitives_to_commands(prims, start_pos, start_heading) + + if config["return_report"]: + report = ConsolidationReport( + n_arc_clusters=len(accepted_arc_clusters), + arc_cluster_sizes=[len(c) for c in accepted_arc_clusters], + n_arc_clusters_rejected=len(rejected_arc_clusters), + n_line_clusters=len(accepted_line_clusters), + line_cluster_sizes=[len(c) for c in accepted_line_clusters], + n_line_clusters_rejected=len(rejected_line_clusters), + n_disjoints_fixed=n_disjoints, + max_disjoint_gap=max_gap, + ) + return out, report + return out, None diff --git a/arc_line_vectorization_suede/vectorize/high_geometry/junction_graph.py b/arc_line_vectorization_suede/vectorize/high_geometry/junction_graph.py new file mode 100644 index 0000000..9a0777f --- /dev/null +++ b/arc_line_vectorization_suede/vectorize/high_geometry/junction_graph.py @@ -0,0 +1,162 @@ +"""Track 2D spatial coincidence of primitive endpoints, so that any +operation that moves one endpoint can propagate the move to all +endpoints that originally coincided with it. + +The motivating problem is "mismatched junctions across strokes". In +the input drawing, two strokes (say, a wheel rim and a spoke) often +share a junction at a point. After segmentation/fusion they end up in +different strokes -- the rim is one stroke and each spoke is another, +with pen-up traversals in between. The shared junction now exists as +TWO independent endpoints, one in each stroke. Later, when the rim +gets snapped onto a consensus circle by the consolidator, ONLY the +rim's endpoint moves; the spoke's endpoint stays where it was drawn, +producing a visible mismatch (spoke overshoots or falls short of the +new rim). + +The junction graph clusters all primitive endpoints by 2D proximity +once at the start. Subsequent endpoint moves go through ``move()``, +which propagates to every endpoint sharing the same junction. The +cluster membership is built from ORIGINAL positions and is immutable +afterward -- endpoints that originally coincided always travel +together, even after they've drifted through several operations. + +Typical usage during consolidation:: + + jg = JunctionGraph(prims, epsilon=3.0) + moved_others = jg.move(prims, idx, "start", new_pos) + for (oi, _) in moved_others: + if prims[oi].kind == "arc": + refit_arc_through_endpoints(prims[oi]) + +The graph itself does not know how to refit arcs after their endpoints +move; the caller is responsible for that. This keeps the graph generic +and its responsibilities focused. +""" + +from __future__ import annotations +from typing import Any, Dict, List, Sequence, Tuple + +import numpy as np +from numpy.typing import NDArray + + +class JunctionGraph: + """Spatial-coincidence map for primitive endpoints. + + Endpoints are identified by ``(prim_idx, "start"|"end")`` tuples. + Two endpoints share a junction iff their L2 distance was below + ``epsilon`` at construction time. Junction membership is fixed + after construction; it does not depend on current positions. + """ + + def __init__(self, prims: Sequence[Any], epsilon: float = 3.0): + self._epsilon = float(epsilon) + + endpoints: List[Tuple[int, str, NDArray]] = [] + for i, p in enumerate(prims): + endpoints.append((i, "start", np.asarray(p.start, dtype=float))) + endpoints.append((i, "end", np.asarray(p.end, dtype=float))) + + n = len(endpoints) + parent = list(range(n)) + + def find(x: int) -> int: + while parent[x] != x: + parent[x] = parent[parent[x]] + x = parent[x] + return x + + def union(a: int, b: int) -> None: + ra, rb = find(a), find(b) + if ra != rb: + parent[ra] = rb + + # O(n^2) pairwise distance check. n is twice the number of + # primitives -- usually a few hundred, never thousands. A spatial + # hash would speed this up but isn't worth the code. + eps_sq = self._epsilon * self._epsilon + for ii in range(n): + pi = endpoints[ii][2] + pi_x, pi_y = float(pi[0]), float(pi[1]) + for jj in range(ii + 1, n): + pj = endpoints[jj][2] + dx = pi_x - float(pj[0]) + dy = pi_y - float(pj[1]) + if dx * dx + dy * dy < eps_sq: + union(ii, jj) + + # Compress find paths and build the (junction_id -> members) and + # (endpoint -> junction_id) tables. + self._junction_of: Dict[Tuple[int, str], int] = {} + self._members: Dict[int, List[Tuple[int, str]]] = {} + for ii, (i, attr, _) in enumerate(endpoints): + root = find(ii) + self._junction_of[(i, attr)] = root + self._members.setdefault(root, []).append((i, attr)) + + @property + def epsilon(self) -> float: + return self._epsilon + + def n_junctions(self) -> int: + """Total number of distinct junctions, shared or not.""" + return len(self._members) + + def n_shared_junctions(self) -> int: + """Number of junctions where two or more endpoints meet. This is + the count of "real" junctions in the drawing -- single-membership + junctions are just stroke endpoints that don't touch anything.""" + return sum(1 for v in self._members.values() if len(v) >= 2) + + def shared_junction_sizes(self) -> List[int]: + """Member counts of every shared junction, sorted descending. A + size-3 entry means three endpoints meet at one point in the + original drawing.""" + sizes = [len(v) for v in self._members.values() if len(v) >= 2] + sizes.sort(reverse=True) + return sizes + + def members(self, prim_idx: int, attr: str) -> List[Tuple[int, str]]: + """All endpoints sharing the junction with the given one, + including itself. If the queried endpoint is not in the graph + (which shouldn't happen for a graph built from the same prims), + returns just the queried endpoint.""" + j = self._junction_of.get((prim_idx, attr)) + if j is None: + return [(prim_idx, attr)] + return list(self._members[j]) + + def has_neighbours(self, prim_idx: int, attr: str) -> bool: + """True iff this endpoint shares its junction with at least one + other primitive's endpoint.""" + j = self._junction_of.get((prim_idx, attr)) + return j is not None and len(self._members[j]) >= 2 + + def move( + self, + prims: Sequence[Any], + prim_idx: int, + attr: str, + new_pos: NDArray, + ) -> List[Tuple[int, str]]: + """Set the position of ``(prim_idx, attr)`` and every endpoint + sharing its junction to ``new_pos``. Returns the list of + ``(other_idx, other_attr)`` that were ALSO moved -- i.e., the + propagation set, not including the requested endpoint itself. + + After this call, every junction-mate of ``(prim_idx, attr)`` has + ``new_pos`` as its position. The caller is expected to refit any + arc primitives in the propagation set whose center+radius are no + longer consistent with their new endpoints. + """ + np_pos = np.asarray(new_pos, dtype=float) + moved_others: List[Tuple[int, str]] = [] + for other_idx, other_attr in self.members(prim_idx, attr): + p = prims[other_idx] + if other_attr == "start": + p.start = np_pos.copy() + else: + p.end = np_pos.copy() + if (other_idx, other_attr) != (prim_idx, attr): + moved_others.append((other_idx, other_attr)) + return moved_others diff --git a/arc_line_vectorization_suede/vectorize/low_geometry/__init__.py b/arc_line_vectorization_suede/vectorize/low_geometry/__init__.py new file mode 100644 index 0000000..ec339f7 --- /dev/null +++ b/arc_line_vectorization_suede/vectorize/low_geometry/__init__.py @@ -0,0 +1,427 @@ +"""Vectorize: convert a StrokeGraph into a sequence of robot drawing +commands. + +Top-level pipeline (see ``solve.py`` for orchestration details): + +1. Chain subdivision per polyline (``fit_polyline``). +2. Flat parameter manifest assembly. +3. First joint solve with junction-derived hard-ish constraints. +4. Beautification detection + second solve. +5. Eulerian routing → robot commands. + +The ``Vectorize`` class is a one-shot pipeline: instantiate it with +inputs and read the result fields. Intermediate state is exposed so +the caller can render diagnostics or feed individual phases into +external visualization. + +Result fields: + +* ``self.chains`` — initial per-segment chains (pre-solve). +* ``self.primitives_initial`` — flat primitive list before any solve. +* ``self.soft_initial`` — soft constraints from junction translation. +* ``self.primitives_fitted`` — primitives after the first solve. +* ``self.soft_beautified`` — constraints augmented with beautification. +* ``self.primitives_consolidated`` — primitives after second solve. +* ``self.commands`` — robot commands from ``primitives_fitted``. +* ``self.consolidated`` — robot commands from ``primitives_consolidated``. +""" + +from __future__ import annotations +import math +from dataclasses import dataclass, field +from typing import Dict, List, Optional, Sequence, TypedDict + +import numpy as np +from numpy.typing import NDArray + +# Relative imports from the same subpackage. +from ...commands import DrawingCommand +from ...graph import StrokeGraph + +from .beautify import BeautifyTolerances, detect, merge_arc_pairs, merge_into +from .fitting import ChainPiece, fit_polyline +from .manifest import ( + Coincide, + G1Smooth, + OnCurve, + SoftConstraints, + pack, + parameter_count, + parameter_scales, + unpack, +) +from .primitives import Arc, Circle, Line, Primitive, tangent_at_end +from .residuals import Weights, assemble_residuals +from .routing import order_primitives, to_commands +from .solve import ( + FitConfig, + FittedSegment, + SolveConfig, + SolveResult, + assign_global_ids, + build_chains, + build_junction_constraints, + solve_once, + _bbox_diag, +) + +# --------------------------------------------------------------------------- +# Public configuration (typed-dict form for parity with the rest of the +# codebase's Config namespaces) + + +class FitDict(TypedDict, total=False): + line_tol: float + arc_tol: float + lam_rel: float + min_len: int + use_dp: bool + max_window: int + closed_tol: float + smooth_junction_deg_threshold: float + + +class SolveDict(TypedDict, total=False): + weights_data: float + weights_coincide: float + weights_on_curve: float + weights_g1: float + weights_parallel: float + weights_perpendicular: float + weights_equal_radius: float + weights_concentric: float + weights_radius_reg: float + max_iters: int + method: str # 'trf' or 'lm' + + +class BeautifyDict(TypedDict, total=False): + enabled: bool + parallel_rad: float + perp_rad: float + radius_rel: float + center_abs: float + min_radius: float + min_line_length: float + + +class RouteDict(TypedDict, total=False): + snap_tol: float + pen_up_join_tol: float + + +# --------------------------------------------------------------------------- + + +def _fit_config_from(d: Optional[dict]) -> FitConfig: + cfg = FitConfig() + if d: + for k, v in d.items(): + if hasattr(cfg, k): + setattr(cfg, k, v) + return cfg + + +def _solve_config_from(d: Optional[dict]) -> SolveConfig: + cfg = SolveConfig() + if d is None: + return cfg + w = Weights() + name_map = { + "data": "data", + "coincide": "coincide", + "on_curve": "on_curve", + "g1": "g1", + "parallel": "parallel", + "perpendicular": "perpendicular", + "equal_radius": "equal_radius", + "concentric": "concentric", + "radius_reg": "radius_reg", + } + for key, attr in name_map.items(): + full = f"weights_{key}" + if full in d: + setattr(w, attr, d[full]) + cfg.weights = w + if "max_iters" in d: + cfg.max_iters = d["max_iters"] + if "method" in d: + cfg.method = d["method"] + return cfg + + +def _beautify_tols_from(d: Optional[dict]) -> BeautifyTolerances: + tol = BeautifyTolerances() + if d: + for k, v in d.items(): + if hasattr(tol, k): + setattr(tol, k, v) + return tol + + +# --------------------------------------------------------------------------- + + +class Vectorize: + """Orchestrate the full vectorization pipeline. + + Args: + graph: a built StrokeGraph (polylines + junctions). + start_pos: robot's starting position (image coordinates). + start_heading: robot's starting heading, in radians, in the + x-right / y-down frame (positive = CCW in image = CW on + screen). + fit: chain subdivision configuration. + solve: optimization configuration (per-category weights, etc.). + beautify: beautification tolerance settings. Set ``enabled`` + to False to skip the second solve. + route: routing options (endpoint snap tolerance). + """ + + class Config: + Fit = FitDict + Solve = SolveDict + Beautify = BeautifyDict + Route = RouteDict + + def __init__( + self, + graph: StrokeGraph, + start_pos: NDArray[np.float64], + start_heading: float = 0.0, + fit: Optional[FitDict] = None, + solve: Optional[SolveDict] = None, + beautify: Optional[BeautifyDict] = None, + route: Optional[RouteDict] = None, + ): + self.graph = graph + self.start_pos = np.asarray(start_pos, dtype=float) + self.start_heading = float(start_heading) + + self.fit_config = _fit_config_from(fit) + self.solve_config = _solve_config_from(solve) + self.beautify_enabled = beautify is None or beautify.get("enabled", True) + # Strip the 'enabled' flag before passing to tolerance ctor. + beautify_clean = ( + {k: v for k, v in beautify.items() if k != "enabled"} if beautify else None + ) + self.beautify_tols = _beautify_tols_from(beautify_clean) + self.route_config = dict(route or {}) + + self._run() + + # ---------------------------------------------------------------- + + def _run(self) -> None: + # Phase 1: chain subdivision. + self.fitted_segments: List[FittedSegment] = build_chains( + self.graph, self.fit_config + ) + self.primitives_initial: List[Primitive] = assign_global_ids( + self.fitted_segments + ) + + # Per-primitive source-point map (global ID -> NDArray of source + # pixels assigned to that primitive). + source_points: Dict[int, NDArray] = {} + for seg in self.fitted_segments: + for pid, src in zip(seg.primitive_ids, seg.source_points): + source_points[pid] = src + self.source_points = source_points + + # Phase 2: junction-derived constraints. + soft = build_junction_constraints( + self.graph, + self.fitted_segments, + smooth_junction_deg=self.fit_config.smooth_junction_deg_threshold, + primitives=self.primitives_initial, + ) + # Add internal chain joint constraints. Coincide ALWAYS — we + # want consecutive primitives to share an endpoint. G1 ONLY + # when the joint is actually smooth: top-down splitting puts + # chain breakpoints at sharp corners (the apex of a cat ear, + # the cusp of a heart), where the two adjoining primitives + # have materially different tangents. Adding G1 there forces + # the joint to smooth out — which rounds off the corner and + # turns a triangular ear into a trapezoidal blob. So we + # measure the tangent deflection at the joint and skip G1 + # past the threshold. + # + # EXCEPTION: when one of the two consecutive primitives is a + # Circle (chain produced by sub-loop extraction — the loop + # part fits as a Circle, the rest fits as an Arc), Coincide + # would force the Circle's theta=0 to match the Arc's + # endpoint. theta=0 is the arbitrary convention point on the + # Circle (center + (r, 0)), so a hard match there pulls the + # entire circle to satisfy it (vasesun's sun went from r=90 + # to r=1077 because the solver shifted the center 1500px + # away to put theta=0 at the stem-top junction). Use + # OnCurve instead — the Arc's endpoint must lie SOMEWHERE + # on the Circle's perimeter, not at a specific point. + smooth_thresh_rad = math.radians(self.fit_config.smooth_junction_deg_threshold) + for seg in self.fitted_segments: + pids = seg.primitive_ids + for k in range(len(pids) - 1): + a_pid = pids[k] + b_pid = pids[k + 1] + a_is_circle = isinstance(self.primitives_initial[a_pid], Circle) + b_is_circle = isinstance(self.primitives_initial[b_pid], Circle) + if a_is_circle and not b_is_circle: + soft.on_curve.append( + OnCurve(terminating=(b_pid, "start"), host=a_pid) + ) + elif b_is_circle and not a_is_circle: + soft.on_curve.append( + OnCurve(terminating=(a_pid, "end"), host=b_pid) + ) + elif a_is_circle and b_is_circle: + # Two adjacent Circles in a chain is unusual but + # not impossible. Treat as concentric (their centers + # coincide) rather than endpoint-coincide. + soft.coincide.append(Coincide((a_pid, "end"), (b_pid, "start"))) + else: + soft.coincide.append(Coincide((a_pid, "end"), (b_pid, "start"))) + # G1 only applies between primitives with meaningful + # tangent endpoints (skip if either is a Circle, since + # tangent_at_end on a Circle is at the arbitrary theta=0 + # point and isn't meaningful for chain joints). + if a_is_circle or b_is_circle: + continue + t_end = tangent_at_end(self.primitives_initial[a_pid], "end") + t_start = tangent_at_end(self.primitives_initial[b_pid], "start") + dot = float(np.clip(np.dot(t_end, t_start), -1.0, 1.0)) + deflection = math.acos(dot) + if deflection < smooth_thresh_rad: + soft.g1.append(G1Smooth(a=a_pid, alpha_a=1.0, b=b_pid, alpha_b=0.0)) + self.soft_initial = soft + + # Phase 3: first solve. + pos_scale = _bbox_diag(self.graph) + if self.primitives_initial: + primitives_fitted, result = solve_once( + self.primitives_initial, + source_points, + self.soft_initial, + self.solve_config.weights, + pos_scale=pos_scale, + max_iters=self.solve_config.max_iters, + method=self.solve_config.method, + ) + else: + primitives_fitted = [] + result = SolveResult([], [], soft, True, 0.0, 0) + self.primitives_fitted = primitives_fitted + self.solve_result = result + + # Phase 4: beautification + re-solve. + if self.beautify_enabled and primitives_fitted: + additions = detect(primitives_fitted, self.beautify_tols) + soft_b = SoftConstraints( + coincide=list(self.soft_initial.coincide), + on_curve=list(self.soft_initial.on_curve), + g1=list(self.soft_initial.g1), + parallel=list(self.soft_initial.parallel), + perpendicular=list(self.soft_initial.perpendicular), + equal_radius=list(self.soft_initial.equal_radius), + concentric=list(self.soft_initial.concentric), + ) + merge_into(soft_b, additions) + self.soft_beautified = soft_b + primitives_consolidated, result_b = solve_once( + primitives_fitted, + source_points, + soft_b, + self.solve_config.weights, + pos_scale=pos_scale, + max_iters=self.solve_config.max_iters, + method=self.solve_config.method, + ) + self.primitives_consolidated = primitives_consolidated + self.solve_result_consolidated = result_b + else: + self.soft_beautified = self.soft_initial + self.primitives_consolidated = primitives_fitted + self.solve_result_consolidated = self.solve_result + + # Phase 4.5: merge arc pairs that approximate a single circle. + # The upstream segmenter sometimes breaks a closed shape into + # two ~180° polylines (e.g. the bird's head outline in + # birdlove), and each becomes a separate Arc with a slightly + # different center and radius. Soft EqualRadius / Concentric + # constraints during the solve nudge them toward agreement + # but don't fully merge them. This pass actually consolidates + # the pair into a single Circle when they together cover ≥320° + # and their endpoints close up. + consolidated_after_merge, merged_pairs = merge_arc_pairs( + self.primitives_consolidated + ) + self.merged_arc_pairs = merged_pairs + if merged_pairs: + self.primitives_consolidated = consolidated_after_merge + + # Phase 5: routing + command emission. + snap_tol = float(self.route_config.get("snap_tol", 1.5)) + pen_up_join_tol = float(self.route_config.get("pen_up_join_tol", 0.5)) + + self.tour = order_primitives( + self.primitives_fitted, self.start_pos, snap_tol=snap_tol + ) + self.commands: Sequence[DrawingCommand] = to_commands( + self.primitives_fitted, + self.tour, + self.start_pos, + self.start_heading, + pen_up_join_tol=pen_up_join_tol, + ) + + self.tour_consolidated = order_primitives( + self.primitives_consolidated, self.start_pos, snap_tol=snap_tol + ) + self.consolidated: Sequence[DrawingCommand] = to_commands( + self.primitives_consolidated, + self.tour_consolidated, + self.start_pos, + self.start_heading, + pen_up_join_tol=pen_up_join_tol, + ) + + # ---------------------------------------------------------------- + # Diagnostics + + def primitives_by_polyline(self, which: str = "fitted") -> List[List[Primitive]]: + """Return per-polyline grouped primitives for diagnostic + rendering. + + ``which`` ∈ {'initial', 'fitted', 'consolidated'} picks which + primitive snapshot to slice. + """ + if which == "initial": + flat = self.primitives_initial + elif which == "fitted": + flat = self.primitives_fitted + elif which == "consolidated": + flat = self.primitives_consolidated + else: + raise ValueError(f"unknown snapshot {which!r}") + + out: List[List[Primitive]] = [] + for seg in self.fitted_segments: + out.append([flat[i] for i in seg.primitive_ids]) + return out + + def stats(self) -> str: + n_lines = sum(1 for p in self.primitives_fitted if isinstance(p, Line)) + n_arcs = sum(1 for p in self.primitives_fitted if isinstance(p, Arc)) + n_circles = sum(1 for p in self.primitives_fitted if isinstance(p, Circle)) + n_pen_ups = sum( + 1 for c in self.commands if c["kind"] == "line" and not c["penDown"] + ) + return ( + f"{len(self.primitives_fitted)} primitives " + f"({n_lines} lines, {n_arcs} arcs, {n_circles} circles) " + f"in {len(self.fitted_segments)} chains, " + f"{len(self.commands)} commands " + f"({n_pen_ups} pen-ups), " + f"first-solve cost={self.solve_result.cost:.2f}, " + f"consolidated cost={self.solve_result_consolidated.cost:.2f}" + ) diff --git a/arc_line_vectorization_suede/vectorize/low_geometry/beautify.py b/arc_line_vectorization_suede/vectorize/low_geometry/beautify.py new file mode 100644 index 0000000..bfe7da6 --- /dev/null +++ b/arc_line_vectorization_suede/vectorize/low_geometry/beautify.py @@ -0,0 +1,260 @@ +"""Beautification pass: detect candidate soft constraints. + +After the first solve, scan all primitives for near-relationships and +add them as constraints so the second solve snaps them exact (or close +to it). The four relationships we look for: + +* **Parallel.** Two lines whose direction angle differs by < ``parallel_tol`` rad. +* **Perpendicular.** Two lines whose directions differ by < ``perp_tol`` rad from 90°. +* **Equal radius.** Two arcs/circles whose radii differ by less than + ``radius_tol`` fractionally. +* **Concentric.** Two arcs/circles whose centers are within + ``center_tol`` absolute pixels. + +Tolerances should be calibrated to the upstream noise of hand-drawn +strokes: roughly the natural mode of the residual histogram. Defaults +are starting points; tune on real input. +""" + +from __future__ import annotations +from dataclasses import dataclass +from itertools import combinations +from typing import List, Tuple + +import numpy as np + +from .manifest import ( + Concentric, EqualRadius, Parallel, Perpendicular, SoftConstraints, +) +from .primitives import Arc, Circle, Line, Primitive + + +@dataclass +class BeautifyTolerances: + parallel_rad: float = 0.10 # ~5.7° + perp_rad: float = 0.10 + radius_rel: float = 0.07 # 7% relative radius difference + center_abs: float = 4.0 # absolute pixels + min_radius: float = 2.0 # don't beautify tiny arcs + min_line_length: float = 5.0 # don't beautify tiny lines + + +def _line_angle(line: Line) -> float: + """Direction angle in [-π/2, π/2). Lines have no head/tail; +π is + the same direction as 0, so we fold to a half-circle for + comparisons. + """ + d = line.direction() + return float(np.arctan2(d[1], d[0])) + + +def _radius(prim: Primitive) -> float: + if isinstance(prim, Circle): + return float(prim.radius) + if isinstance(prim, Arc): + return float(prim.radius()) + return float("inf") + + +def _center(prim: Primitive): + if isinstance(prim, Circle): + return prim.center + if isinstance(prim, Arc): + return prim.center() + return None + + +def detect( + prims: List[Primitive], + tol: BeautifyTolerances, +) -> SoftConstraints: + """Return a SoftConstraints bundle of detected candidates. + + Note: only returns the four beautification fields. Caller is + expected to merge these into the existing constraint bundle. + """ + out = SoftConstraints.empty() + + lines: List[Tuple[int, Line]] = [ + (i, p) for i, p in enumerate(prims) + if isinstance(p, Line) and p.length() >= tol.min_line_length + ] + + for (i, a), (j, b) in combinations(lines, 2): + ai = _line_angle(a) + aj = _line_angle(b) + # Fold to the unsigned smallest angle between two undirected + # lines (both ranges of length π). + delta = (ai - aj + np.pi / 2.0) % np.pi - np.pi / 2.0 + absdelta = abs(delta) + if absdelta < tol.parallel_rad: + out.parallel.append(Parallel(i, j)) + elif abs(absdelta - np.pi / 2.0) < tol.perp_rad: + out.perpendicular.append(Perpendicular(i, j)) + + curved: List[Tuple[int, Primitive]] = [ + (i, p) for i, p in enumerate(prims) + if isinstance(p, (Arc, Circle)) and _radius(p) >= tol.min_radius + and np.isfinite(_radius(p)) + ] + + for (i, a), (j, b) in combinations(curved, 2): + ra = _radius(a) + rb = _radius(b) + denom = max(ra, rb) + if denom > 0 and abs(ra - rb) / denom < tol.radius_rel: + out.equal_radius.append(EqualRadius(i, j)) + ca = _center(a) + cb = _center(b) + if ca is not None and cb is not None: + if float(np.linalg.norm(ca - cb)) < tol.center_abs: + out.concentric.append(Concentric(i, j)) + + return out + + +def merge_into(target: SoftConstraints, additions: SoftConstraints) -> None: + """Merge beautification additions into ``target`` in place.""" + target.parallel.extend(additions.parallel) + target.perpendicular.extend(additions.perpendicular) + target.equal_radius.extend(additions.equal_radius) + target.concentric.extend(additions.concentric) + + +def merge_arc_pairs( + prims: List[Primitive], + center_tol_rel: float = 1.0, + radius_tol_rel: float = 0.10, + endpoint_tol_rel: float = 0.20, + full_circle_sweep_thresh_deg: float = 320.0, + sweep_match_tol_deg: float = 40.0, +) -> Tuple[List[Primitive], List[Tuple[int, int]]]: + """Detect and merge pairs of Arc primitives that approximate two + halves of the same circle. + + Two common forms appear in real data: + + 1. **Chained halves.** Arc A runs from p0 to p1, arc B runs from + p1' ≈ p1 back to p0' ≈ p0, going the other way around. The two + chain end-to-end. + 2. **Parallel halves.** Both arcs share start and end endpoints, + and bulge in opposite directions. This happens when upstream + segmentation produces TWO polylines tracing the same loop in + opposite directions (e.g. the bird's head outline in birdlove + gets split into upper and lower halves, each starting at the + chord endpoint nearest the beak and ending at the chord + endpoint nearest the body). + + Either way, the criteria are: + * radii within ``radius_tol_rel`` fractionally + * centers within ``center_tol_rel * mean_radius`` (loose because + each half-fit pulled its center toward its own data) + * endpoints connect (one of the two patterns above) + * combined sweep ≥ ``full_circle_sweep_thresh_deg`` + + If all hold, the pair is replaced with a single Circle at the + chord midpoint (which is closer to the true center than either + individual fit when the two fits split the difference). + + Returns ``(new_primitives, merged_pairs)``. + """ + import math + n = len(prims) + if n < 2: + return list(prims), [] + + arc_pids = [ + i for i, p in enumerate(prims) + if isinstance(p, Arc) and np.isfinite(_radius(p)) and _radius(p) > 2.0 + ] + + used = [False] * n + merged_pairs: List[Tuple[int, int]] = [] + replacement: dict = {} + + for ii in range(len(arc_pids)): + i = arc_pids[ii] + if used[i]: + continue + a = prims[i] + assert isinstance(a, Arc) + ra = _radius(a) + ca = _center(a) + + for jj in range(ii + 1, len(arc_pids)): + j = arc_pids[jj] + if used[j]: + continue + b = prims[j] + assert isinstance(b, Arc) + rb = _radius(b) + cb = _center(b) + r_mean = 0.5 * (ra + rb) + + if abs(ra - rb) / max(ra, rb) > radius_tol_rel: + continue + if float(np.linalg.norm(ca - cb)) > center_tol_rel * r_mean: + continue + + tol_pt = max(endpoint_tol_rel * r_mean, 3.0) + d_e_s = float(np.linalg.norm(a.p1 - b.p0)) + d_e_e = float(np.linalg.norm(a.p1 - b.p1)) + d_s_s = float(np.linalg.norm(a.p0 - b.p0)) + d_s_e = float(np.linalg.norm(a.p0 - b.p1)) + + chained = (d_e_s < tol_pt and d_s_e < tol_pt) or \ + (d_e_e < tol_pt and d_s_s < tol_pt) + # Parallel halves: both endpoints coincide as a pair. + parallel = (d_s_s < tol_pt and d_e_e < tol_pt) or \ + (d_s_e < tol_pt and d_e_s < tol_pt) + if not (chained or parallel): + continue + + sweep_a_abs = abs(math.degrees(a.sweep())) + sweep_b_abs = abs(math.degrees(b.sweep())) + sweep_total = sweep_a_abs + sweep_b_abs + if sweep_total < full_circle_sweep_thresh_deg: + continue + + # CRUCIAL: a true circle traced as two halves has + # |sweep_a| ≈ |sweep_b| ≈ 180° (each half covers half + # the perimeter). A crescent moon's outer arc wraps more + # than 180° and its inner arc is well under 180°, so + # |sweep_a| - |sweep_b| can be 70-100° even when the + # radii happen to be similar. The eggmoon crescent has + # radii 103 vs 99 (4% diff, passes the radius gate) but + # sweeps 224° vs 150° — clearly not the same circle. + # Require the sweeps to be within sweep_match_tol_deg. + if abs(sweep_a_abs - sweep_b_abs) > sweep_match_tol_deg: + continue + + # Compute the geometrically-correct center from the chord + # midpoint. For the parallel-halves case (both arcs share + # endpoints), the two fits typically pull their centers + # in opposite directions along the perpendicular bisector + # of the chord. The true center lies on that bisector at + # the right distance, and the simple average of the two + # noisy fits is usually a much better estimate than + # either fit alone. + new_center = 0.5 * (ca + cb) + new_radius = r_mean + + replacement[i] = Circle(center=new_center, radius=new_radius) + replacement[j] = None + used[i] = True + used[j] = True + merged_pairs.append((i, j)) + break + + if not merged_pairs: + return list(prims), [] + + new_prims: List[Primitive] = [] + for pid, p in enumerate(prims): + if pid in replacement: + rep = replacement[pid] + if rep is not None: + new_prims.append(rep) + else: + new_prims.append(p) + return new_prims, merged_pairs diff --git a/arc_line_vectorization_suede/vectorize/low_geometry/fitting.py b/arc_line_vectorization_suede/vectorize/low_geometry/fitting.py new file mode 100644 index 0000000..5afe4a7 --- /dev/null +++ b/arc_line_vectorization_suede/vectorize/low_geometry/fitting.py @@ -0,0 +1,846 @@ +"""Per-segment fitting and chain subdivision. + +Two layers: + +* Single-primitive fitters: ``fit_line``, ``fit_circle``, ``fit_arc``. + These are the building blocks both for terminal classification and as + the per-window cost function in chain subdivision. +* ``fit_segment_chain`` / ``fit_segment_topdown``: split a long polyline + into a sequence of primitives that share endpoints, trading off + fidelity against complexity. + +Single-primitive fits return a ``(primitive, rms)`` pair so the caller +can decide whether the fit was good enough to commit to. + +Chain subdivision: a single segment from the upstream pipeline may be a +long fluid stroke that needs multiple primitives. The DP variant +minimizes ``Σ SSE_i + λ * n_primitives`` over all chain decompositions, +which is the MDL-style "fidelity vs simplicity" tradeoff from Favreau +et al. Top-down recursive split (``fit_segment_topdown``) is the +simpler, near-linear fallback that often produces cleaner splits at +high-curvature points because the split location IS the worst-fit +point. +""" + +from __future__ import annotations +import math +from dataclasses import dataclass +from typing import List, Optional, Tuple + +import numpy as np +from numpy.typing import NDArray +from scipy.optimize import least_squares + +from .primitives import Arc, Circle, Line, Primitive + + +_EPS = 1e-9 + + +# --------------------------------------------------------------------------- +# Polyline utilities + + +def polyline_length(pts: NDArray[np.float64]) -> float: + if len(pts) < 2: + return 0.0 + diffs = np.diff(pts, axis=0) + return float(np.sum(np.linalg.norm(diffs, axis=1))) + + +def is_closed_polyline(pts: NDArray[np.float64], tol: float = 1.5) -> bool: + if len(pts) < 4: + return False + return float(np.linalg.norm(pts[0] - pts[-1])) < tol + + +def is_near_closed_polyline( + pts: NDArray[np.float64], + gap_ratio: float = 0.10, + abs_tol: float = 10.0, +) -> bool: + """Like ``is_closed_polyline`` but also catches polylines where the + endpoints have a small gap *relative to* the polyline's arc length. + Hand-drawn wheels and other closed shapes often have a 5-30 px gap + where the stroke didn't quite meet itself; we want to recognize these + as Circles, not as 350° arcs. + + A polyline is "near-closed" if EITHER: + * absolute endpoint gap < ``abs_tol`` (covers small drawings), OR + * gap / polyline_length < ``gap_ratio`` (covers larger drawings). + """ + if len(pts) < 4: + return False + gap = float(np.linalg.norm(pts[0] - pts[-1])) + if gap < abs_tol: + return True + L = polyline_length(pts) + if L < _EPS: + return False + return gap / L < gap_ratio + + +def find_corners( + pts: NDArray[np.float64], + sharp_threshold_deg: float = 80.0, + soft_threshold_deg: float = 50.0, + sustain_indices: int = 3, + min_separation: Optional[int] = None, + window: Optional[int] = None, +) -> List[int]: + """Like ``count_corners`` but returns the *indices* of detected + corners (into ``pts``). See ``count_corners`` for the algorithm. + + A corner is reported when EITHER: + * the turn at this index exceeds ``sharp_threshold_deg`` (a + clear, sharp corner that doesn't need any neighborhood + check — e.g., a cat-ear apex at ~130°); OR + * the turn exceeds ``soft_threshold_deg`` AND every turn + within ``sustain_indices`` on each side ALSO exceeds + ``soft_threshold_deg`` (a gentle but sustained corner — a + house corner at ~80° has neighbors also at ~75-83° for many + indices, whereas a noise spike on a circle has a high turn + at one index with neighbors at <25°). + + The neighbor-sustained check is the key discriminator between + real corners (which are wide in turn-angle-space because the + polyline has a sustained direction change over several samples) + and per-pixel noise on a hand-drawn curve (which produces + isolated spikes whose neighbors drop back to baseline). + + The tangent window scales with polyline length: + ``w = max(8, min(30, n // 30))``. Short polylines use ``w=8`` + (avoids chord-to-chord false positives on small clean circles); + long polylines use up to ``w=30`` to recover corners the + upstream skeletonization rounded over ~10 pixels. + """ + n = len(pts) + if window is None: + w = max(8, min(30, n // 30)) + else: + w = window + if min_separation is None: + min_separation = max(12, w + 4) + if n < 2 * w + 1: + return [] + + left = pts[w:n - w] - pts[: n - 2 * w] + right = pts[2 * w:] - pts[w:n - w] + ll = np.linalg.norm(left, axis=1) + rl = np.linalg.norm(right, axis=1) + valid = (ll > _EPS) & (rl > _EPS) + left = left.copy() + right = right.copy() + left[valid] = left[valid] / ll[valid, None] + right[valid] = right[valid] / rl[valid, None] + + dots = np.einsum("ij,ij->i", left, right) + dots = np.clip(dots, -1.0, 1.0) + turns = np.arccos(dots) # radians + + sharp_thresh = math.radians(sharp_threshold_deg) + soft_thresh = math.radians(soft_threshold_deg) + n_turns = len(turns) + + corners: List[int] = [] + i = 0 + while i < n_turns: + t = turns[i] + is_sharp = t > sharp_thresh + # Sustained-soft: turn is above soft threshold AND so are its + # immediate neighbors at +/-sustain_indices. + is_sustained = False + if t > soft_thresh: + lo_idx = i - sustain_indices + hi_idx = i + sustain_indices + if (lo_idx >= 0 and hi_idx < n_turns + and turns[lo_idx] > soft_thresh + and turns[hi_idx] > soft_thresh): + is_sustained = True + if is_sharp or is_sustained: + # Non-max suppression in [i, i+min_separation). + j_end = min(n_turns, i + min_separation) + best = i + for j in range(i + 1, j_end): + if turns[j] > turns[best]: + best = j + corners.append(best + w) + i = best + min_separation + else: + i += 1 + return corners + + +def count_corners( + pts: NDArray[np.float64], + sharp_threshold_deg: float = 80.0, + soft_threshold_deg: float = 50.0, + sustain_indices: int = 3, + min_separation: Optional[int] = None, + window: Optional[int] = None, +) -> int: + """Count the number of corners along a polyline. See + ``find_corners`` for the algorithm. + """ + return len(find_corners( + pts, sharp_threshold_deg, soft_threshold_deg, sustain_indices, + min_separation, window, + )) + + +# --------------------------------------------------------------------------- +# Single-primitive fitters + + +def fit_line(pts: NDArray[np.float64]) -> Tuple[Line, float]: + """ODR line fit. Returns (Line, rms_perp_residual). + + Math: minimize Σ ‖(pᵢ - μ) - ((pᵢ - μ)·d) d‖² over unit d. Closed-form + via SVD of the centered point matrix; the leading right singular + vector is the optimal direction. + """ + if len(pts) < 2: + return Line(pts[0].copy(), pts[0].copy()), 0.0 + mu = pts.mean(axis=0) + centered = pts - mu + _, _, vh = np.linalg.svd(centered, full_matrices=False) + d = vh[0] + t = centered @ d + p0 = mu + t.min() * d + p1 = mu + t.max() * d + perp = centered - np.outer(t, d) + rms = float(np.sqrt((perp ** 2).sum(axis=1).mean())) + return Line(p0, p1), rms + + +def fit_circle_kasa( + pts: NDArray[np.float64], +) -> Tuple[NDArray[np.float64], float]: + """Algebraic circle fit (Kasa). Fast initializer for geometric refine. + + Minimizes Σ (xᵢ² + yᵢ² + a xᵢ + b yᵢ + c)² with + a = -2 cx, b = -2 cy, c = cx² + cy² - r². Linear in (a, b, c). + """ + x, y = pts[:, 0], pts[:, 1] + M = np.column_stack([x, y, np.ones_like(x)]) + rhs = -(x * x + y * y) + sol, *_ = np.linalg.lstsq(M, rhs, rcond=None) + a, b, c = sol + cx, cy = -a / 2.0, -b / 2.0 + r_sq = cx * cx + cy * cy - c + if r_sq < 0: + # Numerical degeneracy (collinear points); return huge radius. + return np.array([cx, cy]), 1e9 + return np.array([cx, cy]), float(np.sqrt(r_sq)) + + +def fit_circle_geometric( + pts: NDArray[np.float64], + c0: NDArray[np.float64], + r0: float, +) -> Tuple[NDArray[np.float64], float, float]: + """Geometric circle refine. Returns (center, radius, rms). + + Minimizes Σ (‖pᵢ - c‖ - r)² which is the true geometric residual. + Kasa's bias (underestimates radius for short arcs — see Chernov) is + fixed up by this stage. + """ + + def residuals(params): + return np.linalg.norm(pts - params[:2], axis=1) - params[2] + + x0 = np.array([c0[0], c0[1], r0]) + try: + sol = least_squares(residuals, x0, method="lm", max_nfev=50) + c = sol.x[:2] + r = float(sol.x[2]) + except Exception: + c, r = c0, r0 + res = np.linalg.norm(pts - c, axis=1) - r + rms = float(np.sqrt((res ** 2).mean())) + return c, r, rms + + +def fit_circle( + pts: NDArray[np.float64], +) -> Tuple[NDArray[np.float64], float, float]: + """Combined Kasa + geometric refine. Returns (center, radius, rms).""" + if len(pts) < 3: + # Degenerate; can't fit a circle to <3 points. + return np.array([0.0, 0.0]), 0.0, float("inf") + c0, r0 = fit_circle_kasa(pts) + return fit_circle_geometric(pts, c0, r0) + + +def fit_arc(pts: NDArray[np.float64]) -> Tuple[Optional[Arc], float]: + """Fit a circle to the points, then build an arc using the first and + last points as endpoints. Returns (Arc, rms) or (None, inf) if fit + fails. + + Sweep direction is disambiguated by looking at the midpoint sample's + side of the chord. Sweep magnitude assumes the MINOR arc; if the + sample falls on the wrong side we flip to the major arc. + """ + if len(pts) < 3: + return None, float("inf") + c, r, rms = fit_circle(pts) + if not np.isfinite(r) or r < _EPS: + return None, float("inf") + + p0 = pts[0] + p1 = pts[-1] + chord = p1 - p0 + L = float(np.linalg.norm(chord)) + if L < _EPS: + # Endpoints coincide — treat as full circle, not an arc. + return None, float("inf") + + # Determine sweep direction. Cross product chord × (mid - p0): + # if positive (in y-down) and the geometric center is on the same + # side as the mid sample, sweep is positive (CCW in image coords). + mid_sample_idx = len(pts) // 2 + mid_sample = pts[mid_sample_idx] + side_mid = chord[0] * (mid_sample[1] - p0[1]) - chord[1] * (mid_sample[0] - p0[0]) + side_center = chord[0] * (c[1] - p0[1]) - chord[1] * (c[0] - p0[0]) + + # half-chord angle: sin(θ) = (L/2) / r. Clamp for numerical safety. + half_chord_over_r = min(1.0, max(-1.0, L / (2.0 * r))) + half_sweep_minor = float(np.arcsin(half_chord_over_r)) + + # If the mid sample is on the OPPOSITE side of the chord from the + # center, the arc passes around the far side -> major arc. + is_major = (side_mid * side_center > 0.0) + sweep_mag = ( + 2.0 * (np.pi - half_sweep_minor) if is_major else 2.0 * half_sweep_minor + ) + + # An arc with |sweep| > ~240° means the polyline traces a large + # majority of a circle but doesn't quite close. Fitting such a + # polyline as a single Arc produces a near-degenerate bulge + # magnitude (the endpoints are close together, so a tiny + # perturbation flips the arc through 180°). The joint solver + # tends to "fix" these by collapsing the major arc into a much + # smaller minor arc with a different radius, which radically + # changes the shape (the tree in treecar1 went from a round + # outline to a narrow teardrop because a -264° arc fitting the + # tree's bottom collapsed to a -72° arc). Better to let the + # polyline either be treated as a Circle (handled in + # fit_polyline's near-closed branch) or be split by chain + # subdivision into multiple smaller arcs that ARE stable. + if sweep_mag > math.radians(240.0): + return None, float("inf") + + # Direction: traversal goes p0 -> mid -> p1. The arc bulges toward + # the side where ``mid_sample`` lies. With our convention + # ``normal = (-chord_y, chord_x)``, the Arc's ``center()`` puts the + # center on the POSITIVE-cross side of the chord, so the arc + # midpoint is on the NEGATIVE-cross side. So: + # side_mid < 0 → arc midpoint on negative-cross side → bulge > 0 + # side_mid > 0 → arc midpoint on positive-cross side → bulge < 0 + sweep_sign = -1.0 if side_mid > 0 else 1.0 + sweep = sweep_sign * sweep_mag + bulge = float(np.tan(sweep / 4.0)) + + arc = Arc(p0.copy(), p1.copy(), bulge) + return arc, rms + + +def fit_full_circle( + pts: NDArray[np.float64], +) -> Tuple[Optional[Circle], float]: + """Fit a Circle (not Arc) to a closed polyline. Returns (Circle, rms).""" + if len(pts) < 3: + return None, float("inf") + c, r, rms = fit_circle(pts) + if not np.isfinite(r) or r < _EPS: + return None, float("inf") + return Circle(c, r), rms + + +# --------------------------------------------------------------------------- +# Chain subdivision + + +@dataclass(frozen=True) +class ChainPiece: + start_idx: int # inclusive + end_idx: int # exclusive + primitive: Primitive + + +class FitFailure(RuntimeError): + pass + + +def _segment_extent(pts: NDArray[np.float64]) -> float: + """A length scale to normalize tolerances by. Use bounding-box + diagonal because it's stable for both straight and curvy strokes — + arc length over-rewards wiggly fits. + """ + if len(pts) < 2: + return 1.0 + span = pts.max(axis=0) - pts.min(axis=0) + diag = float(np.linalg.norm(span)) + return max(diag, 1.0) + + +def fit_single_primitive( + pts: NDArray[np.float64], + line_tol_abs: float, + arc_tol_abs: float, +) -> Tuple[Optional[Primitive], float]: + """Try line and arc; return whichever fits within tolerance with the + smaller SSE. Returns ``(primitive, sse)`` or ``(None, inf)``. + + Tolerances are absolute pixel RMS thresholds. + """ + n = len(pts) + if n < 2: + return None, float("inf") + + line, line_rms = fit_line(pts) + line_ok = line_rms < line_tol_abs + line_sse = (line_rms ** 2) * n + + if n < 3: + if line_ok: + return line, line_sse + return None, float("inf") + + arc, arc_rms = fit_arc(pts) + arc_ok = arc is not None and arc_rms < arc_tol_abs + arc_sse = (arc_rms ** 2) * n if arc is not None else float("inf") + + # Prefer line when both work and line SSE isn't much worse — fewer + # parameters, simpler downstream routing, doesn't suffer from bulge + # numerics. The "1.5x" gives arcs a fair shot when the curve is real. + if line_ok and (not arc_ok or line_sse <= 1.5 * arc_sse): + return line, line_sse + if arc_ok: + return arc, arc_sse + if line_ok: + return line, line_sse + return None, float("inf") + + +def fit_segment_topdown( + pts: NDArray[np.float64], + line_tol_abs: float, + arc_tol_abs: float, + min_len: int = 5, + max_depth: int = 12, +) -> List[ChainPiece]: + """Recursive split-and-fit. + + Fit one primitive to the whole window. If it fits within tolerance, + accept. Otherwise split at the index of maximum residual and + recurse on each half. O(N log N) typical. + + This is the plan's recommended starting point — it tends to put + splits at semantically meaningful places (the worst-fit point + *is* the corner) and avoids DP's pathological ties between + near-equivalent splits. + """ + pieces: List[ChainPiece] = [] + + def recurse(lo: int, hi: int, depth: int) -> None: + if hi - lo < min_len or depth >= max_depth: + # Forced terminal — fit whatever single primitive we can. + prim, _ = fit_single_primitive( + pts[lo:hi], line_tol_abs * 4, arc_tol_abs * 4 + ) + if prim is None: + # Last-ditch: straight line through endpoints. + prim = Line(pts[lo].copy(), pts[hi - 1].copy()) + pieces.append(ChainPiece(lo, hi, prim)) + return + + prim, _ = fit_single_primitive(pts[lo:hi], line_tol_abs, arc_tol_abs) + if prim is not None: + pieces.append(ChainPiece(lo, hi, prim)) + return + + # Find split point: index of maximum residual from the best of + # line / arc fits even though neither was good enough. + line, _ = fit_line(pts[lo:hi]) + line_res = np.abs(line.perpendicular_distance(pts[lo:hi])) + if hi - lo >= 3: + c, r, _ = fit_circle(pts[lo:hi]) + arc_res = np.abs(np.linalg.norm(pts[lo:hi] - c, axis=1) - r) + res = np.minimum(line_res, arc_res) + else: + res = line_res + + # Split at the worst-fit interior index; clamp to keep both + # children at least ``min_len`` long. + worst = int(np.argmax(res)) + worst += lo + worst = max(lo + min_len, min(hi - min_len, worst)) + if worst <= lo or worst >= hi: + # Couldn't find a valid split; accept the whole window as-is. + prim = line + pieces.append(ChainPiece(lo, hi, prim)) + return + + recurse(lo, worst + 1, depth + 1) # include the split point in both + recurse(worst, hi, depth + 1) + + recurse(0, len(pts), 0) + # Stitch: sort + ensure indices form a contiguous chain + pieces.sort(key=lambda p: p.start_idx) + return pieces + + +def fit_segment_dp( + pts: NDArray[np.float64], + line_tol_abs: float, + arc_tol_abs: float, + lam: float, + min_len: int = 5, + max_window: Optional[int] = None, +) -> List[ChainPiece]: + """Global DP chain subdivision. + + ``cost(i, j)`` = best single-primitive SSE for ``pts[i:j]``, + or ``+inf`` if neither line nor arc fits within tolerance. + ``best(j)`` = min total cost to fit ``pts[0:j]``. + + Recurrence: ``best(j) = min over i < j of [best(i) + cost(i,j) + lam]``. + + ``lam`` is the MDL penalty per primitive — larger ``lam`` → fewer, + looser primitives. Start around ``2 * line_tol_abs²`` and tune. + + ``max_window`` caps the maximum (j - i) considered, which makes the + inner loop O(N · max_window) instead of O(N²). For most real + drawings a window of 200-500 source points is plenty. + """ + N = len(pts) + if N < min_len: + prim, _ = fit_single_primitive(pts, line_tol_abs * 4, arc_tol_abs * 4) + if prim is None: + prim = Line(pts[0].copy(), pts[-1].copy()) + return [ChainPiece(0, N, prim)] + + cache: dict = {} + + def get_cost(i: int, j: int) -> Tuple[float, Optional[Primitive]]: + if (i, j) in cache: + return cache[(i, j)] + if j - i < min_len: + cache[(i, j)] = (float("inf"), None) + return cache[(i, j)] + prim, sse = fit_single_primitive(pts[i:j], line_tol_abs, arc_tol_abs) + cache[(i, j)] = (sse, prim) + return cache[(i, j)] + + best = [0.0] + [float("inf")] * N + split = [-1] * (N + 1) + + if max_window is None: + max_window = N + + for j in range(min_len, N + 1): + i_lo = max(0, j - max_window) + for i in range(i_lo, j - min_len + 1): + if not np.isfinite(best[i]): + continue + sse, _ = get_cost(i, j) + if not np.isfinite(sse): + continue + total = best[i] + sse + lam + if total < best[j]: + best[j] = total + split[j] = i + + # Reconstruct the chain by backtracking from N. + chain: List[ChainPiece] = [] + j = N + while j > 0: + i = split[j] + if i < 0: + # No valid split found — usually means the tolerance is too + # tight for this stroke. Fall back to top-down on the + # unfit suffix. + fallback = fit_segment_topdown( + pts[:j], line_tol_abs, arc_tol_abs, min_len + ) + chain = fallback + chain + return chain + _, prim = get_cost(i, j) + if prim is None: + raise FitFailure( + f"DP picked split ({i}, {j}) with no fitted primitive" + ) + chain.append(ChainPiece(i, j, prim)) + j = i + chain.reverse() + return chain + + +def find_closed_subloop( + pts: NDArray[np.float64], + min_size: int = 100, + closure_threshold_rel: float = 0.06, + min_extent_rel: float = 0.20, + max_circle_rms_rel: float = 0.06, +) -> Optional[Tuple[int, int]]: + """Find a near-closed sub-loop within an open polyline that + ALSO fits a single Circle well. + + Looks for (i, j) with ``j - i >= min_size`` such that: + 1. ``||pts[j] - pts[i]||`` < ``closure_threshold_rel * polyline_extent`` + (the endpoints meet) + 2. ``extent(pts[i:j+1])`` >= ``min_extent_rel * polyline_extent`` + (the loop has meaningful spatial size; rules out tiny + self-crossings) + 3. The sub-loop fits a circle with RMS below + ``max_circle_rms_rel * loop_extent`` (the loop is actually + circular — rules out V-shapes or U-shapes whose endpoints + happen to be close but whose interior path isn't a circle) + + Returns ``(i, j)`` for the LARGEST qualifying sub-loop, or + ``None``. The largest is preferred because outer/wider loops + are usually the intended shape (a wheel rim, not a small + embedded swirl). + + Example: bikelove's right-wheel polyline poly[14] traces the + rim for ~650 indices and then continues into the bottom + squiggle. Detecting [0, 650] as a circular sub-loop lets the + rim become a Circle while the squiggle gets fit separately. + """ + n = len(pts) + if n < 2 * min_size: + return None + extent = _segment_extent(pts) + if extent < 1.0: + return None + closure_thresh = closure_threshold_rel * extent + extent_thresh = min_extent_rel * extent + + # Coarse search on a downsampled grid. + stride = max(1, n // 200) + candidates: List[Tuple[int, int, int]] = [] # (size, i, j) + for i in range(0, n - min_size, stride): + for j in range(n - 1, i + min_size, -stride): + d = float(np.linalg.norm(pts[j] - pts[i])) + if d < closure_thresh: + sub = pts[i:j + 1] + sub_ext = _segment_extent(sub) + if sub_ext >= extent_thresh: + candidates.append((j - i, i, j)) + break # take largest j for this i + + if not candidates: + return None + + # Try the largest candidates first; require a clean circle fit + # AND a reasonable aspect ratio (a true wheel/sun is roughly + # circular, not stretched). + candidates.sort(reverse=True) + for size, i, j in candidates: + sub = pts[i:j + 1] + circle, rms = fit_full_circle(sub) + if circle is None: + continue + loop_ext = _segment_extent(sub) + if rms >= max_circle_rms_rel * loop_ext: + continue + xs = sub[:, 0] + ys = sub[:, 1] + bbox_x = float(xs.max() - xs.min()) + bbox_y = float(ys.max() - ys.min()) + if min(bbox_x, bbox_y) < _EPS: + continue + aspect = max(bbox_x, bbox_y) / min(bbox_x, bbox_y) + if aspect >= 1.25: + continue + # Refine: shift i and j by +/-stride to find the exact + # closest pair. + i0, j0 = i, j + best_d = float(np.linalg.norm(pts[j0] - pts[i0])) + refined_i, refined_j = i0, j0 + for di in range(-stride, stride + 1): + ii = i0 + di + if ii < 0 or ii >= n: + continue + for dj in range(-stride, stride + 1): + jj = j0 + dj + if jj < 0 or jj >= n or jj - ii < min_size: + continue + d = float(np.linalg.norm(pts[jj] - pts[ii])) + if d < best_d: + best_d = d + refined_i, refined_j = ii, jj + return (refined_i, refined_j) + return None + + +def fit_polyline( + pts: NDArray[np.float64], + line_tol: float = 0.005, + arc_tol: float = 0.01, + lam_rel: float = 4.0, + min_len: int = 5, + closed_tol: float = 1.5, + use_dp: bool = True, + max_window: Optional[int] = 256, + circle_rms_rel: float = 0.06, +) -> List[ChainPiece]: + """Top-level: fit a polyline into a chain of primitives. + + * If the polyline is closed and fits well as a single circle, + shortcut to a one-piece chain. The "well" tolerance here is + ``circle_rms_rel * extent`` — much looser than ``arc_tol`` + because hand-drawn closed loops (wheels, balloons, hearts) + are usually meant to read as a single circle even when they + wobble by several percent of the polyline's extent. Trying to + preserve the wobble by chain-subdividing produces 10-piece + fragmentations that look much worse than a clean circle. + * Otherwise, try ``fit_single_primitive`` on the whole stroke + first — most segments are single primitives once chromosomes + and crossings have been removed upstream. + * Otherwise, run chain subdivision (DP by default, top-down as + fallback). + + ``line_tol`` and ``arc_tol`` are RELATIVE to the bounding-box + diagonal, so the same parameters work across drawings of + different scale. ``lam_rel`` is a multiplier on + ``(line_tol * extent)²`` to set the per-primitive MDL penalty. + """ + n = len(pts) + if n < 2: + return [] + extent = _segment_extent(pts) + line_tol_abs = line_tol * extent + arc_tol_abs = arc_tol * extent + lam = lam_rel * (line_tol_abs ** 2) + circle_rms_abs = circle_rms_rel * extent + + # Closed-loop fast path: try fitting as a single circle. Use the + # near-closed detector (gap small relative to arc length OR + # absolutely small) so hand-drawn wheels with a 10-30 px gap also + # get caught — fit_arc would otherwise produce a near-360° arc + # which is much worse visually than a Circle. + # + # BUT: only use the Circle shortcut if the polyline is genuinely + # corner-free. A cat-head outline that includes ears is closed AND + # the rms-of-circle-fit is moderate (the ears are short relative + # to the head circumference), so the RMS gate alone would happily + # erase the ears. Counting corners catches this — ear tips show up + # as 80°+ tangent-direction changes that a true circle never has. + corners = find_corners(pts) + if is_near_closed_polyline(pts) and not corners: + circle, rms = fit_full_circle(pts) + if circle is not None and rms < circle_rms_abs: + # Additionally check the polyline's aspect ratio. A closed + # polyline that traces an oval (e.g., the catcar cat-head, + # 444x339 pixels = aspect 1.31) can still fit a circle with + # low rms (3.5% in that case, below the 6% gate), but + # rendering it as a perfect circle is visually wrong — the + # circle's radius is the AVERAGE of the oval's axes, so the + # circle extends past the polyline in the shorter dimension. + # In catcar this pushes the cat-head circle down through the + # roof of the car (a clear visual overlap). Reject the + # Circle shortcut when the bounding-box aspect ratio is too + # far from 1, and let chain subdivision fit the oval with + # multiple arcs instead. + xs = pts[:, 0] + ys = pts[:, 1] + bbox_x = float(xs.max() - xs.min()) + bbox_y = float(ys.max() - ys.min()) + if min(bbox_x, bbox_y) > _EPS: + aspect = max(bbox_x, bbox_y) / min(bbox_x, bbox_y) + else: + aspect = 1.0 + if aspect < 1.25: + return [ChainPiece(0, n, circle)] + # Otherwise fall through; the DP will handle rounded rectangles. + + # If the polyline HAS corners (ear apexes, heart cusps), split + # explicitly at each corner index before running chain + # subdivision. The chain subdivider only minimizes residuals — it + # doesn't know about corners. A corner spans 2-3 indices in the + # polyline and is essentially unfit-able by any single primitive, + # so without an explicit split, top-down/DP both end up producing + # 10+ tiny pieces around the corner trying to thread the needle. + # Splitting first lets each side of the corner be fit cleanly as + # a single primitive. + # + # ALSO: look for a near-closed sub-loop within the polyline (e.g., + # the bikelove right wheel rim is the first ~650 indices of a + # 994-pt polyline that continues into the bottom squiggle). If a + # sub-loop exists, split at its start and end indices so the loop + # portion is processed as its own near-closed sub-polyline (which + # will then hit the closed-circle shortcut). + subloop = find_closed_subloop(pts) if not is_near_closed_polyline(pts) else None + splits: List[int] = list(corners) + if subloop is not None: + splits.extend(subloop) + splits = sorted(set(splits)) + if splits: + pieces: List[ChainPiece] = [] + split_idxs = [0] + splits + [n] + # Remove duplicates while preserving order + split_idxs = sorted(set(split_idxs)) + for lo, hi in zip(split_idxs[:-1], split_idxs[1:]): + sub = pts[lo:hi + 1] if hi < n else pts[lo:] + if len(sub) < 2: + continue + sub_chain = fit_polyline( + sub, line_tol=line_tol, arc_tol=arc_tol, + lam_rel=lam_rel, min_len=min_len, closed_tol=closed_tol, + use_dp=use_dp, max_window=max_window, + circle_rms_rel=circle_rms_rel, + ) + for cp in sub_chain: + pieces.append( + ChainPiece(cp.start_idx + lo, cp.end_idx + lo, cp.primitive) + ) + return pieces + + # Corner-free, non-closed polyline: try fitting it as a single Arc + # with a relaxed tolerance, similar to the closed-loop circle + # shortcut. The threshold is intentionally TIGHTER than the + # closed-loop circle shortcut because most non-circle polylines + # (heart halves, bike-frame contours, leaf outlines) fit a single + # arc within 5-10% rms but aren't really arcs — they have + # systematic curvature variation that a single arc averages away. + # If we accept those as one arc we lose the heart cusps, + # rectangle corners, leaf points, etc. Only accept the shortcut + # when the fit is so close to a real arc that chain subdivision + # couldn't do meaningfully better. + # + # Critically, the threshold is BOTH 4% of extent AND capped at an + # absolute pixel ceiling. A pure relative threshold relaxes + # linearly with extent, so a 700-pixel polyline gets a 28-px + # tolerance — that's enough to swallow significant shape detail + # (heartman's body sub-segments were 706 pts with arc-fit rms of + # 23 px, just under 28 px, so the whole body collapsed to one + # arc with 23 px of accumulated deviation). The absolute cap of + # ~12 px keeps the shortcut "this is essentially noise on a clean + # arc" for polylines of any length. + arc_rms_abs = min(0.04 * extent, 12.0) + if not corners: + arc, arc_rms = fit_arc(pts) + if (arc is not None and arc_rms < arc_rms_abs and + arc.chord() > line_tol_abs * 2): + return [ChainPiece(0, n, arc)] + + # Single-primitive shortcut (tight tolerance, line OR arc). + single, _ = fit_single_primitive(pts, line_tol_abs, arc_tol_abs) + if single is not None: + return [ChainPiece(0, n, single)] + + # Chain subdivision. + if use_dp: + try: + chain = fit_segment_dp( + pts, line_tol_abs, arc_tol_abs, lam, + min_len=min_len, max_window=max_window, + ) + except FitFailure: + chain = fit_segment_topdown(pts, line_tol_abs, arc_tol_abs, min_len) + else: + chain = fit_segment_topdown(pts, line_tol_abs, arc_tol_abs, min_len) + + return chain diff --git a/arc_line_vectorization_suede/vectorize/low_geometry/manifest.py b/arc_line_vectorization_suede/vectorize/low_geometry/manifest.py new file mode 100644 index 0000000..d6362ef --- /dev/null +++ b/arc_line_vectorization_suede/vectorize/low_geometry/manifest.py @@ -0,0 +1,240 @@ +"""Variable manifest — the parameter vector layout. + +We follow the plan's recommended fallback ("Replace topo-sorted manifest +with implicit constraints"): every primitive owns its own parameters, +and high-weight soft constraints enforce coincidence and on-curve +relationships. This is dramatically simpler to implement and debug +than the handle/topo-sort approach, at the cost of needing larger +weights for the coincidence terms. + +Parameter layout per primitive: + +* ``Line``: 4 params — (x0, y0, x1, y1) +* ``Arc``: 5 params — (x0, y0, x1, y1, bulge) +* ``Circle``: 3 params — (cx, cy, radius) + +The manifest exposes: + +* ``pack(prims)`` / ``unpack(x)`` — convert between primitive list and + the flat parameter vector ``scipy.optimize.least_squares`` consumes. +* ``parameter_scales(prims)`` — per-parameter typical magnitudes, used + to construct an ``x_scale`` array. The optimizer's ``x_scale='jac'`` + works well most of the time but fails on the rare unused parameter, + so we keep an explicit fallback. +""" + +from __future__ import annotations +from dataclasses import dataclass +from typing import List, Tuple + +import numpy as np +from numpy.typing import NDArray + +from .primitives import Arc, Circle, Line, Primitive + + +# --------------------------------------------------------------------------- +# Constraint dataclasses +# +# Each constraint refers to primitives by their index in the manifest's +# ``primitives`` list. Endpoint identifiers are the string "start" or +# "end" (a Circle has neither — junctions onto a circle use OnCircle, +# which doesn't pick an endpoint on the host). + +EndPoint = Tuple[int, str] # (prim_id, "start"|"end") + + +@dataclass +class Coincide: + """Two endpoints must coincide. Used at terminate-into-each-other + junctions and at internal chain joints.""" + + a: EndPoint + b: EndPoint + + +@dataclass +class OnCurve: + """An endpoint of ``terminating`` lies on the curve ``host``. + + For a Line host, this is the perpendicular distance to the line. + For an Arc/Circle host, this is the radial distance to the circle. + """ + + terminating: EndPoint + host: int + + +@dataclass +class G1Smooth: + """Two primitives meeting at an endpoint should have parallel + tangents at that endpoint. + + ``alpha_a`` ∈ {0.0, 1.0} picks which end of primitive ``a`` the + constraint applies to; same for ``b``. + """ + + a: int + alpha_a: float + b: int + alpha_b: float + + +@dataclass +class Parallel: + a: int + b: int + + +@dataclass +class Perpendicular: + a: int + b: int + + +@dataclass +class EqualRadius: + a: int + b: int + + +@dataclass +class Concentric: + a: int + b: int + + +@dataclass +class SoftConstraints: + """Bundle of all soft constraints applied at solve time. Mutable; + beautification mutates the bundle and re-solves.""" + + coincide: List[Coincide] + on_curve: List[OnCurve] + g1: List[G1Smooth] + parallel: List[Parallel] + perpendicular: List[Perpendicular] + equal_radius: List[EqualRadius] + concentric: List[Concentric] + + @classmethod + def empty(cls) -> "SoftConstraints": + return cls([], [], [], [], [], [], []) + + +# --------------------------------------------------------------------------- +# Parameter packing + + +def _param_count(prim: Primitive) -> int: + if isinstance(prim, Line): + return 4 + if isinstance(prim, Arc): + return 5 + if isinstance(prim, Circle): + return 3 + raise TypeError(f"unknown primitive {type(prim)!r}") + + +def parameter_count(prims: List[Primitive]) -> int: + return sum(_param_count(p) for p in prims) + + +def _offsets(prims: List[Primitive]) -> List[int]: + offs: List[int] = [] + cur = 0 + for p in prims: + offs.append(cur) + cur += _param_count(p) + return offs + + +def pack(prims: List[Primitive]) -> NDArray[np.float64]: + """Flatten primitives into a single parameter vector.""" + out = np.empty(parameter_count(prims), dtype=np.float64) + cur = 0 + for p in prims: + if isinstance(p, Line): + out[cur:cur + 2] = p.p0 + out[cur + 2:cur + 4] = p.p1 + cur += 4 + elif isinstance(p, Arc): + out[cur:cur + 2] = p.p0 + out[cur + 2:cur + 4] = p.p1 + out[cur + 4] = p.bulge + cur += 5 + elif isinstance(p, Circle): + out[cur:cur + 2] = p.center + out[cur + 2] = p.radius + cur += 3 + else: + raise TypeError(type(p)) + return out + + +def unpack( + x: NDArray[np.float64], + template: List[Primitive], +) -> List[Primitive]: + """Reconstruct primitives from a parameter vector. ``template`` is + the original primitive list — used to pick which class to build at + each slot. Returns a fresh list of new primitives. + """ + out: List[Primitive] = [] + cur = 0 + for p in template: + if isinstance(p, Line): + out.append(Line(x[cur:cur + 2].copy(), x[cur + 2:cur + 4].copy())) + cur += 4 + elif isinstance(p, Arc): + out.append( + Arc(x[cur:cur + 2].copy(), x[cur + 2:cur + 4].copy(), + float(x[cur + 4])) + ) + cur += 5 + elif isinstance(p, Circle): + out.append(Circle(x[cur:cur + 2].copy(), float(x[cur + 2]))) + cur += 3 + else: + raise TypeError(type(p)) + return out + + +def parameter_scales( + prims: List[Primitive], + pos_scale: float, +) -> NDArray[np.float64]: + """Per-parameter typical magnitude. ``pos_scale`` is e.g. the + bounding-box diagonal of the drawing. + + * Positions: ``pos_scale`` + * Radius: ``pos_scale`` (same units as positions) + * Bulge: 1.0 (dimensionless; ranges roughly [-1, 1] for sweeps ≤π) + """ + out = np.empty(parameter_count(prims), dtype=np.float64) + cur = 0 + for p in prims: + if isinstance(p, Line): + out[cur:cur + 4] = pos_scale + cur += 4 + elif isinstance(p, Arc): + out[cur:cur + 4] = pos_scale + out[cur + 4] = 1.0 + cur += 5 + elif isinstance(p, Circle): + out[cur:cur + 2] = pos_scale + out[cur + 2] = pos_scale + cur += 3 + else: + raise TypeError(type(p)) + return out + + +def endpoint_position(prim: Primitive, end: str) -> NDArray[np.float64]: + """For Line / Arc only. For Circle, callers should use OnCurve + constraints instead.""" + if isinstance(prim, (Line, Arc)): + return prim.p0 if end == "start" else prim.p1 + raise TypeError( + f"endpoint_position not defined for {type(prim).__name__}" + ) diff --git a/arc_line_vectorization_suede/vectorize/low_geometry/primitives.py b/arc_line_vectorization_suede/vectorize/low_geometry/primitives.py new file mode 100644 index 0000000..13178cd --- /dev/null +++ b/arc_line_vectorization_suede/vectorize/low_geometry/primitives.py @@ -0,0 +1,212 @@ +"""Geometric primitives used during vectorization. + +Coordinate convention: image coordinates with x right, y down. +Cross product ``a × b = a[0]*b[1] - a[1]*b[0]``. Positive cross in this +y-down frame corresponds to a clockwise rotation when rendered to a +screen (matching SVG). + +Internal parameterization: + +* ``Line`` — two endpoints. +* ``Arc`` — two endpoints + a scalar ``bulge = tan(sweep / 4)``. Sign of + bulge encodes direction; bulge → 0 degenerates gracefully to a line. + Choice of bulge keeps endpoints first-class which matches the + endpoint-coincidence constraints we add later. +* ``Circle`` — center + radius. Has its own type because + arc-with-coincident-endpoints is singular. +""" + +from __future__ import annotations +from dataclasses import dataclass +from typing import Union + +import numpy as np +from numpy.typing import NDArray + + +_EPS = 1e-9 + + +@dataclass +class Line: + p0: NDArray[np.float64] # shape (2,) + p1: NDArray[np.float64] # shape (2,) + + def length(self) -> float: + return float(np.linalg.norm(self.p1 - self.p0)) + + def direction(self) -> NDArray[np.float64]: + d = self.p1 - self.p0 + n = float(np.linalg.norm(d)) + if n < _EPS: + return np.array([1.0, 0.0]) + return d / n + + def normal(self) -> NDArray[np.float64]: + d = self.direction() + return np.array([-d[1], d[0]]) + + def point_at(self, t: float) -> NDArray[np.float64]: + return (1.0 - t) * self.p0 + t * self.p1 + + def tangent_at(self, t: float) -> NDArray[np.float64]: + return self.direction() + + def perpendicular_distance( + self, pts: NDArray[np.float64] + ) -> NDArray[np.float64]: + """Signed perpendicular distance from each point to the line.""" + n = self.normal() + return (pts - self.p0) @ n + + +@dataclass +class Arc: + p0: NDArray[np.float64] + p1: NDArray[np.float64] + bulge: float + + def chord(self) -> float: + return float(np.linalg.norm(self.p1 - self.p0)) + + def sweep(self) -> float: + """Signed sweep angle, in radians, from p0 to p1.""" + return 4.0 * float(np.arctan(self.bulge)) + + def radius(self) -> float: + L = self.chord() + if abs(self.bulge) < _EPS: + return float("inf") + return L * (1.0 + self.bulge * self.bulge) / (4.0 * abs(self.bulge)) + + def center(self) -> NDArray[np.float64]: + L = self.chord() + if L < _EPS or abs(self.bulge) < _EPS: + # Degenerate; return midpoint to avoid divide-by-zero. + return 0.5 * (self.p0 + self.p1) + mid = 0.5 * (self.p0 + self.p1) + chord_dir = (self.p1 - self.p0) / L + # Normal pointing 90° CCW in y-down frame (i.e. visually 90° CW + # when rendered to screen). The signed `h` selects the side. + normal = np.array([-chord_dir[1], chord_dir[0]]) + h = L * (1.0 - self.bulge * self.bulge) / (4.0 * self.bulge) + return mid + h * normal + + def theta0(self) -> float: + c = self.center() + return float(np.arctan2(self.p0[1] - c[1], self.p0[0] - c[0])) + + def theta1(self) -> float: + return self.theta0() + self.sweep() + + def point_at(self, alpha: float) -> NDArray[np.float64]: + """Point on arc at parameter alpha ∈ [0, 1].""" + c = self.center() + r = self.radius() + if not np.isfinite(r): + return (1.0 - alpha) * self.p0 + alpha * self.p1 + theta = self.theta0() + alpha * self.sweep() + return c + r * np.array([np.cos(theta), np.sin(theta)]) + + def tangent_at(self, alpha: float) -> NDArray[np.float64]: + """Unit tangent at parameter alpha ∈ [0, 1], pointing along + traversal direction (p0 → p1). + """ + sweep = self.sweep() + if abs(sweep) < _EPS: + # Effectively a line. + d = self.p1 - self.p0 + n = float(np.linalg.norm(d)) + return d / n if n > _EPS else np.array([1.0, 0.0]) + c = self.center() + theta = self.theta0() + alpha * sweep + sign = 1.0 if sweep > 0 else -1.0 + return sign * np.array([-np.sin(theta), np.cos(theta)]) + + def ccw(self) -> bool: + return self.sweep() > 0 + + def signed_distance(self, pts: NDArray[np.float64]) -> NDArray[np.float64]: + """For each point, ‖pt − center‖ − radius.""" + c = self.center() + r = self.radius() + if not np.isfinite(r): + # Degenerate: fall back to line perpendicular distance. + return _line_perp(self.p0, self.p1, pts) + return np.linalg.norm(pts - c, axis=1) - r + + def is_full_circle(self) -> bool: + return abs(self.sweep()) > 2.0 * np.pi - 1e-3 + + +@dataclass +class Circle: + center: NDArray[np.float64] + radius: float + + def point_at(self, alpha: float) -> NDArray[np.float64]: + """Point on circle, traversed CCW starting at angle 0.""" + theta = 2.0 * np.pi * alpha + return self.center + self.radius * np.array( + [np.cos(theta), np.sin(theta)] + ) + + def tangent_at(self, alpha: float) -> NDArray[np.float64]: + theta = 2.0 * np.pi * alpha + return np.array([-np.sin(theta), np.cos(theta)]) + + def signed_distance(self, pts: NDArray[np.float64]) -> NDArray[np.float64]: + return np.linalg.norm(pts - self.center, axis=1) - self.radius + + +Primitive = Union[Line, Arc, Circle] + + +# --------------------------------------------------------------------------- + + +def _line_perp( + p0: NDArray[np.float64], + p1: NDArray[np.float64], + pts: NDArray[np.float64], +) -> NDArray[np.float64]: + """Signed perpendicular distance — module-private helper used to keep + ``Arc.signed_distance`` working in the degenerate (zero-bulge) case. + """ + d = p1 - p0 + n = float(np.linalg.norm(d)) + if n < _EPS: + return np.linalg.norm(pts - p0, axis=1) + u = d / n + perp = np.array([-u[1], u[0]]) + return (pts - p0) @ perp + + +def endpoint(prim: Primitive, end: str) -> NDArray[np.float64]: + """Return the start or end position of a primitive. + + For ``Circle`` (full circle) the choice is arbitrary — start and end + coincide. We return the rightmost point on the circle (theta = 0) for + both 'start' and 'end' so routing can detect a closed-loop and emit + a single arc-360° command. + """ + if isinstance(prim, Line): + return prim.p0 if end == "start" else prim.p1 + if isinstance(prim, Arc): + return prim.p0 if end == "start" else prim.p1 + if isinstance(prim, Circle): + return prim.point_at(0.0) + raise TypeError(f"unknown primitive type {type(prim)!r}") + + +def tangent_at_end(prim: Primitive, end: str) -> NDArray[np.float64]: + """Unit tangent at the start or end of a primitive, pointing in the + traversal direction. + """ + if isinstance(prim, Line): + return prim.direction() + if isinstance(prim, Arc): + return prim.tangent_at(0.0 if end == "start" else 1.0) + if isinstance(prim, Circle): + return prim.tangent_at(0.0) + raise TypeError(f"unknown primitive type {type(prim)!r}") diff --git a/arc_line_vectorization_suede/vectorize/low_geometry/residuals.py b/arc_line_vectorization_suede/vectorize/low_geometry/residuals.py new file mode 100644 index 0000000..66f1b64 --- /dev/null +++ b/arc_line_vectorization_suede/vectorize/low_geometry/residuals.py @@ -0,0 +1,311 @@ +"""Residual assembly for the joint least-squares solve. + +``scipy.optimize.least_squares`` minimizes ``½ Σ rᵢ²`` over a flat +parameter vector. We produce the residual vector as a concatenation of: + +1. **Data residuals.** Each fitted primitive sees the source points that + were assigned to it during chain subdivision. For lines the residual + is perpendicular distance; for arcs/circles, signed radial distance. + +2. **Coincidence residuals.** Two endpoints that should coincide + contribute ``(p_a - p_b)`` (two scalars). High weight = exact in + the limit. + +3. **On-curve residuals.** An endpoint of one primitive lying on another + contributes one scalar — perpendicular distance for a line host, + radial distance for an arc/circle host. + +4. **G1 residuals.** Cross product of unit tangents at the meeting + endpoint. ``sin(angle)`` between them — smooth, no branch cut, + exactly the right "small when aligned" quantity. Using ``arctan2`` + differences would introduce wrap-around discontinuities. + +5. **Beautification residuals.** Parallel = cross of directions = 0. + Perpendicular = dot of directions = 0. Equal radius = + difference of radii. Concentric = difference of centers. + +6. **Regularization.** Small log-radius pull-back to the initial fit so + that an underdetermined arc doesn't drift to a giant near-line. +""" + +from __future__ import annotations +from dataclasses import dataclass +from typing import Dict, List + +import numpy as np +from numpy.typing import NDArray + +from .manifest import ( + Coincide, Concentric, EqualRadius, G1Smooth, OnCurve, Parallel, + Perpendicular, SoftConstraints, unpack, endpoint_position, +) +from .primitives import Arc, Circle, Line, Primitive + + +_EPS = 1e-9 + + +@dataclass +class Weights: + """Tunable weight scaling per residual category. + + Defaults reflect the protocol in the implementation recommendations: + coincidence is essentially-hard (very large), G1 is moderate, data + is unit, beautification is mid, regularization is small. + """ + + data: float = 1.0 + coincide: float = 200.0 + on_curve: float = 200.0 + g1: float = 10.0 + parallel: float = 5.0 + perpendicular: float = 5.0 + equal_radius: float = 2.0 + concentric: float = 5.0 + radius_reg: float = 0.05 + # Endpoint anchor: pull prim.p0 toward source_pts[0] and prim.p1 + # toward source_pts[-1]. Data residuals only constrain points to + # lie on the underlying circle/line — they don't fix angular + # position on a circle or position along a line. Without this + # term, Coincide can slide endpoints along the fitted curve as + # long as the chain still closes. Anchor weight should be small + # relative to data (so the fit isn't dragged off the source) but + # strong enough to lock angular position. + endpoint_anchor: float = 5.0 + # Bulge magnitude penalty: prevents an arc with few source points + # from running away to a near-complete circle (|bulge| → many, + # |sweep| → 360°). With the residual = |bulge|, a typical fitted + # bulge of 0.3 contributes ~(0.3)² to the cost, which is + # negligible next to data residuals. But |bulge|=10 contributes + # 100, which the optimizer feels strongly. Below |bulge|=1 + # (sweep < 180°) the penalty is essentially free. + bulge_reg: float = 0.5 + + +# --------------------------------------------------------------------------- +# Per-primitive endpoint and tangent helpers (numerical-stability layer) +# +# We compute these inline rather than calling primitive methods so that +# the residual function works with whatever `unpack` produced — the +# bulge value in the parameter vector may be transiently near zero +# mid-iteration even when the fit ultimately resolves to a curve. + + +def _endpoint(prim: Primitive, end: str) -> NDArray[np.float64]: + if isinstance(prim, (Line, Arc)): + return prim.p0 if end == "start" else prim.p1 + # Circle: no real endpoint; pick theta=0 by convention. + if isinstance(prim, Circle): + return prim.center + np.array([prim.radius, 0.0]) + raise TypeError(type(prim)) + + +def _tangent_at(prim: Primitive, alpha: float) -> NDArray[np.float64]: + """Unit tangent at parameter alpha ∈ [0, 1]. + + Arc.tangent_at can numerically degenerate for very small bulge; in + that case we fall back to the chord direction (the limit as + sweep → 0). This keeps the cross-product G1 residual smooth across + the bulge=0 transition. + """ + if isinstance(prim, Line): + return prim.direction() + if isinstance(prim, Arc): + if abs(prim.bulge) < 1e-3: + d = prim.p1 - prim.p0 + n = float(np.linalg.norm(d)) + if n > _EPS: + return d / n + return np.array([1.0, 0.0]) + return prim.tangent_at(alpha) + if isinstance(prim, Circle): + return prim.tangent_at(alpha) + raise TypeError(type(prim)) + + +def _data_residual( + prim: Primitive, pts: NDArray[np.float64] +) -> NDArray[np.float64]: + if isinstance(prim, Line): + return prim.perpendicular_distance(pts) + if isinstance(prim, Arc): + return prim.signed_distance(pts) + if isinstance(prim, Circle): + return prim.signed_distance(pts) + raise TypeError(type(prim)) + + +def _on_curve_residual( + p: NDArray[np.float64], host: Primitive +) -> NDArray[np.float64]: + """Single scalar: distance from point p to the host curve.""" + if isinstance(host, Line): + return np.array([host.perpendicular_distance(p[None, :])[0]]) + if isinstance(host, Arc): + # Use radial distance to the host's circle. (We don't constrain + # to the arc's sweep range; OnCurve constraints come from the + # junction graph which has already verified the point lies in + # the arc's source-pixel range.) + c = host.center() + r = host.radius() + return np.array([float(np.linalg.norm(p - c) - r)]) + if isinstance(host, Circle): + return np.array( + [float(np.linalg.norm(p - host.center) - host.radius)] + ) + raise TypeError(type(host)) + + +def _radius(prim: Primitive) -> float: + if isinstance(prim, Circle): + return float(prim.radius) + if isinstance(prim, Arc): + return float(prim.radius()) + raise TypeError(type(prim)) + + +def _center(prim: Primitive) -> NDArray[np.float64]: + if isinstance(prim, Circle): + return prim.center + if isinstance(prim, Arc): + return prim.center() + raise TypeError(type(prim)) + + +# --------------------------------------------------------------------------- +# Residual assembly + + +def assemble_residuals( + x: NDArray[np.float64], + template: List[Primitive], + source_points: Dict[int, NDArray[np.float64]], + soft: SoftConstraints, + weights: Weights, + initial_radii: Dict[int, float], +) -> NDArray[np.float64]: + """Build the full residual vector. + + Args: + x: parameter vector to evaluate at. + template: primitives, in manifest order, used to know which + shape to rebuild at each slot. + source_points: ``{prim_id: pts}`` mapping each primitive to + the source points it should fit. Primitives without source + points (e.g. those derived from elsewhere) can be omitted. + soft: bundle of soft constraints. + weights: per-category multiplicative weights. + initial_radii: ``{prim_id: r0}`` for arcs/circles — used by + the log-radius regularization term. + """ + prims = unpack(x, template) + parts: List[NDArray[np.float64]] = [] + + # 1. Data residuals. + for pid, pts in source_points.items(): + if len(pts) == 0: + continue + r = _data_residual(prims[pid], pts) + parts.append(weights.data * r) + + # 1b. Endpoint anchors — pull prim.p0/p1 toward source_pts[0]/[-1]. + # Two 2-vector residuals per primitive (Circle has no endpoints + # and is skipped). This is what gives the optimizer something to + # work against when it would otherwise rotate an arc's endpoints + # along the fitted circle. + for pid, pts in source_points.items(): + if len(pts) == 0: + continue + prim = prims[pid] + if isinstance(prim, (Line, Arc)): + parts.append(weights.endpoint_anchor * (prim.p0 - pts[0])) + parts.append(weights.endpoint_anchor * (prim.p1 - pts[-1])) + + # 2. Coincidence residuals — 2 scalars per pair. + for c in soft.coincide: + pa = _endpoint(prims[c.a[0]], c.a[1]) + pb = _endpoint(prims[c.b[0]], c.b[1]) + parts.append(weights.coincide * (pa - pb)) + + # 3. On-curve residuals — 1 scalar per constraint. + for c in soft.on_curve: + p = _endpoint(prims[c.terminating[0]], c.terminating[1]) + parts.append(weights.on_curve * _on_curve_residual(p, prims[c.host])) + + # 4. G1 residuals — 1 scalar per constraint (cross product). + for c in soft.g1: + ta = _tangent_at(prims[c.a], c.alpha_a) + tb = _tangent_at(prims[c.b], c.alpha_b) + # If alpha_a = 1.0 we're at the "outgoing" end of A pointing + # away from the join; if alpha_b = 0.0 we're at the "incoming" + # end of B pointing into B. They should be parallel (sin = 0). + cross = float(ta[0] * tb[1] - ta[1] * tb[0]) + parts.append(np.array([weights.g1 * cross])) + + # 5. Beautification residuals. + for c in soft.parallel: + a = prims[c.a] + b = prims[c.b] + if not isinstance(a, Line) or not isinstance(b, Line): + continue + da = a.direction() + db = b.direction() + cross = float(da[0] * db[1] - da[1] * db[0]) + parts.append(np.array([weights.parallel * cross])) + + for c in soft.perpendicular: + a = prims[c.a] + b = prims[c.b] + if not isinstance(a, Line) or not isinstance(b, Line): + continue + da = a.direction() + db = b.direction() + dot = float(da @ db) + parts.append(np.array([weights.perpendicular * dot])) + + for c in soft.equal_radius: + ra = _radius(prims[c.a]) + rb = _radius(prims[c.b]) + if not np.isfinite(ra) or not np.isfinite(rb): + continue + parts.append(np.array([weights.equal_radius * (ra - rb)])) + + for c in soft.concentric: + ca = _center(prims[c.a]) + cb = _center(prims[c.b]) + parts.append(weights.concentric * (ca - cb)) + + # 6. Regularization — log-radius pull-back to the initial fit. + # Log so that the residual scales proportionally with relative + # radius change. Without this an arc can drift to absurd radii + # when its source points are nearly colinear. + for pid, r0 in initial_radii.items(): + if r0 <= _EPS: + continue + prim = prims[pid] + if isinstance(prim, (Arc, Circle)): + r = _radius(prim) + if r <= _EPS or not np.isfinite(r): + continue + parts.append(np.array([weights.radius_reg * np.log(r / r0)])) + + # 7. Bulge-magnitude regularization — one residual per Arc. + # Prevents an under-constrained arc (few or near-colinear source + # points) from running away to |bulge| → many (sweep → 360°), + # which the optimizer otherwise has no reason to avoid because + # the data residuals only enforce "on circle", not "shortest arc". + # We use residual = bulge² so the *cost* contribution is ∝ bulge⁴: + # at |bulge| ≤ 1 (sweep ≤ 180°) the penalty is negligible (cost + # ≤ 0.5 * w²), but at |bulge| = 5 it jumps to 0.5*w²*625 and at + # |bulge| = 55 it dominates everything. The Jacobian + # d(b²)/db = 2b pulls back to zero smoothly from either sign. + # Order matters: this MUST appear last so the sparsity in + # solve.py matches. + for pid, prim in enumerate(prims): + if isinstance(prim, Arc): + b = float(prim.bulge) + parts.append(np.array([weights.bulge_reg * b * b])) + + if not parts: + return np.zeros(0) + return np.concatenate(parts) diff --git a/arc_line_vectorization_suede/vectorize/low_geometry/routing.py b/arc_line_vectorization_suede/vectorize/low_geometry/routing.py new file mode 100644 index 0000000..1e980cc --- /dev/null +++ b/arc_line_vectorization_suede/vectorize/low_geometry/routing.py @@ -0,0 +1,427 @@ +"""Sequencing primitives into a pen-up-minimizing tour and emitting +robot commands. + +Graph problem: + +* **Vertices** = endpoint locations (with tolerance-based clustering). + A Circle has no endpoints — it's a "loop edge" attached to a single + virtual vertex at its rightmost point. +* **Edges** = primitives. Each edge is undirected (we can traverse a + Line/Arc in either direction; the bulge / start-end pair gets + flipped if we go end → start). +* **Goal** = minimum-pen-up tour. This is an Eulerian path if the + graph has 0 or 2 odd-degree vertices; otherwise it's the Chinese + Postman problem. NetworkX's ``eulerize`` adds duplicate edges to + fix the parity, which translates back to extra pen-up moves. + +Emitted commands match the robot's interface in ``commands.py``: + +* ``{"kind": "line", "distance": …, "penDown": …}`` — drive forward. +* ``{"kind": "spin", "degrees": …}`` — rotate in place by signed degrees. +* ``{"kind": "arc", "radius": …, "degrees": …}`` — drive an arc. + +Heading convention (per ``commands.py``): positive degrees = CCW in +image coords (which renders as CW on screen). For arcs, ``radius`` is +always positive; the sign of ``degrees`` selects direction. +""" + +from __future__ import annotations +from typing import Dict, List, Optional, Sequence, Tuple + +import math +import numpy as np +import networkx as nx +from numpy.typing import NDArray + +from ...commands import DrawingCommand +from .primitives import Arc, Circle, Line, Primitive, endpoint, tangent_at_end + +# --------------------------------------------------------------------------- +# Endpoint clustering + + +def _endpoint_locations( + prims: Sequence[Primitive], snap_tol: float +) -> Tuple[Dict[int, Tuple[int, int]], List[NDArray[np.float64]]]: + """Cluster all primitive endpoints into a small set of canonical + vertex locations. + + Returns: + edge_to_verts: ``{prim_id: (start_vert, end_vert)}`` + vert_positions: list of (x, y) centroids, indexed by vertex id + + Circles get a synthetic vertex at their theta=0 point on both ends + so they show up as a self-loop in the multigraph. + """ + pts: List[NDArray[np.float64]] = [] + refs: List[Tuple[int, str]] = [] + for pid, p in enumerate(prims): + if isinstance(p, Circle): + v = endpoint(p, "start") # theta = 0 + pts.append(v) + refs.append((pid, "start")) + pts.append(v) + refs.append((pid, "end")) + else: + pts.append(endpoint(p, "start")) + refs.append((pid, "start")) + pts.append(endpoint(p, "end")) + refs.append((pid, "end")) + + # Simple greedy clustering: O(N²) but N is small (<2*n_primitives). + cluster: List[int] = [-1] * len(pts) + centroids: List[List[NDArray[np.float64]]] = [] + for i, p in enumerate(pts): + best = -1 + best_d = snap_tol + for ci, ps in enumerate(centroids): + mean = np.mean(np.stack(ps, axis=0), axis=0) + d = float(np.linalg.norm(p - mean)) + if d < best_d: + best = ci + best_d = d + if best >= 0: + cluster[i] = best + centroids[best].append(p) + else: + cluster[i] = len(centroids) + centroids.append([p]) + + vert_positions = [np.mean(np.stack(ps, axis=0), axis=0) for ps in centroids] + edge_to_verts: Dict[int, Tuple[int, int]] = {} + for i, (pid, end) in enumerate(refs): + vs = edge_to_verts.get(pid, (-1, -1)) + if end == "start": + edge_to_verts[pid] = (cluster[i], vs[1]) + else: + edge_to_verts[pid] = (vs[0], cluster[i]) + return edge_to_verts, vert_positions + + +def _build_multigraph( + edge_to_verts: Dict[int, Tuple[int, int]], + prims: Sequence[Primitive], +) -> nx.MultiGraph: + G = nx.MultiGraph() + for pid, (u, v) in edge_to_verts.items(): + # Include all vertices, even if isolated, so eulerize() works. + G.add_node(u) + G.add_node(v) + # Edge weight = arc-length-ish, useful for eulerize() to prefer + # short duplications. + p = prims[pid] + if isinstance(p, Line): + w = p.length() + elif isinstance(p, Arc): + w = abs(p.sweep()) * ( + p.radius() if math.isfinite(p.radius()) else p.chord() + ) + elif isinstance(p, Circle): + w = 2.0 * math.pi * p.radius + else: + w = 0.0 + G.add_edge(u, v, key=pid, weight=float(w)) + return G + + +def _pick_start_vertex( + G: nx.MultiGraph, + vert_positions: List[NDArray[np.float64]], + start_pos: NDArray[np.float64], +) -> int: + """For an Eulerian *path* graph (exactly two odd vertices), the path + must begin at one of them. Pick the odd vertex closer to the robot's + starting position; otherwise pick the vertex closest to start_pos. + """ + odd = [v for v in G.nodes if G.degree(v) % 2 == 1] + candidates = odd if odd else list(G.nodes) + if not candidates: + return 0 + best = candidates[0] + best_d = float(np.linalg.norm(vert_positions[best] - start_pos)) + for v in candidates[1:]: + d = float(np.linalg.norm(vert_positions[v] - start_pos)) + if d < best_d: + best = v + best_d = d + return best + + +def _eulerize_per_component(G: nx.MultiGraph) -> nx.MultiGraph: + """Apply ``nx.eulerize`` per connected component, then return the + union. NetworkX 3.x's ``eulerize`` requires a connected graph; for + drawings with multiple disconnected groups we must split first. + """ + out = nx.MultiGraph() + for comp_nodes in nx.connected_components(G): + sub = G.subgraph(comp_nodes).copy() + if sub.number_of_edges() == 0: + out.add_nodes_from(sub.nodes()) + continue + if not nx.is_eulerian(sub) and not nx.has_eulerian_path(sub): + sub = nx.eulerize(sub) + out = nx.compose(out, sub) + return out + + +def order_primitives( + prims: Sequence[Primitive], + start_pos: NDArray[np.float64], + snap_tol: float = 1.0, +) -> List[Tuple[int, bool]]: + """Return a tour ``[(prim_id, reverse), ...]`` that visits every + primitive exactly once, with pen-up jumps inserted between + disconnected components and parity-fix duplications when needed. + + ``reverse`` is True if the primitive should be traversed end → + start. Lines and arcs reverse cleanly; circles ignore the flag. + """ + if not prims: + return [] + + edge_to_verts, vert_positions = _endpoint_locations(prims, snap_tol) + G = _build_multigraph(edge_to_verts, prims) + G2 = _eulerize_per_component(G) + + # Within each connected component, run an Eulerian path / circuit. + tour: List[Tuple[int, int, int]] = [] # (u, v, key) + + # Process components in an order that minimizes pen-up travel from + # the current position. + cur_pos = np.asarray(start_pos, dtype=float) + components = list(nx.connected_components(G2)) + used_components: List[bool] = [False] * len(components) + while True: + # Pick the unused component whose closest vertex is nearest to + # cur_pos. + best_ci = -1 + best_d = float("inf") + best_start = -1 + for ci, comp in enumerate(components): + if used_components[ci]: + continue + sub = G2.subgraph(comp).copy() + if sub.number_of_edges() == 0: + continue + start_v = _pick_start_vertex(sub, vert_positions, cur_pos) + d = float(np.linalg.norm(vert_positions[start_v] - cur_pos)) + if d < best_d: + best_d = d + best_ci = ci + best_start = start_v + if best_ci < 0: + break + + sub = G2.subgraph(components[best_ci]).copy() + used_components[best_ci] = True + + # Eulerian path / circuit on this component. + if nx.is_eulerian(sub): + path = list(nx.eulerian_circuit(sub, source=best_start, keys=True)) + else: + path = list(nx.eulerian_path(sub, source=best_start, keys=True)) + + for u, v, key in path: + tour.append((u, v, key)) + cur_pos = vert_positions[v] + + # Resolve direction per edge. Synthetic eulerize-added edges have + # keys not present in our original primitives → those become + # pen-up jumps and we skip them in the emitted command sequence. + original_keys = set(edge_to_verts.keys()) + out: List[Tuple[int, bool]] = [] + for u, v, key in tour: + if key not in original_keys: + # Synthetic edge (parity fix); ignored — the pen-up handling + # in to_commands() will emit a pen-up + move when needed. + continue + orig_u, orig_v = edge_to_verts[key] + reverse = u == orig_v and v == orig_u and orig_u != orig_v + out.append((key, reverse)) + return out + + +# --------------------------------------------------------------------------- +# Command emission + + +def _wrap_to_pi(angle: float) -> float: + """Wrap a radian angle to (-π, π].""" + return (angle + math.pi) % (2.0 * math.pi) - math.pi + + +def _reverse_primitive(prim: Primitive) -> Primitive: + """Return a primitive traversed in the opposite direction.""" + if isinstance(prim, Line): + return Line(prim.p1.copy(), prim.p0.copy()) + if isinstance(prim, Arc): + # Flip endpoints AND bulge sign — same arc, reversed traversal. + return Arc(prim.p1.copy(), prim.p0.copy(), -prim.bulge) + if isinstance(prim, Circle): + # Circles are loops; "reverse" doesn't change anything. + return prim + raise TypeError(type(prim)) + + +def _heading_change_deg(from_heading: float, to_heading: float) -> float: + """Signed degrees needed to spin from one heading to another, picking + the shortest rotation. + """ + delta = _wrap_to_pi(to_heading - from_heading) + return float(math.degrees(delta)) + + +def _heading_of(vec: NDArray[np.float64]) -> float: + return float(math.atan2(vec[1], vec[0])) + + +def to_commands( + prims: Sequence[Primitive], + tour: List[Tuple[int, bool]], + start_pos: NDArray[np.float64], + start_heading: float, # radians + pen_up_join_tol: float = 0.5, +) -> Sequence[DrawingCommand]: + """Emit a robot command sequence for the ordered tour. + + A pen-down move is emitted as a sequence: optional spin to align + with the primitive's start tangent, then either a line (with pen + down) or an arc command. Pen-up jumps between disconnected + components are emitted as: pen-up line + spin/arc to align with + the next primitive's start. + """ + cmds: Sequence[DrawingCommand] = [] + cur_pos = np.asarray(start_pos, dtype=float) + cur_heading = float(start_heading) + + for k, (pid, reverse) in enumerate(tour): + prim = prims[pid] + if reverse: + prim = _reverse_primitive(prim) + + # Pen position at primitive start. + if isinstance(prim, Circle): + start = prim.point_at(0.0) + else: + start = endpoint(prim, "start") + + # If the pen isn't already at the primitive's start, jump + # there with a pen-up line. Pen-up still requires correct + # heading first. + gap = float(np.linalg.norm(start - cur_pos)) + if gap > pen_up_join_tol: + target_heading = _heading_of(start - cur_pos) + spin_deg = _heading_change_deg(cur_heading, target_heading) + if abs(spin_deg) > 1e-3: + cmds.append({"kind": "spin", "degrees": float(spin_deg)}) + cur_heading = target_heading + cmds.append({"kind": "line", "distance": gap, "penDown": False}) + cur_pos = start.copy() + + # Decide whether to emit this primitive as an arc or a line. + # This MUST be settled before the alignment spin: a major arc + # (sweep > 180°) has a start tangent pointing nearly opposite + # its chord, so if we spin to the arc's tangent and then emit + # a chord-line, the simulator drives backwards relative to the + # actual chord direction. Pick the spin target based on what + # we're actually emitting. + emit_as_line: Optional[float] = None # distance, if degraded to line + if isinstance(prim, Arc): + sweep_deg = float(math.degrees(prim.sweep())) + r = prim.radius() + sweep_rad = prim.sweep() + if math.isfinite(r) and r > 0: + sagitta = abs(r * (1.0 - math.cos(sweep_rad / 2.0))) + chord_over_r = prim.chord() / r + else: + sagitta = 0.0 + chord_over_r = 0.0 + # An arc whose sweep approaches 360° with a chord far + # smaller than the diameter is an optimizer runaway: the + # primitive is trying to be a closed circle but Arc isn't + # the right type for that. A legitimate 350° arc would + # still have a chord of about 0.17·r; a runaway with + # bulge=55 has chord/r ≈ 0.04. Only when BOTH conditions + # hold do we degrade. + degenerate = ( + not math.isfinite(r) + or r > 1e6 + or sagitta < 0.5 + or (abs(sweep_deg) > 270.0 and chord_over_r < 0.15) + ) + if degenerate: + emit_as_line = float(prim.chord()) + + # Pick the heading we should face before emitting. For a line + # or a degraded arc, this is the chord (p1 - p0) direction. + # For a true arc, it's the arc's tangent at p0. + if isinstance(prim, Line) or emit_as_line is not None: + chord_vec = prim.p1 - prim.p0 + primitive_start_heading = _heading_of(chord_vec) + elif isinstance(prim, Arc): + start_tan = tangent_at_end(prim, "start") + primitive_start_heading = _heading_of(start_tan) + else: # Circle + # The 360° arc command emitted below assumes the robot is + # heading along the CIRCLE's tangent at point_at(0.0). The + # simulator places the circle's center at heading ± π/2; + # if the heading isn't perpendicular to the radius at the + # start point, the simulator draws the circle in the wrong + # place. Set the heading to the circle's tangent at the + # starting point, going CCW (matches the +360° sweep we + # emit, and the convention used by Arc.tangent_at). + start_tan = prim.tangent_at(0.0) + primitive_start_heading = _heading_of(start_tan) + + spin_deg = _heading_change_deg(cur_heading, primitive_start_heading) + if abs(spin_deg) > 1e-3: + cmds.append({"kind": "spin", "degrees": float(spin_deg)}) + cur_heading = primitive_start_heading + + # Emit the primitive. + if isinstance(prim, Line): + d = prim.length() + if d > 1e-9: + cmds.append({"kind": "line", "distance": float(d), "penDown": True}) + cur_pos = prim.p1.copy() + cur_heading = primitive_start_heading + elif isinstance(prim, Arc): + if emit_as_line is not None: + if emit_as_line > 1e-9: + cmds.append( + { + "kind": "line", + "distance": float(emit_as_line), + "penDown": True, + } + ) + cur_pos = prim.p1.copy() + cur_heading = primitive_start_heading + else: + cmds.append( + { + "kind": "arc", + "radius": float(r), + "degrees": float(sweep_deg), + } + ) + cur_pos = prim.p1.copy() + # Heading at end of arc: + end_tan = tangent_at_end(prim, "end") + cur_heading = _heading_of(end_tan) + elif isinstance(prim, Circle): + # Full circle. Drive a 360° arc (always CCW in image, which + # is visually CW; sign matches Arc convention). Heading + # returns to start so we leave it alone. + cmds.append( + { + "kind": "arc", + "radius": float(prim.radius), + "degrees": 360.0, + } + ) + # Pen position is unchanged (full loop). + else: + raise TypeError(type(prim)) + + return cmds diff --git a/arc_line_vectorization_suede/vectorize/low_geometry/solve.py b/arc_line_vectorization_suede/vectorize/low_geometry/solve.py new file mode 100644 index 0000000..fba7484 --- /dev/null +++ b/arc_line_vectorization_suede/vectorize/low_geometry/solve.py @@ -0,0 +1,641 @@ +"""Top-level orchestration: from polylines + junctions to fitted primitives. + +The flow: + +1. For each polyline, run ``fit_polyline`` to get a chain of primitives + plus per-primitive source-point ranges. +2. Lay out a flat manifest indexing each primitive globally. +3. Build the constraint bundle from the StrokeGraph's junctions: + * Internal chain joints → Coincide + G1 + * Terminate-into-each-other junctions → Coincide between participating endpoints + * Terminate-on (mixed) junctions → OnCurve from terminating endpoint to host + * Smooth (low-deflection) junctions where two strokes meet → G1 +4. Solve once → ``primitives_fitted``. +5. Run beautification on the result, add candidates as soft constraints, + solve again → ``primitives_consolidated``. + +The build is incremental — each step's output is exposed on the result +dataclass so the caller can render diagnostics. +""" + +from __future__ import annotations +from dataclasses import dataclass, field +from typing import Dict, List, Optional, Tuple + +import numpy as np +from numpy.typing import NDArray +from scipy.optimize import least_squares +from scipy.sparse import lil_matrix + +from ...graph import Junction, Role, StrokeGraph + +from .beautify import BeautifyTolerances, detect, merge_into +from .fitting import ChainPiece, fit_polyline, is_closed_polyline +from .manifest import ( + Coincide, + G1Smooth, + OnCurve, + SoftConstraints, + pack, + parameter_count, + parameter_scales, + unpack, +) +from .primitives import Arc, Circle, Line, Primitive +from .residuals import Weights, assemble_residuals + +# --------------------------------------------------------------------------- +# Result containers + + +@dataclass +class FittedSegment: + """The chain output for one polyline, with global primitive IDs.""" + + polyline_index: int + pieces: List[ChainPiece] + primitive_ids: List[int] # global index in the manifest, per piece + source_points: List[NDArray[np.float64]] # one array per piece + + +@dataclass +class SolveResult: + primitives: List[Primitive] + fitted_segments: List[FittedSegment] + soft: SoftConstraints + converged: bool + cost: float + n_iters: int + + +# --------------------------------------------------------------------------- +# Configuration + + +@dataclass +class FitConfig: + line_tol: float = 0.005 # relative to segment bbox diagonal + arc_tol: float = 0.012 + lam_rel: float = 4.0 # MDL penalty multiplier + min_len: int = 5 + use_dp: bool = False # top-down is faster & usually cleaner + max_window: int = 256 + closed_tol: float = 1.5 + # G1 auto-detection at smooth segment-to-segment junctions: + smooth_junction_deg_threshold: float = 25.0 + # Source points are subsampled (uniformly along the chain piece) + # before being handed to the joint solver. Skeletal polylines from + # the upstream are pixel-dense (~1000+ pts/polyline) and that scale + # is well beyond what the data residual needs to nail down a + # primitive — it just makes the optimizer slow. ~30 pts per + # primitive captures shape with sub-pixel residual mean. + source_points_per_primitive: int = 30 + # Cap on the working polyline length passed to ``fit_polyline``. + # Top-down's min_len is in raw indices; on a 1000+ pt polyline + # that means it can fragment a 5 px wiggle into a tiny piece + # while the rest of the stroke fits as one arc. Subsampling first + # makes min_len correspond to a stroke-length fraction instead. + polyline_subsample_cap: int = 200 + + +@dataclass +class SolveConfig: + weights: Weights = field(default_factory=Weights) + # Diminishing returns past ~60-100 iters; capping here keeps the + # worst-case drawing under ~30s for the joint solve. + max_iters: int = 80 + method: str = "trf" # 'trf' supports bounds; 'lm' is faster but unbounded + f_scale: float = 1.0 + + +# --------------------------------------------------------------------------- +# Pipeline + + +def build_chains( + graph: StrokeGraph, + config: FitConfig, +) -> List[FittedSegment]: + """Fit each polyline independently into a chain of primitives.""" + out: List[FittedSegment] = [] + for pi, raw_poly in enumerate(graph.polylines): + if len(raw_poly) < 2: + continue + # Subsample the polyline so that ``min_len`` (in indices) and + # the recursion's "split at worst residual" both operate on a + # stroke-length-fraction scale rather than raw pixel density. + cap_poly = max(2, int(config.polyline_subsample_cap)) + if len(raw_poly) > cap_poly: + sub_idx = np.linspace(0, len(raw_poly) - 1, cap_poly, dtype=int) + poly = raw_poly[sub_idx].copy() + else: + poly = raw_poly + chain = fit_polyline( + poly, + line_tol=config.line_tol, + arc_tol=config.arc_tol, + lam_rel=config.lam_rel, + min_len=config.min_len, + closed_tol=config.closed_tol, + use_dp=config.use_dp, + max_window=config.max_window, + ) + if not chain: + continue + # Source points per piece: subsample further if the piece is + # bigger than the per-primitive cap. The piece indices already + # refer to the subsampled polyline. + src: List[NDArray[np.float64]] = [] + cap = max(2, int(config.source_points_per_primitive)) + for c in chain: + piece = poly[c.start_idx : c.end_idx] + if len(piece) <= cap: + src.append(piece.copy()) + else: + idx = np.linspace(0, len(piece) - 1, cap, dtype=int) + src.append(piece[idx].copy()) + out.append( + FittedSegment( + polyline_index=pi, + pieces=chain, + primitive_ids=[], # filled in below by ``assign_global_ids`` + source_points=src, + ) + ) + return out + + +def assign_global_ids(segments: List[FittedSegment]) -> List[Primitive]: + """Lay out a flat primitive list with global IDs and populate + ``segment.primitive_ids``. Returns the resulting primitive list. + """ + prims: List[Primitive] = [] + for seg in segments: + seg.primitive_ids = [] + for piece in seg.pieces: + seg.primitive_ids.append(len(prims)) + prims.append(piece.primitive) + return prims + + +def _chain_endpoint(seg: FittedSegment, end: str) -> Tuple[int, str]: + """Return ``(global_prim_id, "start"|"end")`` for the segment's + external start/end.""" + if end == "start": + return seg.primitive_ids[0], "start" + return seg.primitive_ids[-1], "end" + + +def _internal_joint_constraints( + seg: FittedSegment, + smooth_g1: bool = True, +) -> Tuple[List[Coincide], List[G1Smooth]]: + """Coincide + G1 between consecutive sub-primitives in a chain. + + Internal joints lie inside one fluid stroke, so they should be + smooth by default. The G1 can be suppressed if downstream knows + a corner sits at this index (rare). + """ + coincide: List[Coincide] = [] + g1: List[G1Smooth] = [] + pids = seg.primitive_ids + for k in range(len(pids) - 1): + coincide.append(Coincide((pids[k], "end"), (pids[k + 1], "start"))) + if smooth_g1: + g1.append(G1Smooth(a=pids[k], alpha_a=1.0, b=pids[k + 1], alpha_b=0.0)) + return coincide, g1 + + +def _polyline_endpoint_to_chain_end(poly_n: int, point_index: int) -> Optional[str]: + """Map a Participation's point_index to the chain end ("start" or + "end") if it sits at one of the polyline endpoints; None otherwise. + """ + if point_index == 0: + return "start" + if point_index == poly_n - 1: + return "end" + return None + + +def _resolve_host_subprimitive( + junction_xy: NDArray[np.float64], + seg: FittedSegment, + polyline: NDArray[np.float64], +) -> int: + """For a terminate-on junction landing on segment ``seg``, find which + sub-primitive in the chain the junction sits closest to. Return the + global primitive ID. + + Why we use source_points rather than the raw polyline: ``Vectorize`` + subsamples each polyline (default cap 200 pts) before running + ``fit_polyline``, so ``seg.pieces``' ``start_idx``/``end_idx`` + are indices into the SUBSAMPLED polyline, not the raw polyline. + Using the raw polyline's argmin would land us at an index in + full-polyline space that doesn't correspond to chain-piece + coordinates (the housesun apex junction landed at full-idx 183, + which falls in chain-piece [138:200] in subsampled space and + returns the FLOOR primitive instead of the left roof slope). + ``source_points`` per piece are the actual data the primitive was + fit to, so finding the piece whose source_points are closest to + the junction location is robust to any subsampling layer. + """ + # If source_points are available, use them — they live in the same + # coordinate space as the junction location. + if seg.source_points and all(len(sp) > 0 for sp in seg.source_points): + best_pid = seg.primitive_ids[0] + best_dist = float("inf") + for k, src in enumerate(seg.source_points): + d = float(np.min(np.linalg.norm(src - junction_xy, axis=1))) + if d < best_dist: + best_dist = d + best_pid = seg.primitive_ids[k] + return best_pid + # Fallback: closest-source-point index into ``polyline``. Assumes + # ``polyline``'s indexing matches ``seg.pieces``' indexing — only + # safe when no subsampling happened upstream. + dists = np.linalg.norm(polyline - junction_xy, axis=1) + i_nearest = int(np.argmin(dists)) + for k, piece in enumerate(seg.pieces): + if piece.start_idx <= i_nearest < piece.end_idx: + return seg.primitive_ids[k] + return seg.primitive_ids[-1] + + +def build_junction_constraints( + graph: StrokeGraph, + segments: List[FittedSegment], + smooth_junction_deg: float, + primitives: Optional[List[Primitive]] = None, +) -> SoftConstraints: + """Translate the StrokeGraph's junctions into soft constraints. + + Rules: + + * If two or more **terminal** participants meet at one junction, + their chain ends are made coincident (pairwise to first). + Exception: when the "base" endpoint sits on a Circle primitive, + use OnCurve for the other endpoints instead — a Circle has no + meaningful endpoint (the convention theta=0 is arbitrary), and + a hard Coincide on that arbitrary point pulls the entire circle + to satisfy the constraint. Putting the other endpoints + ON THE PERIMETER respects the geometry instead. + * If a **terminal** participant meets one or more **crossing** / + **cusp** (interior) participants, the terminal end lies on the + host's sub-primitive (OnCurve). + * If the junction's pair-deflection is below + ``smooth_junction_deg``, a G1 constraint is added between the + participating strokes. + """ + seg_by_poly: Dict[int, FittedSegment] = {s.polyline_index: s for s in segments} + out = SoftConstraints.empty() + + def is_circle(pid: int) -> bool: + if primitives is None: + return False + return 0 <= pid < len(primitives) and isinstance(primitives[pid], Circle) + + for j in graph.junctions: + # Filter participants whose polyline made it into our segment + # set (a polyline could be too short for fit_polyline and got + # dropped). + live = [p for p in j.participants if p.polyline_index in seg_by_poly] + if len(live) < 2: + continue + + # Bucket by terminal-vs-interior. + terminals = [p for p in live if p.role == Role.TERMINAL] + interiors = [p for p in live if p.role != Role.TERMINAL] + + # ------------------------------------------------------------ + # Coincidence: all terminal participants at this junction. + # When the base endpoint lies on a Circle primitive, use + # OnCurve instead of Coincide for the other endpoints. + # ------------------------------------------------------------ + if len(terminals) >= 2: + base = terminals[0] + base_seg = seg_by_poly[base.polyline_index] + base_end = _polyline_endpoint_to_chain_end( + len(graph.polylines[base.polyline_index]), base.point_index + ) + if base_end is not None: + base_ep = _chain_endpoint(base_seg, base_end) + base_is_circle = is_circle(base_ep[0]) + for p in terminals[1:]: + seg = seg_by_poly[p.polyline_index] + end = _polyline_endpoint_to_chain_end( + len(graph.polylines[p.polyline_index]), p.point_index + ) + if end is None: + continue + ep = _chain_endpoint(seg, end) + if base_is_circle: + # Other terminal must lie on the circle. + out.on_curve.append(OnCurve(terminating=ep, host=base_ep[0])) + elif is_circle(ep[0]): + # The "other" terminal is the circle; the base + # must lie on it. + out.on_curve.append(OnCurve(terminating=base_ep, host=ep[0])) + else: + out.coincide.append(Coincide(base_ep, ep)) + + # ------------------------------------------------------------ + # On-curve: each terminal's endpoint lies on each interior's + # host sub-primitive. We pick the FIRST interior host so we + # don't over-constrain (the interior strokes are roughly + # tangent to each other at the junction; one constraint is + # enough to keep the terminal pinned). + # ------------------------------------------------------------ + if terminals and interiors: + host_part = interiors[0] + host_seg = seg_by_poly[host_part.polyline_index] + host_poly = graph.polylines[host_part.polyline_index] + host_id = _resolve_host_subprimitive(j.location, host_seg, host_poly) + for term in terminals: + seg = seg_by_poly[term.polyline_index] + end = _polyline_endpoint_to_chain_end( + len(graph.polylines[term.polyline_index]), term.point_index + ) + if end is None: + continue + ep = _chain_endpoint(seg, end) + out.on_curve.append(OnCurve(terminating=ep, host=host_id)) + + # ------------------------------------------------------------ + # G1: at pure-terminal junctions where exactly two strokes + # meet with a small deflection, add a smoothing constraint + # between their chain-end sub-primitives. + # ------------------------------------------------------------ + if len(terminals) == 2 and not interiors: + a, b = terminals + # Pair deflection: arccos(-tangent_a · tangent_b) — both + # tangent_ins point INTO the polyline body, so an aligned + # "smooth corner" has tangent_a ≈ -tangent_b. + cos = float(np.clip(-a.tangent_in @ b.tangent_in, -1.0, 1.0)) + defl = float(np.degrees(np.arccos(cos))) + if defl < smooth_junction_deg: + seg_a = seg_by_poly[a.polyline_index] + seg_b = seg_by_poly[b.polyline_index] + end_a = _polyline_endpoint_to_chain_end( + len(graph.polylines[a.polyline_index]), a.point_index + ) + end_b = _polyline_endpoint_to_chain_end( + len(graph.polylines[b.polyline_index]), b.point_index + ) + if end_a is not None and end_b is not None: + pid_a, _ = _chain_endpoint(seg_a, end_a) + pid_b, _ = _chain_endpoint(seg_b, end_b) + # alpha = 1 means "outgoing tangent at end" — but + # since one of these is a chain-start, we want to + # flip the sign convention by choosing alpha + # appropriately. + alpha_a = 1.0 if end_a == "end" else 0.0 + alpha_b = 0.0 if end_b == "start" else 1.0 + out.g1.append( + G1Smooth(a=pid_a, alpha_a=alpha_a, b=pid_b, alpha_b=alpha_b) + ) + + return out + + +# --------------------------------------------------------------------------- + + +# --------------------------------------------------------------------------- +# Jacobian sparsity +# +# The residual function evaluates ~600-3000 residuals over ~100-300 +# parameters. Each residual only depends on a handful of parameters +# (the params of one or two primitives). Without a sparsity pattern, +# scipy's TRF finite-difference Jacobian does N_params extra residual +# evaluations per iteration — a 10-50x speedup is on the table. + + +def _param_slice(template: List[Primitive], pid: int) -> Tuple[int, int]: + """Return ``(start, stop)`` parameter indices for primitive ``pid``.""" + cur = 0 + for i, p in enumerate(template): + n = ( + 4 + if isinstance(p, Line) + else 5 if isinstance(p, Arc) else 3 if isinstance(p, Circle) else 0 + ) + if i == pid: + return cur, cur + n + cur += n + raise IndexError(pid) + + +def _build_jacobian_sparsity( + template: List[Primitive], + source_points: Dict[int, NDArray[np.float64]], + soft: SoftConstraints, + initial_radii: Dict[int, float], +) -> "lil_matrix": + """Build the sparsity pattern matching ``assemble_residuals``. + + Layout must mirror the residual block order in residuals.py exactly: + 1. data (per primitive with source points, in dict-iteration order) + 2. coincide (2 scalars each) + 3. on_curve (1 scalar each) + 4. g1 (1 scalar each) + 5. parallel (1 scalar each) + 6. perpendicular (1 scalar each) + 7. equal_radius (1 scalar each) + 8. concentric (2 scalars each) + 9. radius regularization (1 scalar per arc/circle with init radius) + 10. bulge regularization (1 scalar per Arc) + """ + n_params = parameter_count(template) + slices = {i: _param_slice(template, i) for i in range(len(template))} + + rows: List[Tuple[int, int]] = [] # (n_rows, depends_on_primitive_id_or_pair) + + # 1. data + data_meta: List[Tuple[int, int]] = [] + for pid, pts in source_points.items(): + if len(pts) == 0: + continue + data_meta.append((pid, len(pts))) + + # 1b. endpoint anchors — 4 scalar rows per Line/Arc with source pts + anchor_meta: List[int] = [] + for pid, pts in source_points.items(): + if len(pts) == 0: + continue + if isinstance(template[pid], (Line, Arc)): + anchor_meta.append(pid) + + # 10. bulge reg — one row per Arc + bulge_arc_ids = [i for i, p in enumerate(template) if isinstance(p, Arc)] + + n_rows = ( + sum(k for _, k in data_meta) + + 4 * len(anchor_meta) + + 2 * len(soft.coincide) + + len(soft.on_curve) + + len(soft.g1) + + len(soft.parallel) + + len(soft.perpendicular) + + len(soft.equal_radius) + + 2 * len(soft.concentric) + + len(initial_radii) + + len(bulge_arc_ids) + ) + + J = lil_matrix((n_rows, n_params), dtype=np.float64) + + row = 0 + # 1. data residuals + for pid, k in data_meta: + s, e = slices[pid] + for _ in range(k): + J[row, s:e] = 1 + row += 1 + # 1b. anchor residuals — depend on prim's p0/p1 only (first 4 params + # of Line/Arc). To keep it simple and correct, mark the whole slice. + for pid in anchor_meta: + s, e = slices[pid] + for _ in range(4): + J[row, s:e] = 1 + row += 1 + # 2. coincide (2 scalars per pair, depends on endpoints of a & b) + for c in soft.coincide: + sa, ea = slices[c.a[0]] + sb, eb = slices[c.b[0]] + for _ in range(2): + J[row, sa:ea] = 1 + J[row, sb:eb] = 1 + row += 1 + # 3. on-curve (1 scalar per constraint, depends on point primitive + host) + for c in soft.on_curve: + sa, ea = slices[c.terminating[0]] + sb, eb = slices[c.host] + J[row, sa:ea] = 1 + J[row, sb:eb] = 1 + row += 1 + # 4. G1 (1 scalar per constraint, depends on params of both primitives) + for c in soft.g1: + sa, ea = slices[c.a] + sb, eb = slices[c.b] + J[row, sa:ea] = 1 + J[row, sb:eb] = 1 + row += 1 + # 5. parallel + for c in soft.parallel: + sa, ea = slices[c.a] + sb, eb = slices[c.b] + J[row, sa:ea] = 1 + J[row, sb:eb] = 1 + row += 1 + # 6. perpendicular + for c in soft.perpendicular: + sa, ea = slices[c.a] + sb, eb = slices[c.b] + J[row, sa:ea] = 1 + J[row, sb:eb] = 1 + row += 1 + # 7. equal radius + for c in soft.equal_radius: + sa, ea = slices[c.a] + sb, eb = slices[c.b] + J[row, sa:ea] = 1 + J[row, sb:eb] = 1 + row += 1 + # 8. concentric (2 rows per pair) + for c in soft.concentric: + sa, ea = slices[c.a] + sb, eb = slices[c.b] + for _ in range(2): + J[row, sa:ea] = 1 + J[row, sb:eb] = 1 + row += 1 + # 9. radius regularization + for pid in initial_radii.keys(): + s, e = slices[pid] + J[row, s:e] = 1 + row += 1 + # 10. bulge regularization — depends only on the bulge param of the arc + # (it's the 5th param in the arc's 5-param slot). For simplicity mark + # the whole arc slice; over-marking is harmless, only under-marking + # would break the optimizer. + for pid in bulge_arc_ids: + s, e = slices[pid] + J[row, s:e] = 1 + row += 1 + + assert row == n_rows, f"sparsity row count mismatch: {row} vs {n_rows}" + return J + + +# --------------------------------------------------------------------------- + + +def solve_once( + primitives: List[Primitive], + source_points: Dict[int, NDArray[np.float64]], + soft: SoftConstraints, + weights: Weights, + pos_scale: float, + max_iters: int = 200, + method: str = "trf", +) -> Tuple[List[Primitive], SolveResult]: + """Run one least-squares solve. Returns the (primitives, result).""" + if not primitives: + return [], SolveResult([], [], soft, True, 0.0, 0) + + template = primitives + x0 = pack(template) + + # Pre-compute initial radii so the log-radius regularization can + # anchor the curved primitives to their starting fits. + initial_radii: Dict[int, float] = {} + for i, p in enumerate(template): + if isinstance(p, Circle): + initial_radii[i] = max(1e-3, float(p.radius)) + elif isinstance(p, Arc): + r = p.radius() + if np.isfinite(r): + initial_radii[i] = max(1e-3, float(r)) + + def f(x): + return assemble_residuals( + x, template, source_points, soft, weights, initial_radii + ) + + x_scale = parameter_scales(template, pos_scale=pos_scale) + jac_sparsity = ( + _build_jacobian_sparsity(template, source_points, soft, initial_radii) + if method == "trf" + else None + ) + + sol = least_squares( + f, + x0, + method=method, + x_scale=x_scale, + max_nfev=max_iters, + ftol=1e-7, + xtol=1e-7, + gtol=1e-7, + jac_sparsity=jac_sparsity, + ) + + new_prims = unpack(sol.x, template) + res = SolveResult( + primitives=new_prims, + fitted_segments=[], # filled in by caller + soft=soft, + converged=bool(sol.success), + cost=float(sol.cost), + n_iters=int(sol.nfev), + ) + return new_prims, res + + +def _bbox_diag(graph: StrokeGraph) -> float: + pts = np.concatenate(graph.polylines, axis=0) if graph.polylines else None + if pts is None or len(pts) == 0: + return 100.0 + span = pts.max(axis=0) - pts.min(axis=0) + return max(float(np.linalg.norm(span)), 10.0) diff --git a/arc_line_vectorization_suede/visualize.py b/arc_line_vectorization_suede/visualize.py new file mode 100644 index 0000000..61b176c --- /dev/null +++ b/arc_line_vectorization_suede/visualize.py @@ -0,0 +1,505 @@ +"""Render a DrawingCommand list to SVG. + +We simulate the robot to recover the drawn primitives (line segments and +arcs) in order, then write an SVG with each *drawn* primitive coloured by +its index in the draw sequence (rainbow / hue ramp). Pen-up traversals are +shown as faint dashed grey lines so you can sanity-check ordering. +""" + +from __future__ import annotations +import math +from typing import List, Optional, Tuple, Sequence + +import numpy as np +from PIL import Image, ImageDraw + +from .commands import DrawingCommand + + +def _hsl(hue_deg: float, sat: float = 80.0, light: float = 50.0) -> str: + return f"hsl({hue_deg:.1f}, {sat:.0f}%, {light:.0f}%)" + + +def _simulate( + commands: Sequence[DrawingCommand], + start_pos: Tuple[float, float], + start_heading: float, +): + """Replay commands; collect drawn segments, pen-up segments, and bounds.""" + pos = np.array(start_pos, dtype=float) + heading = float(start_heading) + drawn = [] # list of dicts {kind, ...} in draw order + pen_up = [] # list of (p0, p1) tuples + xs, ys = [pos[0]], [pos[1]] + + def update_bounds(px, py): + xs.append(px) + ys.append(py) + + for cmd in commands: + if cmd["kind"] == "spin": + heading += math.radians(cmd["degrees"]) + elif cmd["kind"] == "line": + new_pos = pos + cmd["distance"] * np.array( + [math.cos(heading), math.sin(heading)] + ) + if cmd["penDown"]: + drawn.append( + { + "kind": "line", + "p0": pos.copy(), + "p1": new_pos.copy(), + } + ) + else: + pen_up.append((pos.copy(), new_pos.copy())) + update_bounds(new_pos[0], new_pos[1]) + pos = new_pos + elif cmd["kind"] == "arc": + r = float(cmd["radius"]) + sweep = math.radians(cmd["degrees"]) + ccw = sweep > 0 + # Centre is 90deg to the left of heading for CCW, right for CW. + normal_angle = heading + (math.pi / 2 if ccw else -math.pi / 2) + center = pos + r * np.array( + [math.cos(normal_angle), math.sin(normal_angle)] + ) + start_a = math.atan2(pos[1] - center[1], pos[0] - center[0]) + end_a = start_a + sweep + new_pos = center + r * np.array([math.cos(end_a), math.sin(end_a)]) + drawn.append( + { + "kind": "arc", + "p0": pos.copy(), + "p1": new_pos.copy(), + "center": center.copy(), + "radius": r, + "sweep": sweep, + } + ) + # Sample arc for bounds + n_samp = max(2, int(abs(sweep) * 8)) + for k in range(n_samp + 1): + t = k / n_samp + a = start_a + t * sweep + update_bounds( + center[0] + r * math.cos(a), + center[1] + r * math.sin(a), + ) + pos = new_pos + heading += sweep + else: + raise ValueError(f"Unknown command kind: {cmd!r}") + + return drawn, pen_up, (min(xs), min(ys), max(xs), max(ys)) + + +def _render_drawing_parts( + drawn, + pen_up, + stroke_width, + pen_up_stroke_width, + show_pen_up, +): + """SVG fragments for one drawing -- pen-up dashes, drawn primitives + rainbow-colored by execution order, start dot and end ring -- in the + drawing's native coordinate system. Caller is responsible for the + outer , the background , and any wrapping + that places the drawing in the final layout. + """ + parts = [] + + if show_pen_up: + parts.append('') + for p0, p1 in pen_up: + parts.append( + f' ' + ) + parts.append("") + + n = len(drawn) + parts.append('') + for i, d in enumerate(drawn): + hue = 360.0 * i / max(n, 1) + color = _hsl(hue) + if d["kind"] == "line": + p0, p1 = d["p0"], d["p1"] + parts.append( + f' ' + ) + else: # arc + p0, p1 = d["p0"], d["p1"] + r = d["radius"] + sweep = d["sweep"] + if abs(sweep) >= 2 * math.pi - 1e-3: + center = d["center"] + start_a = math.atan2(p0[1] - center[1], p0[0] - center[0]) + mid_a = start_a + sweep / 2.0 + pmx = center[0] + r * math.cos(mid_a) + pmy = center[1] + r * math.sin(mid_a) + sweep_flag = 1 if sweep > 0 else 0 + parts.append( + f' ' + ) + else: + large_arc = 1 if abs(sweep) > math.pi else 0 + sweep_flag = 1 if sweep > 0 else 0 + parts.append( + f' ' + ) + parts.append("") + + if drawn: + first_p0 = drawn[0]["p0"] + last_p1 = drawn[-1]["p1"] + parts.append( + f'' + ) + parts.append( + f'' + ) + + return parts + + +def commands_to_svg( + commands, + output_path, + start_pos=(0.0, 0.0), + start_heading=0.0, + stroke_width=1.5, + pen_up_stroke_width=0.5, + padding=8.0, + show_pen_up=True, +): + """Render a command list to an SVG file. Returns the SVG string.""" + drawn, pen_up, (minx, miny, maxx, maxy) = _simulate( + commands, start_pos, start_heading + ) + minx -= padding + miny -= padding + maxx += padding + maxy += padding + width = maxx - minx + height = maxy - miny + + parts = [ + f'', + '', + ] + parts.extend( + _render_drawing_parts( + drawn, + pen_up, + stroke_width, + pen_up_stroke_width, + show_pen_up, + ) + ) + parts.append("") + + svg = "\n".join(parts) + with open(output_path, "w") as f: + f.write(svg) + return svg + + +def commands_to_svg_compare( + commands_a: Sequence[DrawingCommand], + commands_b: Sequence[DrawingCommand], + output_path, + label_a="A", + label_b="B", + start_pos=(0.0, 0.0), + start_heading=0.0, + stroke_width=1.5, + pen_up_stroke_width=0.5, + padding=8.0, + panel_gap=24.0, + label_height=28.0, + show_pen_up=True, +): + """Render two command lists side by side at the same scale. + + Both panels share a unified bounding box (the union of each + drawing's padded bbox), so a primitive at drawing-space (X, Y) in + `commands_a` appears at exactly the same panel-relative position + as a primitive at (X, Y) in `commands_b`. That equivalence is what + makes "spot the difference" actually work -- if you rendered each + panel to its own bbox, drawings of slightly different extent would + end up at different scales and the visual diff would be muddled. + Each panel still gets its own rainbow over its own primitives. + """ + drawn_a, pen_up_a, bbox_a = _simulate(commands_a, start_pos, start_heading) + drawn_b, pen_up_b, bbox_b = _simulate(commands_b, start_pos, start_heading) + + minx = min(bbox_a[0], bbox_b[0]) - padding + miny = min(bbox_a[1], bbox_b[1]) - padding + maxx = max(bbox_a[2], bbox_b[2]) + padding + maxy = max(bbox_a[3], bbox_b[3]) + padding + panel_w = maxx - minx + panel_h = maxy - miny + + total_w = 2 * panel_w + panel_gap + total_h = panel_h + label_height + label_baseline = label_height * 0.7 + font_size = label_height * 0.5 + + parts = [ + f'', + '', + f'', + f'{label_a}', + f'{label_b}', + ] + + parts.append(f'') + parts.extend( + _render_drawing_parts( + drawn_a, + pen_up_a, + stroke_width, + pen_up_stroke_width, + show_pen_up, + ) + ) + parts.append("") + + parts.append( + f'' + ) + parts.extend( + _render_drawing_parts( + drawn_b, + pen_up_b, + stroke_width, + pen_up_stroke_width, + show_pen_up, + ) + ) + parts.append("") + + parts.append("") + + svg = "\n".join(parts) + with open(output_path, "w") as f: + f.write(svg) + return svg + + +def _primitive_length(primitive: dict) -> float: + if primitive["kind"] == "line": + return float(np.linalg.norm(primitive["p1"] - primitive["p0"])) + return float(abs(primitive["sweep"]) * primitive["radius"]) + + +def _allocate_frames_by_length( + lengths: List[float], + total_frames: int, +) -> List[int]: + if not lengths: + return [] + if total_frames <= 0: + total_frames = len(lengths) + sum_len = float(sum(lengths)) + if sum_len <= 1e-9: + return [max(1, total_frames // len(lengths))] * len(lengths) + + alloc = [max(1, int(round(total_frames * (L / sum_len)))) for L in lengths] + cur = sum(alloc) + if cur == total_frames: + return alloc + + # Adjust allocation to hit the target frame count exactly. + order = sorted(range(len(lengths)), key=lambda i: lengths[i], reverse=True) + if cur < total_frames: + k = 0 + while cur < total_frames: + alloc[order[k % len(order)]] += 1 + cur += 1 + k += 1 + else: + k = 0 + while cur > total_frames: + idx = order[k % len(order)] + if alloc[idx] > 1: + alloc[idx] -= 1 + cur -= 1 + k += 1 + return alloc + + +def _draw_partial_primitive( + draw: ImageDraw.ImageDraw, + primitive: dict, + t: float, + color: str, + stroke_width: int, + map_pt, +) -> None: + t = float(max(0.0, min(1.0, t))) + if primitive["kind"] == "line": + p0 = primitive["p0"] + p1 = primitive["p1"] + p = p0 + t * (p1 - p0) + draw.line([map_pt(p0), map_pt(p)], fill=color, width=stroke_width) + return + + center = primitive["center"] + radius = float(primitive["radius"]) + sweep = float(primitive["sweep"]) * t + p0 = primitive["p0"] + start_a = math.atan2(p0[1] - center[1], p0[0] - center[0]) + n_samp = max(2, int(abs(sweep) * radius * 0.8)) + pts = [] + for k in range(n_samp + 1): + a = start_a + sweep * (k / n_samp) + p = np.array( + [center[0] + radius * math.cos(a), center[1] + radius * math.sin(a)] + ) + pts.append(map_pt(p)) + draw.line(pts, fill=color, width=stroke_width) + + +def commands_to_svg_gif( + commands: Sequence[DrawingCommand], + output_path: str, + start_pos: Tuple[float, float] = (0.0, 0.0), + start_heading: float = 0.0, + stroke_width: int = 2, + padding: float = 8.0, + scale: float = 4.0, + fps: int = 24, + duration_s: Optional[float] = None, + units_per_second: float = 60.0, + max_total_frames: int = 240, + max_pixels_per_frame: int = 400_000, + show_pen_up: bool = False, + pen_up_stroke_width: int = 1, +) -> str: + """Create an animated GIF of the drawing process in primitive order. + + The geometry and ordering match the SVG simulator: each line/arc is + animated progressively, then the next primitive starts. + """ + drawn, pen_up, (minx, miny, maxx, maxy) = _simulate( + commands, start_pos, start_heading + ) + + minx -= padding + miny -= padding + maxx += padding + maxy += padding + width_px = max(1, int(math.ceil((maxx - minx) * scale))) + height_px = max(1, int(math.ceil((maxy - miny) * scale))) + px_count = width_px * height_px + if px_count > max_pixels_per_frame > 0: + shrink = math.sqrt(max_pixels_per_frame / float(px_count)) + scale *= shrink + width_px = max(1, int(math.ceil((maxx - minx) * scale))) + height_px = max(1, int(math.ceil((maxy - miny) * scale))) + + def map_pt(p: np.ndarray) -> Tuple[float, float]: + return ((float(p[0]) - minx) * scale, (float(p[1]) - miny) * scale) + + lengths = [_primitive_length(d) for d in drawn] + total_length = float(sum(lengths)) + if duration_s is None: + duration_s = max(1.0, total_length / max(1e-6, units_per_second)) + total_frames = max(1, int(round(duration_s * max(1, fps)))) + if max_total_frames > 0: + total_frames = min(total_frames, max_total_frames) + frames_per_primitive = _allocate_frames_by_length(lengths, total_frames) + + base = Image.new("RGB", (width_px, height_px), "white") + base_draw = ImageDraw.Draw(base) + if show_pen_up: + for p0, p1 in pen_up: + base_draw.line( + [map_pt(p0), map_pt(p1)], + fill=(190, 190, 190), + width=pen_up_stroke_width, + ) + + frames: List[Image.Image] = [] + n = len(drawn) + for i, primitive in enumerate(drawn): + hue = 360.0 * i / max(n, 1) + color = _hsl(hue) + n_frames = frames_per_primitive[i] if i < len(frames_per_primitive) else 1 + for k in range(1, n_frames + 1): + frame = base.copy() + draw = ImageDraw.Draw(frame) + _draw_partial_primitive( + draw, + primitive, + t=k / n_frames, + color=color, + stroke_width=stroke_width, + map_pt=map_pt, + ) + frames.append( + frame.convert( + "P", + palette=Image.Palette.ADAPTIVE, + colors=256, + dither=Image.Dither.NONE, + ) + ) + + _draw_partial_primitive( + base_draw, + primitive, + t=1.0, + color=color, + stroke_width=stroke_width, + map_pt=map_pt, + ) + + if not frames: + frames = [ + base.convert( + "P", + palette=Image.Palette.ADAPTIVE, + colors=256, + dither=Image.Dither.NONE, + ) + ] + + frame_ms = max(1, int(round(1000 / max(1, fps)))) + frames[0].save( + output_path, + save_all=True, + append_images=frames[1:], + optimize=False, + duration=frame_ms, + loop=0, + ) + return output_path diff --git a/main.py b/main.py index d441536..e6f4517 100644 --- a/main.py +++ b/main.py @@ -1,6 +1,7 @@ from fastapi import FastAPI, HTTPException, UploadFile, File, Query import httpx import aiohttp +import io import logging from fastapi.responses import FileResponse, StreamingResponse, HTMLResponse from fastapi.staticfiles import StaticFiles @@ -8,6 +9,8 @@ from typing import Optional import azure.cognitiveservices.speech as speechsdk from openai import OpenAI +import numpy as np +from PIL import Image import pyaudio import wave import asyncio @@ -17,6 +20,9 @@ from functools import wraps from fastapi.middleware.cors import CORSMiddleware +from arc_line_vectorization_suede import default_pipeline +from arc_line_vectorization_suede.visualize import commands_to_svg_compare + # Load environment variables load_dotenv() @@ -24,7 +30,7 @@ app = FastAPI( title="Voice Assistant API", description="A voice assistant that converts speech to text, processes it, and returns synthesized speech", - version="1.0.0" + version="1.0.0", ) app.add_middleware( CORSMiddleware, @@ -36,9 +42,9 @@ ) # Initialize API clients -openai_client = OpenAI(api_key=os.getenv('OPENAI_API_KEY')) -azure_speech_key = os.getenv('AZURE_SPEECH_KEY') -azure_service_region = os.getenv('AZURE_SPEECH_REGION') +openai_client = OpenAI(api_key=os.getenv("OPENAI_API_KEY")) +azure_speech_key = os.getenv("AZURE_SPEECH_KEY") +azure_service_region = os.getenv("AZURE_SPEECH_REGION") VOICE_MAP = { 1: "en-US-AnaNeural", @@ -48,16 +54,19 @@ 5: "en-US-BrianMultilingualNeural", 6: "en-US-CoraMultilingualNeural", 7: "en-US-LewisMultilingualNeural", - 8: "en-US-EmmaNeural" + 8: "en-US-EmmaNeural", } + class VoiceAssistantError(Exception): """Custom exception for Voice Assistant errors""" + pass def handle_errors(func): """Decorator for error handling""" + @wraps(func) async def wrapper(*args, **kwargs): try: @@ -66,7 +75,9 @@ async def wrapper(*args, **kwargs): raise HTTPException(status_code=400, detail=str(e)) except Exception as e: raise HTTPException( - status_code=500, detail=f"Internal server error: {str(e)}") + status_code=500, detail=f"Internal server error: {str(e)}" + ) + return wrapper @@ -74,10 +85,9 @@ class VoiceAssistant: def __init__(self): self.conversation_history = [] self.temp_dir = tempfile.mkdtemp() - self.openai_client = OpenAI(api_key=os.getenv('OPENAI_API_KEY')) + self.openai_client = OpenAI(api_key=os.getenv("OPENAI_API_KEY")) self.speech_config = speechsdk.SpeechConfig( - subscription=azure_speech_key, - region=azure_service_region + subscription=azure_speech_key, region=azure_service_region ) self.speech_config.speech_synthesis_voice_name = "en-US-AnaNeural" @@ -98,7 +108,7 @@ async def record_audio(self) -> bytes: channels=self.CHANNELS, rate=self.RATE, input=True, - frames_per_buffer=self.CHUNK + frames_per_buffer=self.CHUNK, ) frames = [] @@ -107,14 +117,14 @@ async def record_audio(self) -> bytes: frames.append(data) temp_path = os.path.join(self.temp_dir, "temp_recording.wav") - wf = wave.open(temp_path, 'wb') + wf = wave.open(temp_path, "wb") wf.setnchannels(self.CHANNELS) wf.setsampwidth(p.get_sample_size(self.FORMAT)) wf.setframerate(self.RATE) - wf.writeframes(b''.join(frames)) + wf.writeframes(b"".join(frames)) wf.close() - with open(temp_path, 'rb') as audio_file: + with open(temp_path, "rb") as audio_file: audio_bytes = audio_file.read() return audio_bytes @@ -172,30 +182,42 @@ async def get_chat_response(self, text: str) -> str: Your goal is to make learning interactive, thought-provoking, and fun. You always encourage creativity and exploration rather than just giving answers. Stay playful, supportive, and engaging—but always remember, you're a robot! """ - self.conversation_history.append({"role": "system", "content": system_prompt}) + self.conversation_history.append( + {"role": "system", "content": system_prompt} + ) self.conversation_history.append({"role": "user", "content": text}) response = self.openai_client.chat.completions.create( - model="gpt-4", - messages=self.conversation_history, - max_tokens=150 + model="gpt-4", messages=self.conversation_history, max_tokens=150 ) assistant_response = response.choices[0].message.content self.conversation_history.append( - {"role": "assistant", "content": assistant_response}) + {"role": "assistant", "content": assistant_response} + ) return assistant_response except Exception as e: raise VoiceAssistantError(f"Chat processing failed: {str(e)}") - async def synthesize_speech(self, text: str, voice: str = "en-US-AnaNeural", pitch: str = "default", rate: Optional[str] = None) -> str: + async def synthesize_speech( + self, + text: str, + voice: str = "en-US-AnaNeural", + pitch: str = "default", + rate: Optional[str] = None, + ) -> str: print("voice", voice) output_path = os.path.join(self.temp_dir, "response.wav") audio_config = speechsdk.audio.AudioOutputConfig(filename=output_path) - speech_config = speechsdk.SpeechConfig(subscription=self.speech_config.subscription_key, region=self.speech_config.region) + speech_config = speechsdk.SpeechConfig( + subscription=self.speech_config.subscription_key, + region=self.speech_config.region, + ) speech_config.speech_synthesis_voice_name = voice - synthesizer = speechsdk.SpeechSynthesizer(speech_config=speech_config, audio_config=audio_config) + synthesizer = speechsdk.SpeechSynthesizer( + speech_config=speech_config, audio_config=audio_config + ) prosody_attrs = f'pitch="{pitch}"' if rate: @@ -217,7 +239,13 @@ async def synthesize_speech(self, text: str, voice: str = "en-US-AnaNeural", pit else: raise VoiceAssistantError("Speech synthesis failed") - async def process_voice_input(self, audio_data: bytes = None, voice: str = "en-US-AnaNeural", pitch: str = "default", rate: Optional[str] = None) -> tuple[str, str]: + async def process_voice_input( + self, + audio_data: bytes = None, + voice: str = "en-US-AnaNeural", + pitch: str = "default", + rate: Optional[str] = None, + ) -> tuple[str, str]: """Process voice input and return response text and audio file path""" try: if audio_data is None: @@ -232,7 +260,13 @@ async def process_voice_input(self, audio_data: bytes = None, voice: str = "en-U except Exception as e: raise VoiceAssistantError(f"Voice processing failed: {str(e)}") - async def process_voice_input_chat(self, audio_data: bytes = None, voice: str = "en-US-AnaNeural", pitch: str = "default", rate: Optional[str] = None) -> tuple[str, str]: + async def process_voice_input_chat( + self, + audio_data: bytes = None, + voice: str = "en-US-AnaNeural", + pitch: str = "default", + rate: Optional[str] = None, + ) -> tuple[str, str]: """Process voice input and return response text and audio file path""" try: if audio_data is None: @@ -245,11 +279,11 @@ async def process_voice_input_chat(self, audio_data: bytes = None, voice: str = except Exception as e: raise VoiceAssistantError(f"Voice processing failed: {str(e)}") - def cleanup(self): """Clean up temporary files""" import shutil + try: shutil.rmtree(self.temp_dir) except Exception: @@ -264,10 +298,14 @@ class ChatResponse(BaseModel): class TextInput(BaseModel): text: str + @app.post("/repeat_after_me") @handle_errors -async def repeat_after_me(audio_file: UploadFile = File(None), voice: int = Query(default=None, description="Voice ID (1-8)"), - pitch: int = Query(default=0, description="Pitch adjustment (e.g., -5 to +5)")): +async def repeat_after_me( + audio_file: UploadFile = File(None), + voice: int = Query(default=None, description="Voice ID (1-8)"), + pitch: int = Query(default=0, description="Pitch adjustment (e.g., -5 to +5)"), +): assistant = VoiceAssistant() try: audio_data = None @@ -279,33 +317,38 @@ async def repeat_after_me(audio_file: UploadFile = File(None), voice: int = Quer pitch_value = "default" else: pitch_value = f"{pitch_value:+d}st" # + sign added for positive numbers - response_text, audio_path = await assistant.process_voice_input_chat(audio_data, voice=voice_value, pitch=pitch_value) + response_text, audio_path = await assistant.process_voice_input_chat( + audio_data, voice=voice_value, pitch=pitch_value + ) - with open(audio_path, 'rb') as f: + with open(audio_path, "rb") as f: audio_content = f.read() assistant.cleanup() - temp_response_path = tempfile.mktemp(suffix='.wav') - with open(temp_response_path, 'wb') as f: + temp_response_path = tempfile.mktemp(suffix=".wav") + with open(temp_response_path, "wb") as f: f.write(audio_content) return FileResponse( path=temp_response_path, media_type="audio/wav", headers={"text-response": response_text}, - filename="response.wav" + filename="response.wav", ) except Exception as e: if assistant: assistant.cleanup() raise VoiceAssistantError(f"Speech synthesis failed: {str(e)}") + @app.post("/speak") @handle_errors -async def speak_endpoint(input_data: TextInput, +async def speak_endpoint( + input_data: TextInput, voice: int = Query(default=None, description="Voice ID (1-8)"), - pitch: int = Query(default=0, description="Pitch adjustment (e.g., -5 to +5)")): + pitch: int = Query(default=0, description="Pitch adjustment (e.g., -5 to +5)"), +): """Convert text to speech and return audio file""" assistant = VoiceAssistant() @@ -317,21 +360,21 @@ async def speak_endpoint(input_data: TextInput, pitch_value = f"{pitch_value:+d}st" # + sign added for positive numbers try: - audio_path = await assistant.synthesize_speech(input_data.text, voice=voice_value, pitch=pitch_value) + audio_path = await assistant.synthesize_speech( + input_data.text, voice=voice_value, pitch=pitch_value + ) - with open(audio_path, 'rb') as f: + with open(audio_path, "rb") as f: audio_content = f.read() assistant.cleanup() - temp_response_path = tempfile.mktemp(suffix='.wav') - with open(temp_response_path, 'wb') as f: + temp_response_path = tempfile.mktemp(suffix=".wav") + with open(temp_response_path, "wb") as f: f.write(audio_content) return FileResponse( - path=temp_response_path, - media_type="audio/wav", - filename="speech.wav" + path=temp_response_path, media_type="audio/wav", filename="speech.wav" ) except Exception as e: if assistant: @@ -344,18 +387,21 @@ async def root(): """Health check endpoint""" return {"status": "ok", "message": "Voice Assistant API is running"} + async def mjpeg_proxy_stream(ip_address: str): stream_url = f"http://{ip_address}:8000/video_feed" async with aiohttp.ClientSession() as session: async with session.get(stream_url) as resp: if resp.status != 200: raise Exception(f"Failed to fetch stream: {resp.status}") - + async for data, _ in resp.content.iter_chunks(): yield data + VIDEO_FEED_URL = "http://192.168.41.214:8000/video_feed" + @app.get("/proxy/video_feed") async def proxy_video_feed(): @@ -367,7 +413,9 @@ async def video_stream(): yield chunk await asyncio.sleep(0.001) - return StreamingResponse(video_stream(), media_type="multipart/x-mixed-replace; boundary=frame") + return StreamingResponse( + video_stream(), media_type="multipart/x-mixed-replace; boundary=frame" + ) @app.get("/mjpeg-viewer", response_class=HTMLResponse) @@ -381,7 +429,6 @@ async def mjpeg_viewer(ip_address: str): """ - # Pitch map function (converts int to SSML pitch string) def map_pitch_value(pitch_int: int) -> str: if pitch_int == 0: @@ -391,13 +438,14 @@ def map_pitch_value(pitch_int: int) -> str: else: return f"{pitch_int * 5}%" + @app.post("/chat", response_model=ChatResponse) @handle_errors async def chat_endpoint( - audio_file: UploadFile = File(None), + audio_file: UploadFile = File(None), voice: int = Query(default=None, description="Voice ID (1-8)"), - pitch: int = Query(default=0, description="Pitch adjustment (e.g., -5 to +5)") - ): + pitch: int = Query(default=0, description="Pitch adjustment (e.g., -5 to +5)"), +): """Process voice input and return response""" assistant = VoiceAssistant() try: @@ -412,28 +460,93 @@ async def chat_endpoint( else: pitch_value = f"{pitch_value:+d}st" # + sign added for positive numbers - response_text, audio_path = await assistant.process_voice_input(audio_data, voice=voice_value, pitch=pitch_value) + response_text, audio_path = await assistant.process_voice_input( + audio_data, voice=voice_value, pitch=pitch_value + ) - with open(audio_path, 'rb') as f: + with open(audio_path, "rb") as f: audio_content = f.read() assistant.cleanup() - temp_response_path = tempfile.mktemp(suffix='.wav') - with open(temp_response_path, 'wb') as f: + temp_response_path = tempfile.mktemp(suffix=".wav") + with open(temp_response_path, "wb") as f: f.write(audio_content) return FileResponse( path=temp_response_path, media_type="audio/wav", headers={"text-response": response_text}, - filename="response.wav" + filename="response.wav", ) except Exception as e: if assistant: assistant.cleanup() raise VoiceAssistantError(f"Chat processing failed: {str(e)}") + +def _commands_to_jsonable(commands): + """Strip numpy scalar wrappers so DrawingCommand dicts serialize cleanly.""" + out = [] + for cmd in commands: + item = {} + for key, value in cmd.items(): + if isinstance(value, np.bool_): + item[key] = bool(value) + elif isinstance(value, np.integer): + item[key] = int(value) + elif isinstance(value, np.floating): + item[key] = float(value) + else: + item[key] = value + out.append(item) + return out + + +def _run_vectorization(image_array: np.ndarray): + _, _, _, low_geometry, high_geometry = default_pipeline(image_array) + with tempfile.NamedTemporaryFile(suffix=".svg", delete=False) as tmp: + tmp_path = tmp.name + try: + svg = commands_to_svg_compare( + low_geometry.consolidated, + high_geometry.commands, + tmp_path, + label_a="low_geometry.consolidated", + label_b="high_geometry.commands", + ) + finally: + try: + os.remove(tmp_path) + except OSError: + pass + return { + "low_geometry_consolidated": _commands_to_jsonable(low_geometry.consolidated), + "high_geometry_commands": _commands_to_jsonable(high_geometry.commands), + "svg": svg, + } + + +@app.post("/vectorize") +@handle_errors +async def vectorize_endpoint(image_file: UploadFile = File(...)): + """Vectorize an uploaded image into robot drawing commands. + + Returns the low-geometry consolidated commands, the high-geometry + commands, and a side-by-side comparison SVG of the two. + """ + image_bytes = await image_file.read() + if not image_bytes: + raise VoiceAssistantError("Empty image upload") + try: + pil_image = Image.open(io.BytesIO(image_bytes)) + pil_image.load() + except Exception as e: + raise VoiceAssistantError(f"Could not decode image: {e}") + image_array = np.asarray(pil_image) + return await asyncio.to_thread(_run_vectorization, image_array) + + def get_static_directory(name: str): return os.path.join(os.getcwd(), name) @@ -441,8 +554,7 @@ def get_static_directory(name: str): def try_mount_static_html(app, name: str, prefix: str = "/"): directory = get_static_directory(name) if os.path.exists(directory): - app.mount(prefix, StaticFiles( - directory=directory, html=True), name=name) + app.mount(prefix, StaticFiles(directory=directory, html=True), name=name) print(f"Mounted {name} at {prefix}") else: print(f"Directory not found: {directory}") @@ -452,4 +564,5 @@ def try_mount_static_html(app, name: str, prefix: str = "/"): if __name__ == "__main__": import uvicorn + uvicorn.run(app, host="0.0.0.0", port=8000) From ea2c39fb950fc3e0ef246d144b6d505d06faecc4 Mon Sep 17 00:00:00 2001 From: "Github Action (authored by pmalacho-mit)" Date: Fri, 15 May 2026 14:32:18 +0000 Subject: [PATCH 3/6] simplified implementation --- arc_line_vectorization_suede/visualize.py | 14 ++++++++------ main.py | 21 ++++++--------------- 2 files changed, 14 insertions(+), 21 deletions(-) diff --git a/arc_line_vectorization_suede/visualize.py b/arc_line_vectorization_suede/visualize.py index 61b176c..a46c764 100644 --- a/arc_line_vectorization_suede/visualize.py +++ b/arc_line_vectorization_suede/visualize.py @@ -176,7 +176,7 @@ def _render_drawing_parts( def commands_to_svg( commands, - output_path, + output_path: Optional[str] = None, start_pos=(0.0, 0.0), start_heading=0.0, stroke_width=1.5, @@ -214,15 +214,16 @@ def commands_to_svg( parts.append("") svg = "\n".join(parts) - with open(output_path, "w") as f: - f.write(svg) + if output_path is not None: + with open(output_path, "w") as f: + f.write(svg) return svg def commands_to_svg_compare( commands_a: Sequence[DrawingCommand], commands_b: Sequence[DrawingCommand], - output_path, + output_path: Optional[str] = None, label_a="A", label_b="B", start_pos=(0.0, 0.0), @@ -309,8 +310,9 @@ def commands_to_svg_compare( parts.append("") svg = "\n".join(parts) - with open(output_path, "w") as f: - f.write(svg) + if output_path is not None: + with open(output_path, "w") as f: + f.write(svg) return svg diff --git a/main.py b/main.py index e6f4517..d3a8443 100644 --- a/main.py +++ b/main.py @@ -505,21 +505,12 @@ def _commands_to_jsonable(commands): def _run_vectorization(image_array: np.ndarray): _, _, _, low_geometry, high_geometry = default_pipeline(image_array) - with tempfile.NamedTemporaryFile(suffix=".svg", delete=False) as tmp: - tmp_path = tmp.name - try: - svg = commands_to_svg_compare( - low_geometry.consolidated, - high_geometry.commands, - tmp_path, - label_a="low_geometry.consolidated", - label_b="high_geometry.commands", - ) - finally: - try: - os.remove(tmp_path) - except OSError: - pass + svg = commands_to_svg_compare( + low_geometry.consolidated, + high_geometry.commands, + label_a="low_geometry.consolidated", + label_b="high_geometry.commands", + ) return { "low_geometry_consolidated": _commands_to_jsonable(low_geometry.consolidated), "high_geometry_commands": _commands_to_jsonable(high_geometry.commands), From 16294105c22616c2d0b0ee9fdf7175c64a9631fc Mon Sep 17 00:00:00 2001 From: "Github Action (authored by pmalacho-mit)" Date: Fri, 15 May 2026 14:41:47 +0000 Subject: [PATCH 4/6] refactored to have vectorization in seperate file --- arc_line_vectorization_suede/__init__.py | 1 + main.py | 42 +++------------------- vectorization.py | 45 ++++++++++++++++++++++++ 3 files changed, 50 insertions(+), 38 deletions(-) create mode 100644 vectorization.py diff --git a/arc_line_vectorization_suede/__init__.py b/arc_line_vectorization_suede/__init__.py index 3889b67..74a3405 100644 --- a/arc_line_vectorization_suede/__init__.py +++ b/arc_line_vectorization_suede/__init__.py @@ -3,6 +3,7 @@ from .graph import StrokeGraph from .vectorize.low_geometry import Vectorize as LowGeometryVectorize from .vectorize.high_geometry import Vectorize as HighGeometryVectorize +from .commands import DrawingCommand import numpy as np diff --git a/main.py b/main.py index d3a8443..a3b856a 100644 --- a/main.py +++ b/main.py @@ -20,8 +20,7 @@ from functools import wraps from fastapi.middleware.cors import CORSMiddleware -from arc_line_vectorization_suede import default_pipeline -from arc_line_vectorization_suede.visualize import commands_to_svg_compare +from vectorization import run_vectorization, VectorizationError # Load environment variables load_dotenv() @@ -485,39 +484,6 @@ async def chat_endpoint( raise VoiceAssistantError(f"Chat processing failed: {str(e)}") -def _commands_to_jsonable(commands): - """Strip numpy scalar wrappers so DrawingCommand dicts serialize cleanly.""" - out = [] - for cmd in commands: - item = {} - for key, value in cmd.items(): - if isinstance(value, np.bool_): - item[key] = bool(value) - elif isinstance(value, np.integer): - item[key] = int(value) - elif isinstance(value, np.floating): - item[key] = float(value) - else: - item[key] = value - out.append(item) - return out - - -def _run_vectorization(image_array: np.ndarray): - _, _, _, low_geometry, high_geometry = default_pipeline(image_array) - svg = commands_to_svg_compare( - low_geometry.consolidated, - high_geometry.commands, - label_a="low_geometry.consolidated", - label_b="high_geometry.commands", - ) - return { - "low_geometry_consolidated": _commands_to_jsonable(low_geometry.consolidated), - "high_geometry_commands": _commands_to_jsonable(high_geometry.commands), - "svg": svg, - } - - @app.post("/vectorize") @handle_errors async def vectorize_endpoint(image_file: UploadFile = File(...)): @@ -528,14 +494,14 @@ async def vectorize_endpoint(image_file: UploadFile = File(...)): """ image_bytes = await image_file.read() if not image_bytes: - raise VoiceAssistantError("Empty image upload") + raise VectorizationError("Empty image upload") try: pil_image = Image.open(io.BytesIO(image_bytes)) pil_image.load() except Exception as e: - raise VoiceAssistantError(f"Could not decode image: {e}") + raise VectorizationError(f"Could not decode image: {e}") image_array = np.asarray(pil_image) - return await asyncio.to_thread(_run_vectorization, image_array) + return await asyncio.to_thread(run_vectorization, image_array) def get_static_directory(name: str): diff --git a/vectorization.py b/vectorization.py new file mode 100644 index 0000000..6ccfd75 --- /dev/null +++ b/vectorization.py @@ -0,0 +1,45 @@ +import numpy as np + +from typing import Sequence + +from arc_line_vectorization_suede import default_pipeline, DrawingCommand +from arc_line_vectorization_suede.visualize import commands_to_svg_compare + + +def _commands_to_jsonable(commands: Sequence[DrawingCommand]): + """Strip numpy scalar wrappers so DrawingCommand dicts serialize cleanly.""" + out = [] + for cmd in commands: + item = {} + for key, value in cmd.items(): + if isinstance(value, np.bool_): + item[key] = bool(value) + elif isinstance(value, np.integer): + item[key] = int(value) + elif isinstance(value, np.floating): + item[key] = float(value) + else: + item[key] = value + out.append(item) + return out + + +def run_vectorization(image_array: np.ndarray): + _, _, _, low_geometry, high_geometry = default_pipeline(image_array) + svg = commands_to_svg_compare( + low_geometry.consolidated, + high_geometry.commands, + label_a="low_geometry.consolidated", + label_b="high_geometry.commands", + ) + return { + "low_geometry_consolidated": _commands_to_jsonable(low_geometry.consolidated), + "high_geometry_commands": _commands_to_jsonable(high_geometry.commands), + "svg": svg, + } + + +class VectorizationError(Exception): + """Custom exception for vectorization errors""" + + pass From 4e366f127ad247aca176d1e87cf68225e0bb7e34 Mon Sep 17 00:00:00 2001 From: "Github Action (authored by pmalacho-mit)" Date: Fri, 15 May 2026 14:43:42 +0000 Subject: [PATCH 5/6] cleaning up payload --- vectorization.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/vectorization.py b/vectorization.py index 6ccfd75..ce38896 100644 --- a/vectorization.py +++ b/vectorization.py @@ -33,8 +33,8 @@ def run_vectorization(image_array: np.ndarray): label_b="high_geometry.commands", ) return { - "low_geometry_consolidated": _commands_to_jsonable(low_geometry.consolidated), - "high_geometry_commands": _commands_to_jsonable(high_geometry.commands), + "low_geometry": _commands_to_jsonable(low_geometry.consolidated), + "high_geometry": _commands_to_jsonable(high_geometry.commands), "svg": svg, } From c799b9f23bbd8f116124f63bee945aeb0d846e01 Mon Sep 17 00:00:00 2001 From: "Github Action (authored by pmalacho-mit)" Date: Sat, 16 May 2026 01:13:29 +0000 Subject: [PATCH 6/6] git subrepo push arc_line_vectorization_suede subrepo: subdir: "arc_line_vectorization_suede" merged: "dbd7ae5" upstream: origin: "https://github.com/mitmedialab/arc-line-vectorization-suede" branch: "release" commit: "dbd7ae5" git-subrepo: version: "0.4.9" origin: "https://github.com/ingydotnet/git-subrepo.git" commit: "5e0f401" --- arc_line_vectorization_suede/.gitrepo | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/arc_line_vectorization_suede/.gitrepo b/arc_line_vectorization_suede/.gitrepo index ed71f30..4113e00 100644 --- a/arc_line_vectorization_suede/.gitrepo +++ b/arc_line_vectorization_suede/.gitrepo @@ -6,7 +6,7 @@ [subrepo] remote = https://github.com/mitmedialab/arc-line-vectorization-suede branch = release - commit = b57054125d4e38c35462839e719384d72ebafc03 - parent = e1a6ea9e6760ea8f3f2cb47e8c7b84ecf8c68181 + commit = dbd7ae57777451681797fc9822edc669973ac572 + parent = 4e366f127ad247aca176d1e87cf68225e0bb7e34 method = merge cmdver = 0.4.9