senoquant 1.0.0b2__py3-none-any.whl → 1.0.0b4__py3-none-any.whl

senoquant/tabs/spots/ufish_utils/core.py

@@ -0,0 +1,387 @@
+"""UFish-based spot enhancement utilities.
+
+This module wraps UFish inference for SenoQuant spot detection workflows.
+It handles:
+
+- optional import of UFish from site-packages or vendored sources,
+- ONNX Runtime execution-provider selection,
+- model/weights caching between calls, and
+- default ONNX weight retrieval from the SenoQuant Hugging Face model repo.
+
+Notes
+-----
+Weight loading priority is:
+
+1. explicit ``UFishConfig.weights_path``,
+2. legacy ``UFishConfig.load_from_internet``, then
+3. default ``ufish.onnx`` resolved via :func:`ensure_hf_model`.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+import logging
+from pathlib import Path
+import sys
+from types import MethodType
+from typing import TYPE_CHECKING, Any, cast
+
+import numpy as np
+
+from senoquant.tabs.segmentation.models.hf import (
+    DEFAULT_REPO_ID,
+    ensure_hf_model,
+)
+
+try:  # pragma: no cover - optional dependency
+    from ufish.api import UFish
+except ImportError:  # pragma: no cover - optional dependency
+    _repo_root = Path(__file__).resolve().parents[5]
+    _vendored_root = _repo_root / "_vendor" / "ufish"
+    if _vendored_root.exists():
+        vendored_root = str(_vendored_root)
+        if vendored_root not in sys.path:
+            sys.path.insert(0, vendored_root)
+        try:
+            from ufish.api import UFish
+        except ImportError:
+            UFish = None
+    else:
+        UFish = None
+
+try:  # pragma: no cover - optional dependency
+    import onnxruntime as ort
+except ImportError:  # pragma: no cover - optional dependency
+    ort = None
+
+if TYPE_CHECKING:  # pragma: no cover - typing only
+    from ufish.api import UFish as UFishType
+
+
+@dataclass(slots=True)
+class UFishConfig:
+    """Configuration for UFish enhancement.
+
+    Parameters
+    ----------
+    weights_path : str or None, optional
+        Explicit local path to ONNX/PyTorch weights. When provided, this path
+        is used directly and takes precedence over all other loading modes.
+    load_from_internet : bool, optional
+        Legacy compatibility mode that calls
+        ``UFish.load_weights_from_internet()`` directly.
+    device : {"cuda", "dml", "mps"} or None, optional
+        Preferred accelerator mode used to influence ONNX Runtime provider
+        ordering when constructing UFish sessions.
+    """
+
+    weights_path: str | None = None
+    load_from_internet: bool = False
+    device: str | None = None
+
+
+class _UFishState:
+    """In-process cache for the UFish model and loaded weights."""
+
+    def __init__(self) -> None:
+        """Initialize empty cached state."""
+        self.model: UFishType | None = None
+        self.weights_loaded = False
+        self.device: str | None = None
+        self.weights_path: str | None = None
+
+
+_UFISH_STATE = _UFishState()
+_UFISH_HF_FILENAME = "ufish.onnx"
+_LOGGER = logging.getLogger(__name__)
+
+
+def _ensure_ufish_available() -> None:
+    """Raise a helpful error when UFish cannot be imported.
+
+    Raises
+    ------
+    ImportError
+        If the ``ufish`` package is unavailable from both normal and vendored
+        import locations.
+    """
+    if UFish is None:  # pragma: no cover - import guard
+        msg = "ufish is required for spot enhancement."
+        raise ImportError(msg)
+
+
+def _resolve_default_weights_path() -> Path:
+    """Resolve the default UFish ONNX path.
+
+    Returns
+    -------
+    pathlib.Path
+        Local path to ``ufish.onnx``. The file is downloaded to the
+        ``ufish_utils`` directory if it does not already exist.
+
+    Raises
+    ------
+    RuntimeError
+        If the Hugging Face download helper is unavailable or download fails.
+    """
+    target_dir = Path(__file__).resolve().parent
+    return ensure_hf_model(
+        _UFISH_HF_FILENAME,
+        target_dir,
+        repo_id=DEFAULT_REPO_ID,
+    )
+
+
+def _preferred_providers() -> list[str]:
+    """Return ONNX Runtime providers ordered by GPU preference.
+
+    Returns
+    -------
+    list[str]
+        Providers available in the current runtime, ordered from most
+        preferred accelerator to CPU fallback.
+    """
+    if ort is None:
+        return []
+    available = set(ort.get_available_providers())
+    preferred = [
+        "CUDAExecutionProvider",
+        "ROCMExecutionProvider",
+        "DmlExecutionProvider",
+        "DirectMLExecutionProvider",
+        "CoreMLExecutionProvider",
+        "CPUExecutionProvider",
+    ]
+    providers = [provider for provider in preferred if provider in available]
+    return providers or list(available)
+
+
+def _select_onnx_providers(device: str | None) -> list[str]:
+    """Choose execution providers for a requested device hint.
+
+    Parameters
+    ----------
+    device : str or None
+        Device hint from :class:`UFishConfig`.
+
+    Returns
+    -------
+    list[str]
+        Provider names to pass to ``onnxruntime.InferenceSession``.
+    """
+    preferred = _preferred_providers()
+    if not device:
+        return preferred
+    if device == "cuda":
+        return [
+            p
+            for p in preferred
+            if p in {"CUDAExecutionProvider", "CPUExecutionProvider"}
+        ]
+    if device == "dml":
+        return [
+            p
+            for p in preferred
+            if p
+            in {
+                "DmlExecutionProvider",
+                "DirectMLExecutionProvider",
+                "CPUExecutionProvider",
+            }
+        ]
+    if device == "mps":
+        return [
+            p
+            for p in preferred
+            if p in {"CoreMLExecutionProvider", "CPUExecutionProvider"}
+        ]
+    return preferred
+
+
+def _patch_onnx_loader(model: UFishType) -> None:
+    """Monkey-patch UFish ONNX loader to use SenoQuant provider selection.
+
+    Parameters
+    ----------
+    model : UFishType
+        UFish instance whose private ``_load_onnx`` method will be replaced.
+    """
+    if ort is None:
+        return
+    ort_any = cast("Any", ort)
+
+    def _load_onnx(
+        self: UFishType,
+        onnx_path: str,
+        providers: list[str] | None = None,
+    ) -> None:
+        providers = providers or _select_onnx_providers(
+            getattr(self, "_device", None),
+        )
+        self.ort_session = ort_any.InferenceSession(
+            str(onnx_path),
+            providers=providers,
+        )
+        self.model = None
+
+    model._load_onnx = MethodType(_load_onnx, model)  # noqa: SLF001
+
+
+def _get_ufish(config: UFishConfig) -> UFishType:
+    """Return a cached UFish instance for the requested configuration.
+
+    Parameters
+    ----------
+    config : UFishConfig
+        Runtime configuration used to determine whether the cached model can be
+        reused or must be re-instantiated.
+
+    Returns
+    -------
+    UFishType
+        Ready-to-use UFish instance with patched ONNX loading behavior.
+    """
+    _ensure_ufish_available()
+    if _UFISH_STATE.model is None or _UFISH_STATE.device != config.device:
+        ufish_cls = cast("type[UFishType]", UFish)
+        ufish_any = cast("Any", ufish_cls)
+        if config.device:
+            _UFISH_STATE.model = ufish_any(device=config.device)
+        else:
+            _UFISH_STATE.model = ufish_any()
+        _patch_onnx_loader(cast("UFishType", _UFISH_STATE.model))
+        _UFISH_STATE.weights_loaded = False
+        _UFISH_STATE.device = config.device
+        _UFISH_STATE.weights_path = None
+    return cast("UFishType", _UFISH_STATE.model)
+
+
+def _ensure_weights(model: UFishType, config: UFishConfig) -> None:
+    """Ensure model weights are loaded according to configuration.
+
+    Parameters
+    ----------
+    model : UFishType
+        Active UFish model instance.
+    config : UFishConfig
+        Weight source and device settings.
+
+    Raises
+    ------
+    RuntimeError
+        If neither Hugging Face/default loading nor fallback loading succeeds.
+    """
+    if config.weights_path:
+        weights_path = Path(config.weights_path).expanduser().resolve()
+        resolved_path = str(weights_path)
+        if _UFISH_STATE.weights_loaded and _UFISH_STATE.weights_path == resolved_path:
+            return
+        model.load_weights(resolved_path)
+        _UFISH_STATE.weights_loaded = True
+        _UFISH_STATE.weights_path = resolved_path
+        return
+
+    if config.load_from_internet:
+        if _UFISH_STATE.weights_loaded and _UFISH_STATE.weights_path == "internet":
+            return
+        model.load_weights_from_internet()
+        _UFISH_STATE.weights_loaded = True
+        _UFISH_STATE.weights_path = "internet"
+        return
+
+    try:
+        weights_path = _resolve_default_weights_path()
+    except RuntimeError:
+        # Keep legacy behavior when HF download dependencies are unavailable.
+        if _UFISH_STATE.weights_loaded and _UFISH_STATE.weights_path == "default":
+            return
+        model.load_weights()
+        _UFISH_STATE.weights_loaded = True
+        _UFISH_STATE.weights_path = "default"
+        return
+    except Exception as exc:
+        try:
+            if _UFISH_STATE.weights_loaded and _UFISH_STATE.weights_path == "default":
+                return
+            model.load_weights()
+            _UFISH_STATE.weights_loaded = True
+            _UFISH_STATE.weights_path = "default"
+            return
+        except Exception:
+            msg = (
+                "Could not load UFish weights from local default or Hugging Face. "
+                "Provide UFishConfig(weights_path=...) or ensure model download access."
+            )
+            raise RuntimeError(msg) from exc
+
+    resolved_path = str(weights_path)
+    if _UFISH_STATE.weights_loaded and _UFISH_STATE.weights_path == resolved_path:
+        return
+    model.load_weights(resolved_path)
+    _UFISH_STATE.weights_loaded = True
+    _UFISH_STATE.weights_path = resolved_path
+
+
+def enhance_image(
+    image: np.ndarray,
+    *,
+    config: UFishConfig | None = None,
+) -> np.ndarray:
+    """Enhance an image using UFish.
+
+    Parameters
+    ----------
+    image : numpy.ndarray
+        Input image array. UFish supports 2D images and common 3D stack
+        layouts handled by UFish internally.
+    config : UFishConfig or None, optional
+        Optional runtime configuration. If omitted, default behavior is used.
+
+    Returns
+    -------
+    numpy.ndarray
+        Enhanced image produced by UFish with the same dimensionality as the
+        input image.
+
+    Raises
+    ------
+    ImportError
+        If UFish cannot be imported.
+    RuntimeError
+        If weights cannot be loaded from configured/default sources.
+    """
+    if config is None:
+        config = UFishConfig()
+    model = _get_ufish(config)
+    _ensure_weights(model, config)
+    image = np.asarray(image)
+    model_any = cast("Any", model)
+    predict_chunks = getattr(model_any, "predict_chunks", None)
+
+    def _run_inference() -> tuple[Any, Any]:
+        if callable(predict_chunks):
+            return predict_chunks(image)
+        return model_any.predict(image)
+
+    try:
+        _pred_spots, enhanced = _run_inference()
+    except Exception as exc:
+        # If ONNX inference fails on GPU, retry once on CPU provider.
+        weight_path = getattr(model_any, "weight_path", None)
+        can_retry_on_cpu = (
+            ort is not None
+            and isinstance(weight_path, str)
+            and weight_path.endswith(".onnx")
+            and hasattr(model_any, "_load_onnx")
+            and "CPUExecutionProvider" in set(ort.get_available_providers())
+        )
+        if not can_retry_on_cpu:
+            raise
+        _LOGGER.warning(
+            "UFish inference failed; retrying with ONNX CPUExecutionProvider. "
+            "Original error: %s",
+            exc,
+        )
+        model_any._load_onnx(weight_path, providers=["CPUExecutionProvider"])
+        _pred_spots, enhanced = _run_inference()
+    return np.asarray(enhanced)

senoquant/tabs/visualization/__init__.py

@@ -0,0 +1 @@
+"""Visualization tab modules."""

senoquant/tabs/visualization/backend.py

@@ -0,0 +1,306 @@
+"""Backend logic for the Visualization tab."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Iterable
+import shutil
+import tempfile
+
+from .plots import PlotConfig
+
+
+@dataclass
+class PlotExportResult:
+    """Output metadata for a single plot export.
+
+    Attributes
+    ----------
+    plot_id : str
+        Stable identifier for the exported plot instance.
+    plot_type : str
+        Plot type name used for routing (e.g., ``"UMAP"``).
+    temp_dir : Path
+        Temporary directory where the plot wrote its outputs.
+    outputs : list of Path
+        Explicit file paths returned by the plot processor.
+    """
+
+    plot_id: str
+    plot_type: str
+    temp_dir: Path
+    outputs: list[Path] = field(default_factory=list)
+
+
+@dataclass
+class VisualizationResult:
+    """Aggregated output information for a visualization run.
+
+    Attributes
+    ----------
+    input_root : Path
+        Root input directory.
+    output_root : Path
+        Root output directory for the run.
+    temp_root : Path
+        Temporary root directory used during processing.
+    plot_outputs : list of PlotExportResult
+        Per-plot export metadata for the run.
+    """
+    input_root: Path
+    output_root: Path
+    temp_root: Path
+    plot_outputs: list[PlotExportResult]
+
+
+class VisualizationBackend:
+    """Backend orchestrator for visualization exports.
+
+    Notes
+    -----
+    Feature export routines live with their feature implementations. The
+    backend iterates through configured feature contexts, asks each feature
+    handler to export into a temporary directory, and then routes those
+    outputs into a final output structure.
+    """
+
+    def __init__(self) -> None:
+        """Initialize the backend state.
+
+        Attributes
+        ----------
+        metrics : list
+            Placeholder container for computed metrics.
+        """
+        self.metrics: list[object] = []
+
+    def process(
+        self,
+        plots: Iterable[object],
+        input_path: Path,
+        output_path: str,
+        output_name: str,
+        export_format: str,
+        markers: list[str] | None = None,
+        thresholds: dict[str, float] | None = None,
+        save: bool = True,
+        cleanup: bool = True,
+    ) -> VisualizationResult:
+        """Run plot exports and route their outputs.
+
+        Parameters
+        ----------
+        plots : iterable of object
+            Plot UI contexts with ``state`` and ``plot_handler``.
+            Each handler should implement ``plot(temp_dir, input_path, export_format)``.
+        input_path : Path
+            Path to the input folder containing quantification files.
+        output_path : str
+            Base output folder path.
+        output_name : str
+            Folder name used to group exported outputs.
+        export_format : str
+            File format requested by the user (``"png"`` or ``"svg"``).
+        markers : list of str, optional
+            List of selected markers to include.
+        thresholds : dict, optional
+            Dictionary of {marker_name: threshold_value} for filtering.
+        save : bool, optional
+            Whether to save (route) the outputs to the final destination immediately.
+        cleanup : bool, optional
+            Whether to delete temporary export folders after routing.
+
+        Returns
+        -------
+        VisualizationResult
+            Output metadata for the completed run.
+
+        Notes
+        -----
+        If a plot export does not return explicit output paths, the backend
+        will move all files found in the plot's temp directory. This allows
+        plot implementations to either return specific files or simply write
+        into the provided temporary directory.
+        """
+        input_root = input_path.parent
+        # Treat `output_path` as the folder and `output_name` as an optional
+        # filename base. Resolve output_root without using output_name so
+        # output_name can be applied as a file name rather than a subfolder.
+        output_root = self._resolve_output_root(output_path, "")
+        output_root.mkdir(parents=True, exist_ok=True)
+        temp_root = Path(tempfile.mkdtemp(prefix="senoquant-plot-"))
+
+        plot_outputs: list[PlotExportResult] = []
+        for context in plots:
+            plot = getattr(context, "state", None)
+            handler = getattr(context, "plot_handler", None)
+            if not isinstance(plot, PlotConfig):
+                continue
+            print(f"[Backend] Processing plot: {plot.type_name}")
+            print(f"[Backend] Handler: {handler}")
+            print(f"[Backend] Handler has plot method: {hasattr(handler, 'plot') if handler else False}")
+            temp_dir = temp_root / plot.plot_id
+            temp_dir.mkdir(parents=True, exist_ok=True)
+            outputs: list[Path] = []
+            if handler is not None and hasattr(handler, "plot"):
+                print(f"[Backend] Calling handler.plot() with input_path={input_path}, format={export_format}")
+                outputs = [
+                    Path(path)
+                    for path in handler.plot(
+                        temp_dir, 
+                        input_path, 
+                        export_format,
+                        markers=markers,
+                        thresholds=thresholds
+                    )
+                ]
+                print(f"[Backend] Handler returned {len(outputs)} outputs: {outputs}")
+            else:
+                print(f"[Backend] Skipping: handler is None or has no plot method")
+            plot_outputs.append(
+                PlotExportResult(
+                    plot_id=plot.plot_id,
+                    plot_type=plot.type_name,
+                    temp_dir=temp_dir,
+                    outputs=outputs,
+                )
+            )
+
+        if save:
+            print(f"[Backend] About to route {len(plot_outputs)} plot outputs")
+            self._route_plot_outputs(output_root, plot_outputs, output_name)
+        if cleanup:
+            shutil.rmtree(temp_root, ignore_errors=True)
+        return VisualizationResult(
+            input_root=input_root,
+            output_root=output_root,
+            temp_root=temp_root,
+            plot_outputs=plot_outputs,
+        )
+
+    def save_result(
+        self, 
+        result: VisualizationResult, 
+        output_path: str, 
+        output_name: str
+    ) -> None:
+        """Save an existing visualization result to the specified output.
+
+        This moves/copies files from the result's temporary directory (or previous location)
+        to the new output path.
+        """
+        output_root = self._resolve_output_root(output_path, "")
+        output_root.mkdir(parents=True, exist_ok=True)
+        result.output_root = output_root
+        self._route_plot_outputs(output_root, result.plot_outputs, output_name)
+
+    def _resolve_output_root(self, output_path: str, output_name: str) -> Path:
+        """Resolve the final output root directory.
+
+        Parameters
+        ----------
+        output_path : str
+            Base output folder path.
+        output_name : str
+            Folder name used to group exported outputs.
+
+        Returns
+        -------
+        Path
+            Resolved output directory path.
+        """
+        if output_path and output_path.strip():
+            base = Path(output_path)
+        else:
+            # Default to repository root (current working directory)
+            base = Path.cwd()
+        if output_name and output_name.strip():
+            return base / output_name
+        return base
+
+    def _route_plot_outputs(
+        self,
+        output_root: Path,
+        plot_outputs: Iterable[PlotExportResult],
+        output_name: str = "",
+    ) -> None:
+        """Move plot outputs from temp folders to the final location.
+
+        Parameters
+        ----------
+        output_root : Path
+            Destination root folder.
+        plot_outputs : iterable of PlotExportResult
+            Export results to route.
+
+        Notes
+        -----
+        When a plot returns no explicit output list, all files present
+        in the temporary directory are routed instead. Subdirectories are
+        not traversed.
+        """
+        for plot_output in plot_outputs:
+            print(f"[Backend] Routing {plot_output.plot_type} to {output_root}")
+            final_paths: list[Path] = []
+            outputs = plot_output.outputs
+            # Choose source list: explicit outputs if provided, otherwise files
+            # from the temp directory.
+            source_files: list[Path] = []
+            if outputs:
+                source_files = [p for p in outputs if Path(p).exists()]
+            else:
+                source_files = [p for p in plot_output.temp_dir.glob("*") if p.is_file()]
+
+            if not source_files:
+                print(f"[Backend]   No files to route for {plot_output.plot_type}")
+                plot_output.outputs = []
+                continue
+
+            # If the caller provided output_name, use it as the base filename.
+            for idx, src in enumerate(source_files):
+                src = Path(src)
+                ext = src.suffix
+                if output_name and output_name.strip():
+                    # If multiple files, append an index to avoid collisions.
+                    if len(source_files) == 1:
+                        dest_name = f"{output_name}{ext}"
+                    else:
+                        dest_name = f"{output_name}_{idx+1}{ext}"
+                else:
+                    # Fallback: prefix with plot type for clarity
+                    safe_type = plot_output.plot_type.replace(' ', '_')
+                    dest_name = f"{safe_type}_{src.name}"
+                dest = output_root / dest_name
+                print(f"[Backend]   Copying {src} -> {dest}")
+                try:
+                    shutil.copy2(str(src), dest)
+                except shutil.SameFileError:
+                    print(f"[Backend]   Skipping copy: source and destination are the same ({dest})")
+                final_paths.append(dest)
+
+            # Update plot_output.outputs to point at final routed files
+            plot_output.outputs = final_paths
+
+    def _plot_dir_name(self, plot_output: PlotExportResult) -> str:
+        """Build a filesystem-friendly folder name for a plot.
+
+        Parameters
+        ----------
+        plot_output : PlotExportResult
+            Export result metadata.
+
+        Returns
+        -------
+        str
+            Directory name for the plot outputs.
+
+        Notes
+        -----
+        Non-alphanumeric characters are replaced to avoid filesystem issues.
+        """
+        name = plot_output.plot_type.strip()
+        safe = "".join(
+            char if char.isalnum() or char in "-_ " else "_" for char in name
+        )
+        return safe.replace(" ", "_").lower()