tono/backend/nodes/neural_classification.py

"""Neural network classification — classify pixels using a simple feedforward network."""

from __future__ import annotations

import numpy as np
from scipy.ndimage import gaussian_filter

from backend.node_registry import register_node
from backend.data_types import DataField
from backend.nodes.helpers import mask_to_bool, bool_to_mask


def _sigmoid(x: np.ndarray) -> np.ndarray:
    """Numerically stable sigmoid."""
    return np.where(
        x >= 0,
        1.0 / (1.0 + np.exp(-x)),
        np.exp(x) / (1.0 + np.exp(x)),
    )


def _extract_features(data: np.ndarray, n_gaussians: int) -> np.ndarray:
    """Build multi-scale Gaussian feature matrix from 2-D data.

    For each scale sigma = 2^i (i = 0 .. n_gaussians-1), compute
    gaussian_filter(data, sigma) and stack as feature columns.
    Each feature is normalised to zero mean and unit variance.
    """
    rows, cols = data.shape
    features = np.empty((rows * cols, n_gaussians), dtype=np.float64)

    for i in range(n_gaussians):
        sigma = 2.0 ** i
        blurred = gaussian_filter(data, sigma).ravel()
        mean = blurred.mean()
        std = blurred.std()
        if std > 0:
            blurred = (blurred - mean) / std
        else:
            blurred = blurred - mean
        features[:, i] = blurred

    return features


def _forward(X: np.ndarray, W1: np.ndarray, b1: np.ndarray,
             W2: np.ndarray, b2: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    """Forward pass through a 2-layer sigmoid network."""
    h = _sigmoid(X @ W1 + b1)
    y = _sigmoid(h @ W2 + b2)
    return h, y


def _train_network(
    X: np.ndarray,
    targets: np.ndarray,
    W1: np.ndarray,
    b1: np.ndarray,
    W2: np.ndarray,
    b2: np.ndarray,
    train_steps: int,
    lr: float = 0.1,
) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    """Train via gradient descent with binary cross-entropy loss."""
    eps = 1e-7
    n = X.shape[0]
    t = targets.reshape(-1, 1)

    for _ in range(train_steps):
        # Forward
        h, y = _forward(X, W1, b1, W2, b2)

        # Clamp to avoid log(0)
        y_clamped = np.clip(y, eps, 1.0 - eps)

        # Backward — output layer
        dy = (y_clamped - t) / (y_clamped * (1.0 - y_clamped) + eps)
        dy *= y * (1.0 - y)  # sigmoid derivative

        dW2 = (h.T @ dy) / n
        db2 = dy.mean(axis=0)

        # Backward — hidden layer
        dh = (dy @ W2.T) * h * (1.0 - h)
        dW1 = (X.T @ dh) / n
        db1 = dh.mean(axis=0)

        # Update
        W1 -= lr * dW1
        b1 -= lr * db1
        W2 -= lr * dW2
        b2 -= lr * db2

    return W1, b1, W2, b2


@register_node(display_name="Neural Classification")
class NeuralClassification:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "field": ("DATA_FIELD",),
                "n_gaussians": ("INT", {"default": 4, "min": 1, "max": 10, "step": 1}),
                "n_hidden": ("INT", {"default": 16, "min": 4, "max": 128, "step": 1}),
                "train_steps": ("INT", {"default": 200, "min": 10, "max": 5000, "step": 1}),
                "seed": ("INT", {"default": 42, "min": 0, "max": 999999, "step": 1}),
            },
            "optional": {
                "training_mask": ("IMAGE",),
            },
        }

    OUTPUTS = (
        ('IMAGE', 'mask'),
        ('DATA_FIELD', 'probability'),
    )
    FUNCTION = "process"

    DESCRIPTION = (
        "Classify surface pixels into two classes using a simple two-layer "
        "feedforward neural network with sigmoid activations. Features are "
        "extracted via multi-scale Gaussian filtering. When a training mask "
        "is provided the network learns from labelled pixels; otherwise it "
        "uses unsupervised self-labelling from the initial random projection. "
        "Equivalent in purpose to Gwyddion's neural.c classifier."
    )

    KEYWORDS = ("machine learning", "ml", "segment", "nn", "feedforward", "classifier")

    def process(
        self,
        field: DataField,
        n_gaussians: int,
        n_hidden: int,
        train_steps: int,
        seed: int,
        training_mask: np.ndarray | None = None,
    ) -> tuple:
        data = np.asarray(field.data, dtype=np.float64)
        yres, xres = data.shape
        n_features = int(n_gaussians)
        n_hidden = int(n_hidden)
        train_steps = int(train_steps)

        # 1. Feature extraction
        X_all = _extract_features(data, n_features)

        # 2. Initialise weights
        rng = np.random.default_rng(int(seed))
        scale1 = np.sqrt(2.0 / n_features)
        W1 = rng.standard_normal((n_features, n_hidden)) * scale1
        b1 = np.zeros(n_hidden)
        scale2 = np.sqrt(2.0 / n_hidden)
        W2 = rng.standard_normal((n_hidden, 1)) * scale2
        b2 = np.zeros(1)

        # 3/4. Training
        if training_mask is not None:
            # Supervised — use labelled pixels
            mask_bool = mask_to_bool(training_mask)
            if mask_bool.shape != data.shape:
                raise ValueError(
                    f"Training mask shape {mask_bool.shape} does not match "
                    f"field shape {data.shape}."
                )
            # Class B = masked (255), class A = unmasked but we need both labels.
            # Pixels that are 0 are class A, pixels that are 255 are class B.
            # We train on ALL pixels that have a definitive label.
            labels_flat = training_mask.ravel().astype(np.float64) / 255.0
            # Use all pixels as training data (0 = class A, 1 = class B)
            X_train = X_all
            targets = labels_flat
            W1, b1, W2, b2 = _train_network(
                X_train, targets, W1, b1, W2, b2, train_steps,
            )
        else:
            # Unsupervised — use random projection to create initial labels,
            # then refine with self-training.
            _, y_init = _forward(X_all, W1, b1, W2, b2)
            self_labels = (y_init.ravel() > 0.5).astype(np.float64)
            # Train on the self-assigned labels for a few iterations
            steps = min(train_steps, 50)
            W1, b1, W2, b2 = _train_network(
                X_all, self_labels, W1, b1, W2, b2, steps,
            )

        # 5. Apply trained network to all pixels
        _, prob_flat = _forward(X_all, W1, b1, W2, b2)
        probability = prob_flat.reshape(yres, xres)

        # 6. Build outputs
        mask = bool_to_mask(probability > 0.5)

        prob_field = DataField(
            data=probability,
            xreal=field.xreal,
            yreal=field.yreal,
            xoff=field.xoff,
            yoff=field.yoff,
            si_unit_xy=field.si_unit_xy,
            si_unit_z="",
            domain="spatial",
        )

        return (mask, prob_field)