Upload 7 files

README.md

@@ -1,20 +1,93 @@
----
-title: SVG SAMurai
-emoji: 🚀
-colorFrom: red
-colorTo: red
-sdk: docker
-app_port: 8501
-tags:
-- streamlit
-pinned: false
-short_description: Tool to turn an image to an SVG with named image sections
-license: mit
----
+# SVG-SAMurai 🗡️
+
+**SVG-SAMurai** is an interactive, Streamlit-based web application that leverages the power of Meta's **Segment Anything Model (SAM)** to transform raster (PNG, JPG) and vector (SVG) images into precisely segmented, editable SVG paths.
+
+Whether you're starting from a flat image or an existing SVG file, SVG-SAMurai allows you to click on regions of interest, predict their boundaries, and inject those precise vector paths back into a master SVG document.
+
+## 🌟 Features
+
+- **Interactive Segmentation:** Click on any part of an uploaded image to instantly generate an accurate mask using the SAM Vision Transformer (`facebook/sam-vit-base`).
+- **Support for Multiple Formats:** Upload PNG, JPG, or SVG files. Vector images are rasterized cleanly in the backend for processing, allowing you to segment them seamlessly.
+- **Smart Vectorization:** Extracted masks are converted into optimized SVG `<path>` elements using the Ramer-Douglas-Peucker algorithm (via OpenCV) for smooth, simplified contours.
+- **Adjustable Simplification:** Fine-tune the vectorization epsilon factor directly from the UI to control the complexity of the generated paths.
+- **Segment Management:** Name your segments and save them into a live-updating SVG document.
+- **In-Memory Caching:** Heavy image embeddings are cached securely using Streamlit's `@st.cache_data` and `@st.cache_resource`, ensuring snappy performance and instant mask prediction on subsequent clicks.
+- **Easy Export:** Download the final composed SVG with all your tagged, labeled segments neatly organized in `<g>` groups.
+
+## 🛠️ Tech Stack
+
+* **Frontend:** [Streamlit](https://streamlit.io/), [streamlit-image-coordinates](https://pypi.org/project/streamlit-image-coordinates/)
+* **Machine Learning:** [PyTorch](https://pytorch.org/), [Hugging Face Transformers](https://huggingface.co/docs/transformers/index) (Segment Anything Model)
+* **Image Processing:** [OpenCV](https://opencv.org/) (Contour extraction & smoothing), [Pillow (PIL)](https://python-pillow.org/)
+* **SVG / DOM Manipulation:** [lxml](https://lxml.de/) (XML parsing and injection), [CairoSVG](https://cairosvg.org/) (SVG rasterization)
+* **Dependency Management:** [Poetry](https://python-poetry.org/)
+
+## 🚀 Quick Start
+
+### Prerequisites
+
+- Python `>=3.10, <3.13` (Required for PyTorch and Triton compatibility)
+- [Poetry](https://python-poetry.org/docs/#installation) installed on your system.
+- System dependencies for CairoSVG and OpenCV (e.g., `libcairo2-dev`, `libgl1-mesa-glx` on Ubuntu/Debian).
+
+### Installation
 
-# Welcome to Streamlit!
+1. **Clone the repository:**
+   ```bash
+   git clone <repository-url>
+   cd svg-samurai
+   ```
 
-Edit `/src/streamlit_app.py` to customize this app to your heart's desire. :heart:
+2. **Install dependencies using Poetry:**
+   ```bash
+   poetry install
+   ```
 
-If you have any questions, checkout our [documentation](https://docs.streamlit.io) and [community
-forums](https://discuss.streamlit.io).
+### Running the App
+
+Start the Streamlit development server:
+
+```bash
+poetry run streamlit run app.py
+```
+
+The application will launch in your default web browser at `http://localhost:8501`.
+
+## 📖 How to Use
+
+1. **Upload an Image:** Use the file uploader to select a PNG, JPG, or SVG file. The app will calculate the complex image embeddings once (this may take a few moments depending on your hardware).
+2. **Select Segments:** Click anywhere on the image in the left panel to prompt the model.
+   - *Tip:* You can toggle the "Next Click is Negative Prompt" checkbox in the sidebar to exclude specific regions from your mask.
+3. **Refine & Save:**
+   - Use the "Undo Last Click" or "Clear Current Selection" buttons to fix mistakes.
+   - Give your highlighted segment a descriptive name (e.g., `car_body`).
+   - Adjust the **Simplification (epsilon)** slider if you want fewer, smoother nodes in your final vector path.
+   - Click **Save Segment to SVG**.
+4. **Download:** Once you have saved all desired segments, click the **Download Final SVG** button to retrieve your newly layered vector graphic.
+
+## 📂 Project Structure
+
+```text
+svg-samurai/
+├── app.py                  # Main Streamlit user interface and application state
+├── pyproject.toml          # Poetry dependencies and project configuration
+├── src/                    # Backend logic
+│   ├── model.py            # PyTorch SAM loading, embedding generation, and mask prediction
+│   ├── vectorizer.py       # OpenCV contour extraction and SVG path conversion
+│   └── xml_manager.py      # lxml DOM manipulation and CairoSVG rasterization utilities
+└── tests/                  # Unit tests for core logic
+    ├── test_model.py
+    ├── test_vectorizer.py
+    └��─ test_xml_manager.py
+```
+
+## 🧪 Testing
+
+The project uses `pytest` for unit testing. To run the test suite, simply execute:
+
+```bash
+poetry run pytest
+```
+
+---
+*Developed with Streamlit and Meta's Segment Anything Model.*

requirements.txt

@@ -1,3 +1,83 @@
-altair
-pandas
-streamlit
+altair==4.2.2
+annotated-doc==0.0.4
+anyio==3.7.1
+attrs==23.1.0
+blinker==1.6.2
+cachetools==5.3.1
+cairocffi==1.6.1
+CairoSVG==2.7.0
+certifi==2023.7.22
+cffi==1.15.1
+charset-normalizer==3.2.0
+click==8.1.7
+cssselect2==0.7.0
+defusedxml==0.7.1
+entrypoints==0.4
+filelock==3.12.2
+fsspec==2023.6.0
+gitdb==4.0.10
+GitPython==3.1.32
+h11==0.14.0
+httpcore==0.17.3
+httpx==0.24.1
+huggingface-hub==0.16.4
+idna==3.4
+iniconfig==2.0.0
+Jinja2==3.1.2
+jsonschema==4.18.6
+jsonschema-specifications==2023.7.1
+lxml==4.9.3
+markdown-it-py==3.0.0
+MarkupSafe==2.1.3
+mdurl==0.1.2
+mpmath==1.3.0
+networkx==3.1
+numpy==1.25.2
+opencv-python==4.8.0.76
+packaging==23.1
+pandas==2.0.3
+Pillow==10.0.0
+pluggy==1.2.0
+protobuf==4.23.4
+pyarrow==12.0.1
+pycparser==2.21
+pydeck==0.8.0b4
+Pygments==2.16.1
+pytest==7.4.0
+pytest-mock==3.11.1
+python-dateutil==2.8.2
+pytz==2023.3
+pytz-deprecation-shim==0.1.0.post0
+PyYAML==6.0.1
+referencing==0.30.2
+regex==2023.8.8
+requests==2.31.0
+rich==13.5.2
+rpds-py==0.9.2
+safetensors==0.3.2
+setuptools==68.0.0
+six==1.16.0
+smmap==5.0.0
+streamlit==1.25.0
+streamlit-image-coordinates==0.1.4
+sympy==1.12
+tenacity==8.2.3
+tinycss2==1.2.1
+tokenizers==0.13.3
+toml==0.10.2
+toolz==0.12.0
+torch==2.0.1
+tornado==6.3.3
+tqdm==4.66.1
+transformers==4.30.2
+triton==2.0.0
+typer==0.9.0
+typing_extensions==4.7.1
+tzdata==2023.3
+tzlocal==5.0.1
+urllib3==2.0.4
+validators==0.21.2
+watchdog==3.0.0
+webencodings==0.5.1
+wheel==0.41.1
+zipp==3.16.2

src/src/model.py

@@ -0,0 +1,95 @@
+import torch
+from transformers import SamModel, SamProcessor
+import streamlit as st
+import numpy as np
+from PIL import Image
+from typing import Tuple, List
+
+
+# Use @st.cache_resource to avoid reloading the model on every rerun
+@st.cache_resource(show_spinner="Loading Segment Anything Model (SAM)...")
+def load_sam_model() -> Tuple[SamModel, SamProcessor, str]:
+    """Loads the SAM model and processor from Hugging Face."""
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    # Using facebook/sam-vit-base as the standard baseline
+    model_id = "facebook/sam-vit-base"
+    model = SamModel.from_pretrained(model_id).to(device)
+    processor = SamProcessor.from_pretrained(model_id)
+    return model, processor, device
+
+
+@st.cache_resource(show_spinner="Computing Image Embeddings...")
+def compute_image_embedding(image: Image.Image) -> torch.Tensor:
+    """
+    Computes and caches the SAM image embedding for a given image.
+    This is the heavy part of the computation.
+    """
+    model, processor, device = load_sam_model()
+
+    # Preprocess the image to get pixel values
+    inputs = processor(images=image, return_tensors="pt").to(device)
+
+    # Compute image embeddings
+    with torch.no_grad():
+        image_embeddings = model.get_image_embeddings(inputs.pixel_values)
+
+    return image_embeddings
+
+
+def predict_mask(
+    image: Image.Image,
+    image_embeddings: torch.Tensor,
+    input_points: List[List[int]],
+    input_labels: List[int],
+) -> np.ndarray:
+    """
+    Predicts a binary mask given the image embeddings and prompt points.
+    input_points: list of [x, y] coordinates
+    input_labels: list of 1 (positive) or 0 (negative) for each point
+    """
+    model, processor, device = load_sam_model()
+
+    # Format inputs for the processor
+    # The processor expects points in the format [[[x1, y1], [x2, y2], ...]]
+    # and labels in [[1, 0, ...]] for a single batch
+    points = [input_points]
+    labels = [input_labels]
+
+    # Preprocess prompts
+    inputs = processor(
+        images=image, input_points=points, input_labels=labels, return_tensors="pt"
+    ).to(device)
+
+    # Run prediction using the cached embeddings
+    with torch.no_grad():
+        outputs = model(
+            image_embeddings=image_embeddings,
+            input_points=inputs.input_points,
+            input_labels=inputs.input_labels,
+            multimask_output=False,  # We only want the best mask
+        )
+
+    # Process the predicted mask back to the original image size
+    # inputs contains original_sizes and reshaped_input_sizes from the processor call
+    masks = processor.image_processor.post_process_masks(
+        outputs.pred_masks.cpu(),
+        inputs["original_sizes"].cpu(),
+        inputs["reshaped_input_sizes"].cpu(),
+    )
+
+    # masks is a list of tensors, get the first one and squeeze it to a 2D array
+    mask = masks[0]
+    # Squeeze out the batch and channel dimensions if present, but keep spatial dims.
+    # Usually shape is (1, 1, H, W) or (1, H, W)
+    if mask.ndim > 2:
+        mask = mask.squeeze()
+        # If the image was 1x1, squeeze might have removed all dimensions.
+        if mask.ndim < 2:
+            mask = mask.view(masks[0].shape[-2], masks[0].shape[-1])
+
+    mask = mask.numpy()
+
+    # The mask is boolean, convert to uint8 for OpenCV (0 and 255)
+    binary_mask = (mask * 255).astype(np.uint8)
+
+    return binary_mask

src/src/vectorizer.py

@@ -0,0 +1,70 @@
+import cv2
+import numpy as np
+
+
+def mask_to_svg_path(mask: np.ndarray, epsilon_factor: float = 0.005) -> str:
+    """
+    Converts a binary mask to an SVG path string.
+
+    Args:
+        mask (np.ndarray): The 2D binary mask.
+        epsilon_factor (float): The factor for approximating the contour with Ramer-Douglas-Peucker algorithm.
+            A higher value means more simplification (fewer points, smaller SVG size).
+
+    Returns:
+        str: An SVG path data string (`M x,y L x,y Z ...`).
+    """
+    if not isinstance(mask, np.ndarray) or mask.ndim != 2:
+        raise ValueError("Mask must be a 2D numpy array.")
+    # 1. Extract Contours
+    # RETR_CCOMP retrieves all of the contours and organizes them into a two-level hierarchy.
+    # At the top level, there are external boundaries of the components.
+    # At the second level, there are boundaries of the holes.
+    contours, hierarchy = cv2.findContours(
+        mask, cv2.RETR_CCOMP, cv2.CHAIN_APPROX_SIMPLE
+    )
+
+    if contours is None or len(contours) == 0:
+        return ""
+
+    path_data = []
+
+    # 2. Iterate through contours and hierarchy to build the path
+    # The hierarchy array has shape (1, num_contours, 4)
+    # The 4 elements are: [Next, Previous, First_Child, Parent]
+    if hierarchy is None:
+        return ""
+
+    hierarchy = hierarchy[0]
+
+    for i, contour in enumerate(contours):
+        # We only want to process the contours if it has at least 3 points
+        if len(contour) < 3:
+            continue
+
+        # 3. Simplify Contour
+        # Calculate epsilon based on the contour's arc length
+        epsilon = epsilon_factor * cv2.arcLength(contour, True)
+        approx = cv2.approxPolyDP(contour, epsilon, True)
+
+        # We want to skip highly simplified contours that are just points or lines
+        if len(approx) < 3:
+            continue
+
+        # 4. Format to SVG path
+        # M = moveto (start point)
+        # L = lineto (subsequent points)
+        # Z = closepath (return to start)
+        pts = approx.reshape(-1, 2)
+
+        # Add the M command for the first point
+        path_data.append(f"M {pts[0][0]},{pts[0][1]}")
+
+        # Add the L commands for the rest
+        for x, y in pts[1:]:
+            path_data.append(f"L {x},{y}")
+
+        # Close the contour
+        path_data.append("Z")
+
+    return " ".join(path_data)

src/src/xml_manager.py

@@ -0,0 +1,94 @@
+from lxml import etree
+import cairosvg
+import io
+from PIL import Image
+from typing import Any
+
+# Namespace for SVG creation
+SVG_NS = "http://www.w3.org/2000/svg"
+NSMAP = {None: SVG_NS}
+
+
+def create_base_svg(width: int, height: int) -> str:
+    """Creates a basic empty SVG string with specified dimensions."""
+    root = etree.Element(
+        "svg",
+        width=str(width),
+        height=str(height),
+        viewBox=f"0 0 {width} {height}",
+        nsmap=NSMAP,
+    )
+    return etree.tostring(root, pretty_print=True, encoding="unicode")
+
+
+def add_path_to_svg(
+    svg_str: str,
+    path_d: str,
+    path_id: str,
+    fill_color: str = "#FF0000",
+    opacity: float = 0.5,
+) -> str:
+    """
+    Injects an SVG `<path>` into an existing SVG string within a `<g>` group using lxml.
+    """
+    if not path_d:
+        return svg_str
+
+    try:
+        # Provide a parser that handles basic errors and mitigates XXE injection
+        parser = etree.XMLParser(recover=True, resolve_entities=False, no_network=True)
+        root = etree.fromstring(
+            svg_str.encode("utf-8", errors="replace"), parser=parser
+        )
+        if root is None:
+            return svg_str
+    except Exception:
+        # If the string isn't an XML document or parsing fails
+        return svg_str
+
+    # Find the correct namespace for the root or default to SVG_NS
+    ns = SVG_NS
+    if root.nsmap and None in root.nsmap:
+        ns = root.nsmap[None]
+    elif root.tag.startswith("{"):
+        ns = root.tag[1:].split("}")[0]
+
+    # Clean the namespace map to avoid redundant ns0 prefixes
+    # Ensure xmlns is explicitly available in nsmap of new elements
+    new_nsmap = {None: ns} if ns else None
+
+    # Create the <g id="path_id">
+    group = etree.SubElement(
+        root, f"{{{ns}}}g" if ns else "g", id=path_id, nsmap=new_nsmap
+    )
+
+    # Create the <path>
+    # Using fill-rule="evenodd" is important when combining outer boundaries and inner holes
+    etree.SubElement(
+        group,
+        f"{{{ns}}}path" if ns else "path",
+        d=path_d,
+        fill=fill_color,
+        opacity=str(opacity),
+        attrib={"fill-rule": "evenodd"},  # Handles holes properly
+    )
+
+    return etree.tostring(root, pretty_print=True, encoding="unicode")
+
+
+def parse_svg_to_image(svg_bytes: bytes) -> Image.Image:
+    """Converts uploaded SVG file bytes into a PIL Image."""
+    # Pass url_fetcher to block network and local file access from within SVG
+    png_bytes = cairosvg.svg2png(
+        bytestring=svg_bytes, url_fetcher=lambda *args, **kwargs: b""
+    )
+    return Image.open(io.BytesIO(png_bytes))
+
+
+def load_image(uploaded_file: Any) -> Image.Image:
+    """Loads an uploaded image (Raster or Vector) and returns a PIL Image."""
+    if getattr(uploaded_file, "type", "") == "image/svg+xml":
+        return parse_svg_to_image(uploaded_file.getvalue())
+    else:
+        # Handle regular rasters (PNG, JPG)
+        return Image.open(uploaded_file).convert("RGB")

src/streamlit_app.py

@@ -1,40 +1,223 @@
-import altair as alt
-import numpy as np
-import pandas as pd
 import streamlit as st
+from PIL import Image
+import numpy as np
+from streamlit_image_coordinates import streamlit_image_coordinates
+
+from src.model import compute_image_embedding, predict_mask
+from src.vectorizer import mask_to_svg_path
+from src.xml_manager import load_image, create_base_svg
+
+st.set_page_config(page_title="SVG-SAMurai", layout="wide", page_icon="🗡️")
+
+# Session State Initialization
+if "image" not in st.session_state:
+    st.session_state.image = None
+if "image_embedding" not in st.session_state:
+    st.session_state.image_embedding = None
+if "points" not in st.session_state:
+    st.session_state.points = []
+if "labels" not in st.session_state:
+    st.session_state.labels = []
+if "current_mask" not in st.session_state:
+    st.session_state.current_mask = None
+if "segments" not in st.session_state:
+    st.session_state.segments = {}
+if "original_svg" not in st.session_state:
+    st.session_state.original_svg = None
+
+st.title("SVG-SAMurai 🗡️")
+st.markdown(
+    "Transform raster and vector images into segmented SVG paths using the **Segment Anything Model (SAM)**."
+)
+
+# File uploader
+uploaded_file = st.file_uploader(
+    "Upload an Image (PNG, JPG, SVG)", type=["png", "jpg", "jpeg", "svg"]
+)
+
+if uploaded_file is not None:
+    # Reset state if a new file is uploaded
+    if (
+        "last_uploaded_file_id" not in st.session_state
+        or st.session_state.last_uploaded_file_id != uploaded_file.file_id
+    ):
+        st.session_state.last_uploaded_file_id = uploaded_file.file_id
+        st.session_state.image = None
+        st.session_state.image_embedding = None
+        st.session_state.points = []
+        st.session_state.labels = []
+        st.session_state.current_mask = None
+        st.session_state.segments = {}
+        st.session_state.original_svg = None
+
+    if st.session_state.image is None:
+        # Load the image
+        with st.spinner("Processing Image..."):
+            image = load_image(uploaded_file)
+            st.session_state.image = image
+
+            # If the original file was an SVG, save its string representation
+            if uploaded_file.type == "image/svg+xml":
+                st.session_state.original_svg = uploaded_file.getvalue().decode(
+                    "utf-8", errors="replace"
+                )
+            else:
+                # Create a blank SVG canvas with the original raster image dimensions
+                width, height = image.size
+                st.session_state.original_svg = create_base_svg(width, height)
+
+            # Compute image embeddings once
+            st.session_state.image_embedding = compute_image_embedding(image)
+            st.success("Image embedded successfully!")
+
+    col1, col2 = st.columns([2, 1])
+
+    with col1:
+        st.subheader("Interactive Segmentation")
+
+        # Display the image with coordinates clicker
+        # If there's a mask, we overlay it
+        display_image = st.session_state.image.copy()
+        if st.session_state.current_mask is not None:
+            # Create a semi-transparent blue overlay for the current mask
+            overlay = np.zeros(
+                (*st.session_state.current_mask.shape, 4), dtype=np.uint8
+            )
+            overlay[st.session_state.current_mask > 0] = [
+                0,
+                0,
+                255,
+                128,
+            ]  # Blue, 50% opacity
+            overlay_image = Image.fromarray(overlay, mode="RGBA")
+            display_image = display_image.convert("RGBA")
+            display_image.paste(overlay_image, (0, 0), overlay_image)
+            display_image = display_image.convert(
+                "RGB"
+            )  # Convert back to RGB for display
+
+        # Show the image using streamlit-image-coordinates
+        # Note: we need to handle scaling if the image is wider than the container
+        # streamlit-image-coordinates scales the image to the container width but gives the
+        # coordinates relative to the original image dimensions.
+        value = streamlit_image_coordinates(display_image, key="image_coord")
+
+        # Handle clicks
+        if value is not None:
+            # streamlit_image_coordinates returns x, y relative to the original image size
+            x, y = value["x"], value["y"]
+
+            # Determine if it's a positive or negative prompt
+            # For simplicity, let's say left click is positive, and we can add a toggle for negative
+            is_positive = st.sidebar.checkbox(
+                "Next Click is Negative Prompt (Exclude)",
+                value=False,
+                key="neg_prompt_toggle",
+            )
+            label = 0 if is_positive else 1
+
+            # Check if this is a new click (prevent reruns from adding the same point repeatedly)
+            new_point = [x, y]
+            if not st.session_state.points or st.session_state.points[-1] != new_point:
+                st.session_state.points.append(new_point)
+                st.session_state.labels.append(label)
+
+                # Predict new mask
+                with st.spinner("Predicting Segment..."):
+                    mask = predict_mask(
+                        st.session_state.image,
+                        st.session_state.image_embedding,
+                        st.session_state.points,
+                        st.session_state.labels,
+                    )
+                    st.session_state.current_mask = mask
+                st.rerun()
+
+        # Tools for interacting with the points
+        col_btn1, col_btn2 = st.columns(2)
+        with col_btn1:
+            if st.button("Undo Last Click"):
+                if st.session_state.points:
+                    st.session_state.points.pop()
+                    st.session_state.labels.pop()
+                    if st.session_state.points:
+                        # Repredict
+                        mask = predict_mask(
+                            st.session_state.image,
+                            st.session_state.image_embedding,
+                            st.session_state.points,
+                            st.session_state.labels,
+                        )
+                        st.session_state.current_mask = mask
+                    else:
+                        st.session_state.current_mask = None
+                    st.rerun()
+
+        with col_btn2:
+            if st.button("Clear Current Selection"):
+                st.session_state.points = []
+                st.session_state.labels = []
+                st.session_state.current_mask = None
+                st.rerun()
+
+    with col2:
+        st.subheader("Segment Management")
+
+        segment_name = st.text_input("Segment Name", placeholder="e.g., car_body")
+        epsilon_factor = st.slider(
+            "Vectorization Simplification (epsilon)",
+            min_value=0.001,
+            max_value=0.05,
+            value=0.005,
+            step=0.001,
+            format="%.3f",
+        )
+
+        if st.button(
+            "Save Segment to SVG",
+            disabled=st.session_state.current_mask is None or not segment_name,
+        ):
+            with st.spinner("Vectorizing..."):
+                # 1. Convert mask to SVG path
+                path_d = mask_to_svg_path(
+                    st.session_state.current_mask, epsilon_factor=epsilon_factor
+                )
+
+                # 2. Add to session state segments dictionary
+                st.session_state.segments[segment_name] = path_d
+
+                # 3. Inject the path into the working SVG string
+                from src.xml_manager import add_path_to_svg
+
+                try:
+                    st.session_state.original_svg = add_path_to_svg(
+                        st.session_state.original_svg,
+                        path_d,
+                        segment_name,
+                        fill_color="#FF0000",
+                        opacity=0.5,
+                    )
+
+                    # Clear current selection for the next segment
+                    st.session_state.points = []
+                    st.session_state.labels = []
+                    st.session_state.current_mask = None
+
+                    st.success(f"Segment '{segment_name}' saved!")
+                    st.rerun()
+                except Exception as e:
+                    st.error(f"Failed to inject SVG: {e}")
+
+        # Display saved segments list
+        if st.session_state.segments:
+            st.write("### Saved Segments")
+            for name in st.session_state.segments.keys():
+                st.markdown(f"- **{name}**")
 
-"""
-# Welcome to Streamlit!
-
-Edit `/streamlit_app.py` to customize this app to your heart's desire :heart:.
-If you have any questions, checkout our [documentation](https://docs.streamlit.io) and [community
-forums](https://discuss.streamlit.io).
-
-In the meantime, below is an example of what you can do with just a few lines of code:
-"""
-
-num_points = st.slider("Number of points in spiral", 1, 10000, 1100)
-num_turns = st.slider("Number of turns in spiral", 1, 300, 31)
-
-indices = np.linspace(0, 1, num_points)
-theta = 2 * np.pi * num_turns * indices
-radius = indices
-
-x = radius * np.cos(theta)
-y = radius * np.sin(theta)
-
-df = pd.DataFrame({
-    "x": x,
-    "y": y,
-    "idx": indices,
-    "rand": np.random.randn(num_points),
-})
-
-st.altair_chart(alt.Chart(df, height=700, width=700)
-    .mark_point(filled=True)
-    .encode(
-        x=alt.X("x", axis=None),
-        y=alt.Y("y", axis=None),
-        color=alt.Color("idx", legend=None, scale=alt.Scale()),
-        size=alt.Size("rand", legend=None, scale=alt.Scale(range=[1, 150])),
-    ))
+            # Provide download button for the final SVG
+            st.download_button(
+                label="Download Final SVG",
+                data=st.session_state.original_svg,
+                file_name="segmented_output.svg",
+                mime="image/svg+xml",
+            )