singlebehaviorlab 2.0.0__tar.gz → 2.2.0__tar.gz

{singlebehaviorlab-2.0.0/singlebehaviorlab.egg-info → singlebehaviorlab-2.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: singlebehaviorlab
-Version: 2.0.0
+Version: 2.2.0
 Summary: Semi-automated behavioral video annotation, training, and analysis tool
 Author: Almir Aljovic
 Maintainer: Almir Aljovic
@@ -164,6 +164,45 @@ singlebehaviorlab
 
 Equivalent module form: `python -m singlebehaviorlab`.
 
+### Headless / server use (CLI)
+
+The same `singlebehaviorlab` command also runs as a headless CLI when a subcommand is supplied. Use it to run the GPU-heavy pipeline steps on a remote machine without opening the GUI:
+
+```bash
+# Train a classifier from an experiment directory
+singlebehaviorlab train --experiment /path/to/my_experiment --profile balanced
+
+# Run a trained model on a long video
+singlebehaviorlab infer --experiment /path/to/my_experiment \
+    --model /path/to/my_experiment/models/behavior_heads/model.pt \
+    --video /data/recording.mp4 \
+    --out /data/recording_inference.json
+
+# Extract VideoPrism embeddings from a video + mask
+singlebehaviorlab register --experiment /path/to/my_experiment \
+    --video /data/recording.mp4 \
+    --mask /data/recording_masks.h5 \
+    --out /data/recording_matrix.npz
+
+# Cluster an embedding matrix (loadable via the GUI's "Load Analysis State")
+singlebehaviorlab cluster --matrix /data/recording_matrix.npz \
+    --metadata /data/recording_matrix_metadata.npz \
+    --out /data/recording_clusters.pkl
+
+# Run SAM2 tracking from a prompts JSON exported in the Segmentation tab
+singlebehaviorlab segment --video /data/recording.mp4 \
+    --prompts /data/prompts.json \
+    --out /data/recording_masks.h5
+```
+
+Run `singlebehaviorlab <command> --help` for the full flag list on each subcommand. The GUI-only steps (labeling, refinement review, cluster inspection) still require the graphical interface; the CLI covers the batch-processing steps where no human input is needed.
+
+> **Full CLI reference:** [**CLI.md**](CLI.md) — detailed per-command docs, file-format reference, Python API, and troubleshooting.
+>
+> **Notebook demos:** [**demo/**](demo/) — two Jupyter notebooks walking through behavior sequencing and segmentation/clustering end-to-end. Drop your own demo video + prompts into `demo/data/` and step through the cells.
+>
+> **Try it in Colab — no install needed:** [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/alms93/SingleBehaviorLab/blob/main/demo/colab_segmentation_clustering.ipynb) Runs the full segmentation + clustering pipeline on a bundled demo video in a free Colab GPU runtime. Install takes ~5–10 minutes; the pipeline itself finishes in a few minutes.
+
 ---
 
 ## 4. First Launch

{singlebehaviorlab-2.0.0 → singlebehaviorlab-2.2.0}/README.md

@@ -97,6 +97,45 @@ singlebehaviorlab
 
 Equivalent module form: `python -m singlebehaviorlab`.
 
+### Headless / server use (CLI)
+
+The same `singlebehaviorlab` command also runs as a headless CLI when a subcommand is supplied. Use it to run the GPU-heavy pipeline steps on a remote machine without opening the GUI:
+
+```bash
+# Train a classifier from an experiment directory
+singlebehaviorlab train --experiment /path/to/my_experiment --profile balanced
+
+# Run a trained model on a long video
+singlebehaviorlab infer --experiment /path/to/my_experiment \
+    --model /path/to/my_experiment/models/behavior_heads/model.pt \
+    --video /data/recording.mp4 \
+    --out /data/recording_inference.json
+
+# Extract VideoPrism embeddings from a video + mask
+singlebehaviorlab register --experiment /path/to/my_experiment \
+    --video /data/recording.mp4 \
+    --mask /data/recording_masks.h5 \
+    --out /data/recording_matrix.npz
+
+# Cluster an embedding matrix (loadable via the GUI's "Load Analysis State")
+singlebehaviorlab cluster --matrix /data/recording_matrix.npz \
+    --metadata /data/recording_matrix_metadata.npz \
+    --out /data/recording_clusters.pkl
+
+# Run SAM2 tracking from a prompts JSON exported in the Segmentation tab
+singlebehaviorlab segment --video /data/recording.mp4 \
+    --prompts /data/prompts.json \
+    --out /data/recording_masks.h5
+```
+
+Run `singlebehaviorlab <command> --help` for the full flag list on each subcommand. The GUI-only steps (labeling, refinement review, cluster inspection) still require the graphical interface; the CLI covers the batch-processing steps where no human input is needed.
+
+> **Full CLI reference:** [**CLI.md**](CLI.md) — detailed per-command docs, file-format reference, Python API, and troubleshooting.
+>
+> **Notebook demos:** [**demo/**](demo/) — two Jupyter notebooks walking through behavior sequencing and segmentation/clustering end-to-end. Drop your own demo video + prompts into `demo/data/` and step through the cells.
+>
+> **Try it in Colab — no install needed:** [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/alms93/SingleBehaviorLab/blob/main/demo/colab_segmentation_clustering.ipynb) Runs the full segmentation + clustering pipeline on a bundled demo video in a free Colab GPU runtime. Install takes ~5–10 minutes; the pipeline itself finishes in a few minutes.
+
 ---
 
 ## 4. First Launch

{singlebehaviorlab-2.0.0 → singlebehaviorlab-2.2.0}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "singlebehaviorlab"
-version = "2.0.0"
+version = "2.2.0"
 description = "Semi-automated behavioral video annotation, training, and analysis tool"
 readme = "README.md"
 license = { file = "LICENSE" }
@@ -54,7 +54,7 @@ Issues = "https://github.com/alms93/SingleBehaviorLab/issues"
 test = ["pytest", "pytest-cov"]
 
 [project.scripts]
-singlebehaviorlab = "singlebehaviorlab.__main__:main"
+singlebehaviorlab = "singlebehaviorlab.cli:main"
 
 [tool.pytest.ini_options]
 testpaths = ["tests"]

singlebehaviorlab-2.2.0/singlebehaviorlab/__init__.py

@@ -0,0 +1,58 @@
+"""SingleBehaviorLab — behavioral video annotation and training tool.
+
+A short Python API re-exports the backend pipeline functions at the package
+top level, so the typical usage is::
+
+    import singlebehaviorlab as sbl
+
+    sbl.segment(video, prompts, out_masks)
+    sbl.register(video, out_masks, out_matrix)
+    sbl.cluster(out_matrix, out_pkl, metadata_path=out_metadata)
+    sbl.plot_umap_clusters(out_pkl, show=True, save="umap.pdf")
+
+    sbl.train(experiment_dir, profile="balanced")
+    sbl.infer(model, video, out_json, experiment_dir=experiment_dir)
+
+The re-exports use lazy attribute loading (PEP 562 ``__getattr__``) so that
+``import singlebehaviorlab`` does not pull in torch, jax, tensorflow, sam2,
+or videoprism. Each symbol triggers its underlying backend module only on
+first access.
+"""
+
+__version__ = "2.2.0"
+__author__ = "Almir Aljovic"
+
+# Mapping of public name → (backend module, attribute name).
+_PUBLIC_API = {
+    "segment": ("singlebehaviorlab.backend.segmentation", "run_sam2_segmentation"),
+    "load_prompts_json": ("singlebehaviorlab.backend.segmentation", "load_prompts_json"),
+    "save_prompts_json": ("singlebehaviorlab.backend.segmentation", "save_prompts_json"),
+    "register": ("singlebehaviorlab.backend.registration", "run_registration"),
+    "RegistrationParams": ("singlebehaviorlab.backend.registration", "RegistrationParams"),
+    "cluster": ("singlebehaviorlab.backend.clustering", "run_clustering"),
+    "ClusteringParams": ("singlebehaviorlab.backend.clustering", "ClusteringParams"),
+    "plot_umap_clusters": ("singlebehaviorlab.backend.clustering", "plot_umap_clusters"),
+    "infer": ("singlebehaviorlab.backend.inference", "run_inference_on_video"),
+    "train": ("singlebehaviorlab.backend.training_runner", "run_training_session"),
+    "load_config": ("singlebehaviorlab.config", "load_config"),
+    "load_demo": ("singlebehaviorlab.demo", "load_demo"),
+    "DEMOS": ("singlebehaviorlab.demo", "DEMOS"),
+}
+
+__all__ = ["__version__", "__author__", *sorted(_PUBLIC_API)]
+
+
+def __getattr__(name):
+    try:
+        module_name, attr = _PUBLIC_API[name]
+    except KeyError as exc:
+        raise AttributeError(f"module 'singlebehaviorlab' has no attribute {name!r}") from exc
+    import importlib
+    module = importlib.import_module(module_name)
+    value = getattr(module, attr)
+    globals()[name] = value
+    return value
+
+
+def __dir__():
+    return sorted(set(list(globals().keys()) + list(_PUBLIC_API.keys())))

{singlebehaviorlab-2.0.0 → singlebehaviorlab-2.2.0}/singlebehaviorlab/__main__.py

@@ -7,71 +7,23 @@ Runs when invoked as:
 """
 
 import logging
-import sys
 import os
+import sys
 
 os.environ.setdefault("TF_CPP_MIN_LOG_LEVEL", "3")
 os.environ.setdefault("TF_ENABLE_ONEDNN_OPTS", "0")
 os.environ.setdefault("GRPC_VERBOSITY", "ERROR")
 os.environ.setdefault("GLOG_minloglevel", "2")
 
-# Let JAX grow GPU memory on demand and leave headroom for PyTorch.
 os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false"
 os.environ["XLA_PYTHON_CLIENT_MEM_FRACTION"] = "0.45"
-# Fall back to driver JIT compilation when ptxas/nvlink is unavailable.
 os.environ["XLA_FLAGS"] = "--xla_gpu_unsafe_fallback_to_driver_on_ptxas_not_found"
 
-import yaml
-from singlebehaviorlab._paths import get_default_config_path, get_experiments_dir
-from singlebehaviorlab.gui.main_window import MainWindow
-
-
-def load_config(config_path: str = None) -> dict:
-    """Load configuration from YAML file."""
-    if config_path is None:
-        config_path = str(get_default_config_path())
-
-    if os.path.exists(config_path):
-        with open(config_path, 'r') as f:
-            config = yaml.safe_load(f) or {}
-    else:
-        config = {}
-
-    # For path keys that are missing or blank, resolve relative to the
-    # experiments directory (pip install) or the package parent (source install).
-    from singlebehaviorlab._paths import get_package_dir
-    base_dir = str(get_package_dir().parent)
-
-    defaults = {
-        "data_dir":        os.path.join(base_dir, "data"),
-        "raw_videos_dir":  os.path.join(base_dir, "data", "raw_videos"),
-        "clips_dir":       os.path.join(base_dir, "data", "clips"),
-        "annotations_dir": os.path.join(base_dir, "data", "annotations"),
-        "models_dir":      os.path.join(base_dir, "models", "behavior_heads"),
-        "backbone_dir":    os.path.join(base_dir, "models", "videoprism_backbone"),
-        "annotation_file": os.path.join(base_dir, "data", "annotations", "annotations.json"),
-    }
-
-    for key, value in defaults.items():
-        if not config.get(key):
-            config[key] = value
-        elif not os.path.isabs(config[key]):
-            config[key] = os.path.join(base_dir, config[key])
+from singlebehaviorlab.config import load_config  # noqa: E402  (re-exported for backward compat)
 
-    experiments_dir = str(get_experiments_dir())
-    if not config.get("experiments_dir"):
-        config["experiments_dir"] = experiments_dir
-    os.makedirs(config["experiments_dir"], exist_ok=True)
 
-    config["config_path"] = config_path
-    config.setdefault("experiment_name", None)
-    config.setdefault("experiment_path", None)
-
-    return config
-
-
-def main():
-    """Application entry point."""
+def run_gui_app():
+    """Launch the PyQt6 graphical interface."""
     logging.basicConfig(
         level=logging.WARNING,
         format="%(levelname)s [%(name)s] %(message)s",
@@ -79,6 +31,7 @@ def main():
     from PyQt6.QtWidgets import (
         QApplication, QDialog, QVBoxLayout, QPushButton, QLabel,
     )
+    from singlebehaviorlab.gui.main_window import MainWindow
 
     app = QApplication(sys.argv)
     app.setApplicationName("SingleBehaviorLab")
@@ -126,5 +79,10 @@ def main():
     sys.exit(app.exec())
 
 
+def main():
+    from singlebehaviorlab.cli import main as cli_main
+    cli_main()
+
+
 if __name__ == "__main__":
     main()

singlebehaviorlab-2.2.0/singlebehaviorlab/backend/clustering.py

@@ -0,0 +1,316 @@
+"""Headless UMAP + Leiden/HDBSCAN clustering for the CLI.
+
+Produces a pickle file that the GUI Clustering tab loads via the existing
+"Load Analysis State" action. The state schema matches
+``clustering_widget._save_analysis_state``.
+"""
+
+from __future__ import annotations
+
+import os
+import pickle
+from dataclasses import dataclass
+from datetime import datetime
+from pathlib import Path
+from typing import Any, Callable, Optional, Union
+
+import numpy as np
+import pandas as pd
+
+__all__ = ["ClusteringParams", "run_clustering", "plot_umap_clusters"]
+
+
+@dataclass
+class ClusteringParams:
+    """Knobs for ``run_clustering``. Defaults mirror the GUI sliders."""
+
+    method: str = "leiden"
+    n_components: int = 2
+    n_neighbors: int = 15
+    min_dist: float = 0.1
+    normalization: str = "standard"  # standard | minmax | l2 | none
+    leiden_resolution: float = 1.0
+    leiden_k: int = 15
+    min_cluster_size: int = 10
+    min_samples: int = 5
+    hdbscan_epsilon: float = 0.0
+
+
+def _load_matrix_metadata(
+    matrix_path: str,
+    metadata_path: Optional[str],
+) -> tuple[pd.DataFrame, Optional[pd.DataFrame]]:
+    if matrix_path.endswith(".npz"):
+        npz = np.load(matrix_path, allow_pickle=True)
+        matrix = npz["matrix"]
+        feature_names = npz["feature_names"]
+        snippet_ids = npz["snippet_ids"] if "snippet_ids" in npz else npz.get("span_ids")
+        if snippet_ids is None:
+            snippet_ids = np.array([f"snippet{i + 1}" for i in range(matrix.shape[1])])
+        matrix_df = pd.DataFrame(matrix, index=feature_names, columns=snippet_ids)
+    elif matrix_path.endswith(".parquet"):
+        matrix_df = pd.read_parquet(matrix_path, engine="pyarrow")
+    else:
+        matrix_df = pd.read_csv(matrix_path, index_col=0)
+
+    metadata_df: Optional[pd.DataFrame] = None
+    if metadata_path:
+        if metadata_path.endswith(".npz"):
+            meta_npz = np.load(metadata_path, allow_pickle=True)
+            metadata_df = pd.DataFrame(meta_npz["metadata"], columns=meta_npz["columns"])
+        elif metadata_path.endswith(".parquet"):
+            metadata_df = pd.read_parquet(metadata_path, engine="pyarrow")
+        else:
+            metadata_df = pd.read_csv(metadata_path)
+    return matrix_df, metadata_df
+
+
+def _normalize(X: pd.DataFrame, method: str) -> pd.DataFrame:
+    X = X.replace([np.inf, -np.inf], np.nan).fillna(0.0)
+    method = method.lower()
+    if method == "standard":
+        from sklearn.preprocessing import StandardScaler
+        arr = StandardScaler().fit_transform(X)
+    elif method == "minmax":
+        from sklearn.preprocessing import MinMaxScaler
+        arr = MinMaxScaler().fit_transform(X)
+    elif method == "l2":
+        from sklearn.preprocessing import Normalizer
+        arr = Normalizer(norm="l2").fit_transform(X)
+    elif method in ("none", "raw"):
+        arr = X.values
+    else:
+        raise ValueError(f"Unknown normalization method: {method}")
+    return pd.DataFrame(arr, index=X.index, columns=X.columns)
+
+
+def _run_umap(
+    data: pd.DataFrame,
+    n_components: int,
+    n_neighbors: int,
+    min_dist: float,
+) -> np.ndarray:
+    import umap
+    reducer = umap.UMAP(
+        n_neighbors=n_neighbors,
+        min_dist=min_dist,
+        n_components=n_components,
+        random_state=42,
+    )
+    return reducer.fit_transform(data.values)
+
+
+def _run_leiden(data: pd.DataFrame, leiden_k: int, resolution: float) -> np.ndarray:
+    from sklearn.neighbors import kneighbors_graph
+    import igraph as ig
+    import leidenalg as la
+
+    knn = kneighbors_graph(data.values, n_neighbors=leiden_k, mode="connectivity", include_self=False)
+    sources, targets = knn.nonzero()
+    edges = list(zip(sources.tolist(), targets.tolist()))
+    graph = ig.Graph(n=data.shape[0], edges=edges, directed=False)
+    partition = la.find_partition(
+        graph,
+        la.RBConfigurationVertexPartition,
+        resolution_parameter=resolution,
+    )
+    return np.array(partition.membership)
+
+
+def _run_hdbscan(data: pd.DataFrame, params: ClusteringParams) -> np.ndarray:
+    import hdbscan
+    clusterer = hdbscan.HDBSCAN(
+        min_cluster_size=params.min_cluster_size,
+        min_samples=params.min_samples,
+        cluster_selection_epsilon=params.hdbscan_epsilon,
+    )
+    return clusterer.fit_predict(data.values)
+
+
+def run_clustering(
+    matrix_path: str | os.PathLike[str],
+    output_path: str | os.PathLike[str],
+    *,
+    metadata_path: Optional[str | os.PathLike[str]] = None,
+    params: Optional[ClusteringParams] = None,
+    log_fn: Optional[Callable[[str], None]] = None,
+) -> str:
+    """Cluster a feature matrix and write a GUI-loadable analysis pickle.
+
+    Returns the written pickle path.
+    """
+    params = params or ClusteringParams()
+
+    matrix_path_str = str(Path(matrix_path).expanduser().resolve())
+    metadata_path_str: Optional[str] = None
+    if metadata_path:
+        metadata_path_str = str(Path(metadata_path).expanduser().resolve())
+
+    output_path_obj = Path(output_path).expanduser().resolve()
+    output_path_obj.parent.mkdir(parents=True, exist_ok=True)
+
+    def _log(msg: str) -> None:
+        if log_fn:
+            log_fn(msg)
+
+    _log(f"Loading matrix: {matrix_path_str}")
+    matrix_df, metadata_df = _load_matrix_metadata(matrix_path_str, metadata_path_str)
+    _log(f"Matrix shape: {matrix_df.shape[0]} features × {matrix_df.shape[1]} samples")
+
+    processed = _normalize(matrix_df.T, params.normalization)
+    _log(f"Processed shape: {processed.shape} (samples × features); normalization={params.normalization}")
+
+    _log(
+        f"Running UMAP (n_neighbors={params.n_neighbors}, "
+        f"min_dist={params.min_dist}, n_components={params.n_components})"
+    )
+    embedding = _run_umap(
+        processed,
+        n_components=params.n_components,
+        n_neighbors=params.n_neighbors,
+        min_dist=params.min_dist,
+    )
+
+    if params.method == "leiden":
+        _log(f"Running Leiden clustering (k={params.leiden_k}, resolution={params.leiden_resolution})")
+        clusters = _run_leiden(processed, params.leiden_k, params.leiden_resolution)
+    elif params.method == "hdbscan":
+        _log(
+            f"Running HDBSCAN (min_cluster_size={params.min_cluster_size}, "
+            f"min_samples={params.min_samples}, epsilon={params.hdbscan_epsilon})"
+        )
+        clusters = _run_hdbscan(processed, params)
+    else:
+        raise ValueError(f"Unknown clustering method: {params.method}")
+
+    unique_clusters = sorted(set(int(c) for c in clusters))
+    _log(f"Clusters found: {len(unique_clusters)} (labels: {unique_clusters})")
+
+    snippet_to_clip_map: dict[str, str] = {}
+    if metadata_df is not None and "clip_path" in metadata_df.columns:
+        snippet_col = (
+            "snippet"
+            if "snippet" in metadata_df.columns
+            else ("span_id" if "span_id" in metadata_df.columns else None)
+        )
+        if snippet_col is not None:
+            for _, row in metadata_df.iterrows():
+                snippet_id = str(row.get(snippet_col, "")).strip()
+                clip_path_val = str(row.get("clip_path", "")).strip()
+                if snippet_id and clip_path_val and os.path.exists(clip_path_val):
+                    snippet_to_clip_map[snippet_id] = clip_path_val
+        _log(f"Built snippet→clip map with {len(snippet_to_clip_map)} entries from metadata.")
+
+    state = {
+        "matrix_data": matrix_df,
+        "metadata": metadata_df,
+        "processed_data": processed,
+        "embedding": embedding,
+        "clusters": clusters,
+        "selected_features": list(matrix_df.index),
+        "snippet_to_clip_map": snippet_to_clip_map,
+        "metadata_file_path": metadata_path_str,
+        "timestamp": datetime.now().strftime("%Y%m%d_%H%M%S"),
+        "version": "1.0",
+    }
+
+    with open(output_path_obj, "wb") as f:
+        pickle.dump(state, f)
+    _log(f"Wrote analysis state: {output_path_obj}")
+    return str(output_path_obj)
+
+
+def plot_umap_clusters(
+    state: Union[dict[str, Any], str, os.PathLike[str]],
+    *,
+    show: bool = False,
+    save: Optional[str | os.PathLike[str]] = None,
+    title: Optional[str] = None,
+    figsize: tuple[float, float] = (8.0, 6.0),
+    point_size: float = 6.0,
+):
+    """Plot the UMAP embedding produced by ``run_clustering``.
+
+    Follows the scanpy ``sc.pl.umap`` convention: pass ``show=True`` to pop up
+    an interactive window, pass ``save='foo.pdf'`` to write the figure to
+    disk (format inferred from the extension). Both can be combined.
+
+    Args:
+        state: Either an analysis-state dict (with ``embedding`` and
+            ``clusters`` keys) or the path to a ``.pkl`` file produced by
+            ``run_clustering``.
+        show: Call ``plt.show()`` to open an interactive window.
+        save: Destination path for the rendered figure. PDF, PNG, and SVG
+            are all supported.
+        title: Optional plot title.
+        figsize: Figure size in inches.
+        point_size: Scatter marker size.
+
+    Returns:
+        The matplotlib ``Figure`` object so the caller can further customise
+        or embed it.
+    """
+    import matplotlib
+    if not show and matplotlib.get_backend().lower() not in {"agg", "pdf", "svg", "ps"}:
+        matplotlib.use("Agg")
+    import matplotlib.pyplot as plt
+
+    if isinstance(state, (str, os.PathLike)):
+        with open(state, "rb") as f:
+            state = pickle.load(f)
+    if not isinstance(state, dict):
+        raise TypeError("state must be a dict or a path to a clustering .pkl file.")
+
+    embedding = np.asarray(state.get("embedding"))
+    clusters = np.asarray(state.get("clusters"))
+    if embedding is None or embedding.ndim != 2 or embedding.shape[1] < 2:
+        raise ValueError("Analysis state is missing a valid 2D+ UMAP embedding.")
+    if clusters is None or len(clusters) != embedding.shape[0]:
+        raise ValueError("Cluster labels are missing or the wrong length.")
+
+    unique_labels = sorted({int(c) for c in clusters})
+    non_noise = [c for c in unique_labels if c >= 0]
+    cmap = plt.get_cmap("tab20", max(len(non_noise), 1))
+
+    fig, ax = plt.subplots(figsize=figsize)
+    noise_mask = clusters < 0
+    if np.any(noise_mask):
+        ax.scatter(
+            embedding[noise_mask, 0],
+            embedding[noise_mask, 1],
+            s=point_size,
+            color="lightgray",
+            alpha=0.5,
+            label="Noise",
+            linewidths=0,
+        )
+    for i, c in enumerate(non_noise):
+        mask = clusters == c
+        ax.scatter(
+            embedding[mask, 0],
+            embedding[mask, 1],
+            s=point_size,
+            color=cmap(i % cmap.N),
+            alpha=0.85,
+            label=f"Cluster {c}",
+            linewidths=0,
+        )
+
+    ax.set_xlabel("UMAP1")
+    ax.set_ylabel("UMAP2")
+    if title:
+        ax.set_title(title)
+    else:
+        ax.set_title(f"UMAP + {'Noise + ' if np.any(noise_mask) else ''}{len(non_noise)} clusters")
+    ax.set_aspect("equal", adjustable="datalim")
+    ax.legend(loc="best", fontsize="x-small", framealpha=0.85, markerscale=1.5)
+    fig.tight_layout()
+
+    if save:
+        save_path = Path(save).expanduser().resolve()
+        save_path.parent.mkdir(parents=True, exist_ok=True)
+        fig.savefig(save_path, bbox_inches="tight")
+    if show:
+        plt.show()
+
+    return fig