Status: Needs Review

This page has not been reviewed for accuracy and completeness. Content may be outdated or contain errors.

---

Visualization Nodes

Visualization nodes cover mask overlays, tracking overlays, false-RGB monitoring, and pipeline comparison artifacts.

Anomaly And Tracking Visualization

anomaly_visualization

Anomaly detection visualization sink nodes for monitoring training progress.

ImageArtifactVizBase

ImageArtifactVizBase(
    max_samples=4,
    log_every_n_batches=1,
    execution_stages=None,
    **kwargs,
)

Bases: Node

Base class for visualization nodes that produce image artifacts.

Source code in cuvis_ai/node/anomaly_visualization.py

def __init__(
    self,
    max_samples: int = 4,
    log_every_n_batches: int = 1,
    execution_stages: set[ExecutionStage] | None = None,
    **kwargs,
) -> None:
    self.max_samples = max_samples
    self.log_every_n_batches = log_every_n_batches
    self._batch_counter = 0
    if execution_stages is None:
        execution_stages = {ExecutionStage.TRAIN, ExecutionStage.VAL, ExecutionStage.TEST}
    super().__init__(
        execution_stages=execution_stages,
        **kwargs,
    )

AnomalyMask

AnomalyMask(channel, up_to=None, **kwargs)

Bases: Node

Visualize anomaly detection with GT and predicted masks.

Creates side-by-side visualizations showing ground truth masks, predicted masks, and overlay comparisons on hyperspectral cube images. The overlay shows:

• Green: True Positives (correct anomaly detection)
• Red: False Positives (false alarms)
• Yellow: False Negatives (missed anomalies)

Also displays IoU and other metrics. Returns a list of Artifact objects for logging to monitoring systems.

Executes during validation and inference stages.

Parameters:

Name	Type	Description	Default
channel	int	Channel index to use for cube visualization (required)	required
up_to	int	Maximum number of images to visualize. If None, visualizes all (default: None)	None

Examples:

>>> decider = BinaryDecider(threshold=0.2)
>>> viz_mask = AnomalyMask(channel=30, up_to=5)
>>> tensorboard_node = TensorBoardMonitorNode(output_dir="./runs")
>>> graph.connect(
...     (logit_head.logits, decider.data),
...     (decider.decisions, viz_mask.decisions),
...     (data_node.mask, viz_mask.mask),
...     (data_node.cube, viz_mask.cube),
...     (viz_mask.artifacts, tensorboard_node.artifacts),
... )

Source code in cuvis_ai/node/anomaly_visualization.py

def __init__(self, channel: int, up_to: int | None = None, **kwargs) -> None:
    self.channel = channel
    self.up_to = up_to

    super().__init__(
        execution_stages={ExecutionStage.VAL, ExecutionStage.TEST, ExecutionStage.INFERENCE},
        channel=channel,
        up_to=up_to,
        **kwargs,
    )

forward

forward(decisions, cube, context, mask=None, scores=None)

Create anomaly mask visualizations with GT/pred comparison.

Parameters:

Name	Type	Description	Default
decisions	Tensor	Binary anomaly decisions [B, H, W, 1]	required
mask	Tensor | None	Ground truth anomaly mask [B, H, W, 1] (optional)	None
cube	Tensor	Original cube [B, H, W, C] for visualization	required
context	Context	Execution context with stage, epoch, batch_idx	required

Returns:

Type	Description
dict	Dictionary with "artifacts" key containing list of Artifact objects

Source code in cuvis_ai/node/anomaly_visualization.py

def forward(
    self,
    decisions: torch.Tensor,
    cube: torch.Tensor,
    context: Context,
    mask: torch.Tensor | None = None,
    scores: torch.Tensor | None = None,
) -> dict:
    """Create anomaly mask visualizations with GT/pred comparison.

    Parameters
    ----------
    decisions : torch.Tensor
        Binary anomaly decisions [B, H, W, 1]
    mask : torch.Tensor | None
        Ground truth anomaly mask [B, H, W, 1] (optional)
    cube : torch.Tensor
        Original cube [B, H, W, C] for visualization
    context : Context
        Execution context with stage, epoch, batch_idx

    Returns
    -------
    dict
        Dictionary with "artifacts" key containing list of Artifact objects
    """
    # Extract context information
    stage = context.stage.value
    epoch = context.epoch
    batch_idx = context.batch_idx

    # Use decisions directly (already binary)
    pred_mask = decisions.float()

    # Convert to numpy and squeeze channel dimension
    pred_mask_np = pred_mask.detach().cpu().numpy().squeeze(-1)  # [B, H, W]
    cube_np = cube.detach().cpu().numpy()  # [B, H, W, C]

    # Determine if we should use ground truth
    # Skip GT comparison if: mask not provided, inference stage, or mask is all zeros
    use_gt = (
        mask is not None and context.stage != ExecutionStage.INFERENCE and mask.any().item()
    )

    # Process ground truth mask if available
    gt_mask_np = None
    batch_iou = None
    if use_gt:
        gt_mask_np = mask.detach().cpu().numpy().squeeze(-1)  # [B, H, W]

        # Add binary mask assertion
        unique_values = np.unique(gt_mask_np)
        if not np.all(np.isin(unique_values, [0, 1, True, False])):
            raise ValueError(
                f"AnomalyMask expects binary masks with only values {{0, 1}}. "
                f"Found unique values: {unique_values}. "
                f"Ensure LentilsAnomolyDataNode is configured with anomaly_class_ids "
                f"to convert multi-class masks to binary."
            )

        # Compute batch-level IoU (matches AnomalyDetectionMetrics computation)
        batch_gt = gt_mask_np > 0.5  # [B, H, W] bool
        batch_pred = pred_mask_np > 0.5  # [B, H, W] bool
        batch_tp = np.logical_and(batch_pred, batch_gt).sum()
        batch_fp = np.logical_and(batch_pred, ~batch_gt).sum()
        batch_fn = np.logical_and(~batch_pred, batch_gt).sum()
        batch_iou = batch_tp / (batch_tp + batch_fp + batch_fn + 1e-8)

    # Determine how many images to visualize from this batch
    batch_size = pred_mask_np.shape[0]
    up_to_batch = batch_size if self.up_to is None else min(batch_size, self.up_to)

    # List to collect artifacts
    artifacts = []

    # Loop through each image in the batch up to the limit
    for i in range(up_to_batch):
        # Get predicted mask for this image
        pred = pred_mask_np[i] > 0.5  # [H, W] bool

        # Get cube channel for visualization
        cube_img = cube_np[i]  # [H, W, C]
        cube_channel = cube_img[:, :, self.channel]

        # Normalize cube channel to [0, 1] for display
        cube_norm = (cube_channel - cube_channel.min()) / (
            cube_channel.max() - cube_channel.min() + 1e-8
        )

        if use_gt:
            # Mode A: Full comparison with ground truth
            assert gt_mask_np is not None, "gt_mask_np should not be None when use_gt is True"
            gt = gt_mask_np[i] > 0.5  # [H, W] bool

            # Compute confusion matrix
            tp = np.logical_and(pred, gt)  # True Positives
            fp = np.logical_and(pred, ~gt)  # False Positives
            fn = np.logical_and(~pred, gt)  # False Negatives
            # Compute metrics
            tp_count = tp.sum()
            fp_count = fp.sum()
            fn_count = fn.sum()

            precision = tp_count / (tp_count + fp_count + 1e-8)
            recall = tp_count / (tp_count + fn_count + 1e-8)
            iou = tp_count / (tp_count + fp_count + fn_count + 1e-8)

            # Create figure with 3 subplots
            fig, axes = plt.subplots(1, 3, figsize=(18, 6))

            # Subplot 1: Ground truth mask
            axes[0].imshow(gt, cmap="gray", aspect="auto")
            axes[0].set_title("Ground Truth Mask")
            axes[0].set_xlabel("Width")
            axes[0].set_ylabel("Height")

            # Subplot 2: Cube with TP/FP/FN overlay
            per_image_ap = None
            if scores is not None:
                raw_scores = scores[i, ..., 0]
                probs = torch.sigmoid(raw_scores).flatten()
                target_tensor = mask[i, ..., 0].flatten().to(dtype=torch.long)
                if probs.numel() == target_tensor.numel():
                    per_image_ap = binary_average_precision(probs, target_tensor).item()

            axes[1].imshow(cube_norm, cmap="gray", aspect="auto")

            # Create color overlay
            overlay = np.zeros((*gt.shape, 4))
            overlay[tp] = [0, 1, 0, 0.6]  # Green: True Positives
            overlay[fp] = [1, 0, 0, 0.6]  # Red: False Positives
            overlay[fn] = [1, 1, 0, 0.6]  # Yellow: False Negatives
            # TN pixels remain transparent (no overlay)

            overlay_title = f"Overlay (Channel {self.channel}) - IoU: {iou:.3f}"
            if per_image_ap is not None:
                overlay_title += f" | AP: {per_image_ap:.3f}"
            overlay_title += "\nGreen=TP, Red=FP, Yellow=FN"

            axes[1].imshow(overlay, aspect="auto")
            axes[1].set_title(overlay_title)
            axes[1].set_xlabel("Width")
            axes[1].set_ylabel("Height")

            # Subplot 3: Predicted mask with metrics in title
            axes[2].imshow(pred, cmap="gray", aspect="auto")

            # Add metrics as title (smaller font)
            metrics_title = (
                f"Predicted Mask\nIoU: {iou:.4f} | Prec: {precision:.4f} | Rec: {recall:.4f}"
            )
            if per_image_ap is not None:
                metrics_title += f" | AP: {per_image_ap:.4f}"
            metrics_title += f"\nBatch IoU: {batch_iou:.4f} (all {batch_size} imgs) | Ch: {self.channel}/{cube_img.shape[2]}"
            axes[2].set_title(metrics_title, fontsize=9)
            axes[2].set_xlabel("Width")
            axes[2].set_ylabel("Height")

            log_msg = f"Created anomaly mask artifact ({i + 1}/{up_to_batch}): IoU: {iou:.3f}"
        else:
            # Mode B: Prediction-only visualization (no ground truth)
            # Create figure with 2 subplots
            fig, axes = plt.subplots(1, 2, figsize=(12, 6))

            # Subplot 1: Cube with predicted overlay
            axes[0].imshow(cube_norm, cmap="gray", aspect="auto")

            # Create prediction overlay (cyan for predicted anomalies)
            overlay = np.zeros((*pred.shape, 4))
            overlay[pred] = [0, 1, 1, 0.6]  # Cyan: Predicted anomalies

            axes[0].imshow(overlay, aspect="auto")
            axes[0].set_title(
                f"Prediction Overlay (Channel {self.channel})\nCyan=Predicted Anomalies"
            )
            axes[0].set_xlabel("Width")
            axes[0].set_ylabel("Height")

            # Subplot 2: Predicted mask
            axes[1].imshow(pred, cmap="gray", aspect="auto")
            axes[1].set_title("Predicted Mask")
            axes[1].set_xlabel("Width")
            axes[1].set_ylabel("Height")

            # Add statistics as text
            pred_pixels = pred.sum()
            total_pixels = pred.size
            pred_ratio = pred_pixels / total_pixels

            stats_text = (
                f"Prediction Stats:\n"
                f"Anomaly pixels: {pred_pixels}\n"
                f"Total pixels: {total_pixels}\n"
                f"Anomaly ratio: {pred_ratio:.4f}\n"
                f"\n"
                f"Channel: {self.channel}/{cube_img.shape[2]}\n"
                f"\n"
                f"Mode: Inference/No GT"
            )

            fig.text(
                0.98,
                0.5,
                stats_text,
                ha="left",
                va="center",
                bbox={
                    "boxstyle": "round",
                    "facecolor": "lightblue",
                    "alpha": 0.5,
                },
                fontfamily="monospace",
            )

            log_msg = (
                f"Created anomaly mask artifact ({i + 1}/{up_to_batch}): prediction-only mode"
            )

        # Add main title with epoch/batch info
        fig.suptitle(
            f"Anomaly Mask Visualization - {stage} E{epoch} B{batch_idx} Img{i}",
            fontsize=14,
            fontweight="bold",
        )

        plt.tight_layout()

        # Convert figure to numpy array (RGB format)
        img_array = fig_to_array(fig, dpi=150)

        # Create Artifact object
        artifact = Artifact(
            name=f"anomaly_mask_img{i:02d}",
            stage=context.stage,
            epoch=context.epoch,
            batch_idx=context.batch_idx,
            value=img_array,
            el_id=i,
            desc=f"Anomaly mask for {stage} epoch {epoch}, batch {batch_idx}, image {i}",
            type=ArtifactType.IMAGE,
        )
        artifacts.append(artifact)

        logger.info(log_msg)

        plt.close(fig)

    # Return artifacts
    return {"artifacts": artifacts}

ScoreHeatmapVisualizer

ScoreHeatmapVisualizer(
    normalize_scores=True, cmap="inferno", up_to=5, **kwargs
)

Bases: Node

Log LAD/RX score heatmaps as TensorBoard artifacts.

Source code in cuvis_ai/node/anomaly_visualization.py

def __init__(
    self,
    normalize_scores: bool = True,
    cmap: str = "inferno",
    up_to: int | None = 5,
    **kwargs,
) -> None:
    self.normalize_scores = normalize_scores
    self.cmap = cmap
    self.up_to = up_to
    super().__init__(
        execution_stages={ExecutionStage.VAL, ExecutionStage.TEST, ExecutionStage.INFERENCE},
        normalize_scores=normalize_scores,
        cmap=cmap,
        up_to=up_to,
        **kwargs,
    )

forward

forward(scores, context)

Generate heatmap visualizations of anomaly scores.

Creates color-mapped heatmaps of anomaly scores for visualization in TensorBoard. Optionally normalizes scores to [0, 1] range for consistent visualization across batches.

Parameters:

Name	Type	Description	Default
scores	Tensor	Anomaly scores [B, H, W, 1] from detection nodes (e.g., RX, LAD).	required
context	Context	Execution context with stage, epoch, batch_idx information.	required

Returns:

Type	Description
dict[str, list[Artifact]]	Dictionary with "artifacts" key containing list of heatmap artifacts.

Source code in cuvis_ai/node/anomaly_visualization.py

def forward(self, scores: torch.Tensor, context: Context) -> dict[str, list[Artifact]]:
    """Generate heatmap visualizations of anomaly scores.

    Creates color-mapped heatmaps of anomaly scores for visualization
    in TensorBoard. Optionally normalizes scores to [0, 1] range for
    consistent visualization across batches.

    Parameters
    ----------
    scores : Tensor
        Anomaly scores [B, H, W, 1] from detection nodes (e.g., RX, LAD).
    context : Context
        Execution context with stage, epoch, batch_idx information.

    Returns
    -------
    dict[str, list[Artifact]]
        Dictionary with "artifacts" key containing list of heatmap artifacts.
    """
    artifacts: list[Artifact] = []
    batch_limit = scores.shape[0] if self.up_to is None else min(scores.shape[0], self.up_to)

    for idx in range(batch_limit):
        score_map = scores[idx, ..., 0].detach().cpu().numpy()

        if self.normalize_scores:
            min_v = float(score_map.min())
            max_v = float(score_map.max())
            if max_v - min_v > 1e-9:
                score_map = (score_map - min_v) / (max_v - min_v)
            else:
                score_map = np.zeros_like(score_map)

        fig, ax = plt.subplots(1, 1, figsize=(4, 4))
        im = ax.imshow(score_map, cmap=self.cmap)
        ax.set_title(f"Score Heatmap #{idx}")
        ax.axis("off")
        fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
        img_array = fig_to_array(fig, dpi=150)
        plt.close(fig)

        artifact = Artifact(
            name=f"score_heatmap_img{idx:02d}",
            value=img_array,
            el_id=idx,
            desc="Anomaly score heatmap",
            type=ArtifactType.IMAGE,
            stage=context.stage,
            epoch=context.epoch,
            batch_idx=context.batch_idx,
        )
        artifacts.append(artifact)

    return {"artifacts": artifacts}

RGBAnomalyMask

RGBAnomalyMask(up_to=None, **kwargs)

Bases: Node

Visualize anomaly detection with GT and predicted masks on RGB images.

Similar to AnomalyMask but designed for RGB images (e.g., from band selectors). Creates side-by-side visualizations showing ground truth masks, predicted masks, and overlay comparisons on RGB images. The overlay shows:

• Green: True Positives (correct anomaly detection)
• Red: False Positives (false alarms)
• Yellow: False Negatives (missed anomalies)

Also displays IoU and other metrics. Returns a list of Artifact objects for logging to monitoring systems.

Executes during validation and inference stages.

Parameters:

Name	Type	Description	Default
up_to	int	Maximum number of images to visualize. If None, visualizes all (default: None)	None

Examples:

>>> decider = BinaryDecider(threshold=0.2)
>>> viz_mask = RGBAnomalyMask(up_to=5)
>>> tensorboard_node = TensorBoardMonitorNode(output_dir="./runs")
>>> graph.connect(
...     (decider.decisions, viz_mask.decisions),
...     (data_node.mask, viz_mask.mask),
...     (band_selector.rgb_image, viz_mask.rgb_image),
...     (viz_mask.artifacts, tensorboard_node.artifacts),
... )

Initialize RGBAnomalyMask visualizer.

Parameters:

Name	Type	Description	Default
up_to	int | None	Maximum number of images to visualize. If None, visualizes all (default: None)	None

Source code in cuvis_ai/node/anomaly_visualization.py

def __init__(self, up_to: int | None = None, **kwargs) -> None:
    """Initialize RGBAnomalyMask visualizer.

    Parameters
    ----------
    up_to : int | None, optional
        Maximum number of images to visualize. If None, visualizes all (default: None)
    """
    self.up_to = up_to
    super().__init__(
        execution_stages={ExecutionStage.VAL, ExecutionStage.TEST, ExecutionStage.INFERENCE},
        up_to=up_to,
        **kwargs,
    )

forward

forward(
    decisions,
    rgb_image,
    mask=None,
    context=None,
    scores=None,
)

Create anomaly mask visualizations with GT/pred comparison on RGB images.

Parameters:

Name	Type	Description	Default
decisions	Tensor	Binary anomaly decisions [B, H, W, 1]	required
rgb_image	Tensor	RGB image [B, H, W, 3] for visualization	required
mask	Tensor | None	Ground truth anomaly mask [B, H, W, 1] (optional)	None
context	Context | None	Execution context with stage, epoch, batch_idx	None
scores	Tensor | None	Optional anomaly logits/scores [B, H, W, 1]	None

Returns:

Type	Description
dict	Dictionary with "artifacts" key containing list of Artifact objects

Source code in cuvis_ai/node/anomaly_visualization.py

def forward(
    self,
    decisions: torch.Tensor,
    rgb_image: torch.Tensor,
    mask: torch.Tensor | None = None,
    context: Context | None = None,
    scores: torch.Tensor | None = None,
) -> dict:
    """Create anomaly mask visualizations with GT/pred comparison on RGB images.

    Parameters
    ----------
    decisions : torch.Tensor
        Binary anomaly decisions [B, H, W, 1]
    rgb_image : torch.Tensor
        RGB image [B, H, W, 3] for visualization
    mask : torch.Tensor | None
        Ground truth anomaly mask [B, H, W, 1] (optional)
    context : Context | None
        Execution context with stage, epoch, batch_idx
    scores : torch.Tensor | None
        Optional anomaly logits/scores [B, H, W, 1]

    Returns
    -------
    dict
        Dictionary with "artifacts" key containing list of Artifact objects
    """
    if context is None:
        raise ValueError("RGBAnomalyMask.forward() requires a Context object")

    # Convert to numpy only at this point (keep on device until last moment)
    pred_mask_np: np.ndarray = tensor_to_numpy(decisions.float().squeeze(-1))  # [B, H, W]
    rgb_np: np.ndarray = tensor_to_numpy(rgb_image)  # [B, H, W, 3]

    # Normalize RGB to [0, 1]
    if rgb_np.max() > 1.0:
        rgb_np = rgb_np / 255.0
    rgb_np = np.clip(rgb_np, 0.0, 1.0)

    # Check if GT available and valid
    use_gt = (
        mask is not None and context.stage != ExecutionStage.INFERENCE and mask.any().item()
    )

    # Validate and convert GT if available
    gt_mask_np: np.ndarray | None = None
    batch_iou: float | None = None
    if use_gt:
        assert mask is not None
        gt_mask_np = tensor_to_numpy(mask.squeeze(-1))  # [B, H, W]
        unique_values = np.unique(gt_mask_np)
        if not np.all(np.isin(unique_values, [0, 1, True, False])):
            raise ValueError(f"RGBAnomalyMask expects binary masks, found: {unique_values}")
        # Compute batch IoU
        batch_pred = pred_mask_np > 0.5
        batch_gt = gt_mask_np > 0.5
        tp = np.logical_and(batch_pred, batch_gt).sum()
        fp = np.logical_and(batch_pred, ~batch_gt).sum()
        fn = np.logical_and(~batch_pred, batch_gt).sum()
        batch_iou = float(tp / (tp + fp + fn + 1e-8))

    batch_size = pred_mask_np.shape[0]
    up_to_batch = min(batch_size, self.up_to or batch_size)
    artifacts = []

    # Loop through images and visualize
    for i in range(up_to_batch):
        pred = pred_mask_np[i] > 0.5
        rgb_img = rgb_np[i]
        gt = gt_mask_np[i] > 0.5 if gt_mask_np is not None else None

        # Compute metrics and AP if GT available
        metrics: dict | None = None
        per_image_ap: float | None = None
        if gt is not None:
            metrics = self._compute_metrics(pred, gt)
            if scores is not None and mask is not None:
                raw_scores = scores[i, ..., 0]
                probs = torch.sigmoid(raw_scores).flatten()
                target_tensor = mask[i, ..., 0].flatten().to(dtype=torch.long)
                if probs.numel() == target_tensor.numel():
                    per_image_ap = binary_average_precision(probs, target_tensor).item()

        # Create figure and plot
        ncols = 3 if gt is not None else 2
        fig, axes = plt.subplots(1, ncols, figsize=(6 * ncols, 6))
        if ncols == 1:
            axes = [axes]

        if gt is not None and metrics is not None and batch_iou is not None:
            self._plot_with_gt(
                axes, rgb_img, pred, gt, metrics, batch_iou, batch_size, per_image_ap
            )
            log_msg = (
                f"Created RGB anomaly mask ({i + 1}/{up_to_batch}): IoU={metrics['iou']:.3f}"
            )
        else:
            self._plot_no_gt(axes, rgb_img, pred)
            log_msg = f"Created RGB anomaly mask ({i + 1}/{up_to_batch}) (no GT)"

        plt.tight_layout()
        img_array = fig_to_array(fig, dpi=150)
        plt.close(fig)

        artifact = Artifact(
            name=f"rgb_anomaly_mask_img{i:02d}",
            value=img_array,
            el_id=i,
            desc=log_msg,
            type=ArtifactType.IMAGE,
            stage=context.stage,
            epoch=context.epoch,
            batch_idx=context.batch_idx,
        )
        artifacts.append(artifact)

    return {"artifacts": artifacts}

ChannelSelectorFalseRGBViz

ChannelSelectorFalseRGBViz(
    mask_overlay_alpha=0.4,
    max_samples=4,
    log_every_n_batches=1,
    **kwargs,
)

Bases: ImageArtifactVizBase

Visualize false RGB output from channel selectors with optional mask overlay.

Produces per-sample image artifacts:

• false_rgb_sample_{b}: Normalized false RGB image [H, W, 3]
• mask_overlay_sample_{b}: False RGB with red alpha-blend on foreground pixels (if mask provided)

Parameters:

Name	Type	Description	Default
mask_overlay_alpha	float	Alpha value for red mask overlay on foreground pixels (default: 0.4).	0.4
max_samples	int	Maximum number of batch elements to visualize (default: 4).	4
log_every_n_batches	int	Log every N-th batch (default: 1).	1

Source code in cuvis_ai/node/anomaly_visualization.py

def __init__(
    self,
    mask_overlay_alpha: float = 0.4,
    max_samples: int = 4,
    log_every_n_batches: int = 1,
    **kwargs,
) -> None:
    self.mask_overlay_alpha = mask_overlay_alpha
    super().__init__(
        max_samples=max_samples,
        log_every_n_batches=log_every_n_batches,
        mask_overlay_alpha=mask_overlay_alpha,
        **kwargs,
    )

forward

forward(rgb_output, context, mask=None, mesu_index=None)

Generate false RGB and mask overlay artifacts.

Parameters:

Name	Type	Description	Default
rgb_output	Tensor	False RGB tensor [B, H, W, 3].	required
context	Context	Execution context with stage, epoch, batch_idx.	required
mask	Tensor | None	Optional segmentation mask [B, H, W].	None
mesu_index	Tensor | None	Optional measurement indices [B] for frame-identified artifact naming.	None

Returns:

Type	Description
dict[str, list[Artifact]]	Dictionary with "artifacts" key containing image artifacts.

Source code in cuvis_ai/node/anomaly_visualization.py

def forward(
    self,
    rgb_output: torch.Tensor,
    context: Context,
    mask: torch.Tensor | None = None,
    mesu_index: torch.Tensor | None = None,
) -> dict[str, list[Artifact]]:
    """Generate false RGB and mask overlay artifacts.

    Parameters
    ----------
    rgb_output : torch.Tensor
        False RGB tensor [B, H, W, 3].
    context : Context
        Execution context with stage, epoch, batch_idx.
    mask : torch.Tensor | None
        Optional segmentation mask [B, H, W].
    mesu_index : torch.Tensor | None
        Optional measurement indices [B] for frame-identified artifact naming.

    Returns
    -------
    dict[str, list[Artifact]]
        Dictionary with "artifacts" key containing image artifacts.
    """
    if not self._should_log():
        return {"artifacts": []}

    batch_size = min(rgb_output.shape[0], self.max_samples)
    artifacts = []

    for b in range(batch_size):
        # Use mesu_index for naming if available, otherwise fall back to batch index
        frame_id = f"mesu_{mesu_index[b].item()}" if mesu_index is not None else f"sample_{b}"

        # Normalized false RGB
        rgb_np = self._normalize_image(rgb_output[b].detach().cpu().numpy())

        artifact_rgb = Artifact(
            name=f"false_rgb_{frame_id}",
            value=rgb_np,
            el_id=b,
            desc=f"False RGB for {frame_id}",
            type=ArtifactType.IMAGE,
            stage=context.stage,
            epoch=context.epoch,
            batch_idx=context.batch_idx,
        )
        artifacts.append(artifact_rgb)

        # Mask overlay (if mask provided and has foreground)
        if mask is not None and mask[b].any():
            overlay = create_mask_overlay(
                torch.from_numpy(rgb_np), mask[b].cpu(), alpha=self.mask_overlay_alpha
            ).numpy()
            artifact_overlay = Artifact(
                name=f"mask_overlay_{frame_id}",
                value=overlay,
                el_id=b,
                desc=f"False RGB with mask overlay for {frame_id}",
                type=ArtifactType.IMAGE,
                stage=context.stage,
                epoch=context.epoch,
                batch_idx=context.batch_idx,
            )
            artifacts.append(artifact_overlay)

    return {"artifacts": artifacts}

MaskOverlayNode

MaskOverlayNode(
    alpha=0.4, overlay_color=(1.0, 0.0, 0.0), **kwargs
)

Bases: Node

Alpha-blend a coloured mask overlay onto RGB frames.

Pure PyTorch processing node (no matplotlib, no gradients). When mask is None or entirely zero the input RGB is passed through unchanged.

Parameters:

Name	Type	Description	Default
alpha	float	Blend factor for the overlay colour (default: 0.4).	0.4
overlay_color	tuple[float, float, float]	RGB overlay colour in [0, 1] (default: red (1, 0, 0)).	(1.0, 0.0, 0.0)

Source code in cuvis_ai/node/anomaly_visualization.py

def __init__(
    self,
    alpha: float = 0.4,
    overlay_color: Sequence[float] = (1.0, 0.0, 0.0),
    **kwargs,
) -> None:
    if len(overlay_color) != 3:
        raise ValueError(
            f"overlay_color must contain exactly 3 channels (R, G, B), got {overlay_color}"
        )

    parsed_overlay_color = tuple(float(channel) for channel in overlay_color)
    if any(channel < 0.0 or channel > 1.0 for channel in parsed_overlay_color):
        raise ValueError(
            f"overlay_color channels must be within [0, 1], got {parsed_overlay_color}"
        )

    self.overlay_color = parsed_overlay_color
    self.alpha = alpha
    super().__init__(alpha=alpha, overlay_color=self.overlay_color, **kwargs)

forward

forward(rgb_image, mask=None, **_)

Apply mask overlay to RGB frames.

Source code in cuvis_ai/node/anomaly_visualization.py

@torch.no_grad()
def forward(
    self,
    rgb_image: torch.Tensor,
    mask: torch.Tensor | None = None,
    **_,
) -> dict[str, torch.Tensor]:
    """Apply mask overlay to RGB frames."""
    if mask is None or not mask.any():
        return {"rgb_with_overlay": rgb_image}
    return {
        "rgb_with_overlay": create_mask_overlay(
            rgb_image, mask, alpha=self.alpha, color=self.overlay_color
        )
    }

TrackingOverlayNode

TrackingOverlayNode(
    alpha=0.4, draw_contours=True, draw_ids=True, **kwargs
)

Bases: Node

Alpha-blend per-object coloured masks onto RGB frames.

Converts a SAM3-style label map (mask) into per-object binary masks and renders a coloured overlay with optional contour lines and object-ID labels using :func:cuvis_ai.utils.torch_draw.overlay_instances.

Parameters:

Name	Type	Description	Default
alpha	float	Blend factor for the overlay colour (default 0.4).	0.4
draw_contours	bool	Draw contour outlines on mask edges (default True).	True
draw_ids	bool	Render numeric object-ID labels above each mask (default True).	True

Source code in cuvis_ai/node/anomaly_visualization.py

def __init__(
    self,
    alpha: float = 0.4,
    draw_contours: bool = True,
    draw_ids: bool = True,
    **kwargs,
) -> None:
    self.alpha = float(alpha)
    self.draw_contours = bool(draw_contours)
    self.draw_ids = bool(draw_ids)
    super().__init__(
        alpha=alpha,
        draw_contours=draw_contours,
        draw_ids=draw_ids,
        **kwargs,
    )

forward

forward(
    rgb_image, mask, object_ids=None, frame_id=None, **_
)

Render coloured per-object mask overlays onto rgb_image.

Parameters:

Name	Type	Description	Default
rgb_image	Tensor	Single RGB frame [1, H, W, 3] float32 in [0, 1].	required
mask	Tensor	SAM3 label map [1, H, W] int32.	required
object_ids	Tensor or None	Active object IDs [1, N] int64. When provided, only these IDs are rendered and the ordering is preserved. When absent, all non-zero unique values in mask are used.	None
frame_id	Tensor or None	Frame / measurement index [1] int64. When provided, the frame number is rendered in the top-left corner.	None

Returns:

Type	Description
dict	{"rgb_with_overlay": torch.Tensor [1, H, W, 3] float32 in [0, 1]}

Source code in cuvis_ai/node/anomaly_visualization.py

@torch.no_grad()
def forward(
    self,
    rgb_image: torch.Tensor,
    mask: torch.Tensor,
    object_ids: torch.Tensor | None = None,
    frame_id: torch.Tensor | None = None,
    **_,
) -> dict[str, torch.Tensor]:
    """Render coloured per-object mask overlays onto *rgb_image*.

    Parameters
    ----------
    rgb_image : torch.Tensor
        Single RGB frame ``[1, H, W, 3]`` float32 in ``[0, 1]``.
    mask : torch.Tensor
        SAM3 label map ``[1, H, W]`` int32.
    object_ids : torch.Tensor or None
        Active object IDs ``[1, N]`` int64.  When provided, only these IDs
        are rendered and the ordering is preserved.  When absent, all
        non-zero unique values in *mask* are used.
    frame_id : torch.Tensor or None
        Frame / measurement index ``[1]`` int64.  When provided, the frame
        number is rendered in the top-left corner.

    Returns
    -------
    dict
        ``{"rgb_with_overlay": torch.Tensor [1, H, W, 3] float32 in [0, 1]}``
    """
    frame_u8 = (rgb_image[0].clamp(0.0, 1.0) * 255.0).to(torch.uint8)  # [H, W, 3]
    # Align mask/object-id tensors to the image device. Some upstream nodes can
    # emit CPU tensors even when the visualization path runs on CUDA.
    mask_t = mask[0].to(frame_u8.device)  # [H, W] int32
    present_ids_t = torch.unique(mask_t)
    present_ids_t = present_ids_t[present_ids_t > 0]

    if object_ids is not None:
        # Some trackers include background label 0 in object_ids; never render it.
        filtered_ids = object_ids[0].to(mask_t.device)
        filtered_ids = filtered_ids[filtered_ids > 0]
        if present_ids_t.numel() > 0:
            filtered_ids = filtered_ids[torch.isin(filtered_ids, present_ids_t)]
        else:
            filtered_ids = filtered_ids[:0]
        ids = []
        seen: set[int] = set()
        for raw_id in filtered_ids.tolist():
            obj_id = int(raw_id)
            if obj_id in seen:
                continue
            seen.add(obj_id)
            ids.append(obj_id)
    else:
        ids = [int(v) for v in present_ids_t.tolist()]

    per_obj_masks: list[tuple[int, torch.Tensor]] = [(oid, mask_t == oid) for oid in ids]

    rendered = overlay_instances(
        frame_u8,
        per_obj_masks,
        alpha=self.alpha,
        draw_edges=self.draw_contours,
        draw_ids=self.draw_ids,
    )

    if frame_id is not None:
        fid = int(frame_id.reshape(-1)[0].item())
        draw_text(rendered, 8, 8, f"frame {fid}", (255, 255, 255), scale=2, bg=True)

    out = rendered.to(torch.float32) / 255.0  # [H, W, 3]
    return {"rgb_with_overlay": out.unsqueeze(0)}  # [1, H, W, 3]

TrackingPointerOverlayNode

TrackingPointerOverlayNode(
    alpha=0.4, draw_contours=True, draw_ids=True, **kwargs
)

Bases: Node

Draw downward triangle pointers for all tracked objects.

The node is composable by design: it renders only the pointer markers on top of an incoming RGB frame and does not perform any mask tinting itself. Colours are derived from object IDs using the same palette as :class:TrackingOverlayNode.

Parameters:

Name	Type	Description	Default
alpha	float	Reserved for API compatibility with :class:TrackingOverlayNode (unused).	0.4
draw_contours	bool	Reserved for API compatibility with :class:TrackingOverlayNode (unused).	True
draw_ids	bool	Reserved for API compatibility with :class:TrackingOverlayNode (unused).	True

Source code in cuvis_ai/node/anomaly_visualization.py

def __init__(
    self,
    alpha: float = 0.4,
    draw_contours: bool = True,
    draw_ids: bool = True,
    **kwargs,
) -> None:
    self.alpha = float(alpha)
    self.draw_contours = bool(draw_contours)
    self.draw_ids = bool(draw_ids)
    super().__init__(
        alpha=alpha,
        draw_contours=draw_contours,
        draw_ids=draw_ids,
        **kwargs,
    )

forward

forward(
    rgb_image, mask, object_ids=None, frame_id=None, **_
)

Render pointer overlays for all objects onto rgb_image.

Source code in cuvis_ai/node/anomaly_visualization.py

@torch.no_grad()
def forward(
    self,
    rgb_image: torch.Tensor,
    mask: torch.Tensor,
    object_ids: torch.Tensor | None = None,
    frame_id: torch.Tensor | None = None,
    **_,
) -> dict[str, torch.Tensor]:
    """Render pointer overlays for all objects onto *rgb_image*."""
    frame_u8 = (rgb_image[0].clamp(0.0, 1.0) * 255.0).to(torch.uint8)
    mask_t = mask[0].to(frame_u8.device)
    if tuple(mask_t.shape) != tuple(frame_u8.shape[:2]):
        raise ValueError(
            f"Mask shape {tuple(mask_t.shape)} does not match image shape {tuple(frame_u8.shape[:2])}."
        )

    present_ids_t = torch.unique(mask_t)
    present_ids_t = present_ids_t[present_ids_t > 0]

    if object_ids is not None:
        filtered_ids = object_ids[0].to(mask_t.device)
        filtered_ids = filtered_ids[filtered_ids > 0]
        if present_ids_t.numel() > 0:
            filtered_ids = filtered_ids[torch.isin(filtered_ids, present_ids_t)]
        else:
            filtered_ids = filtered_ids[:0]
        ids: list[int] = []
        seen: set[int] = set()
        for raw_id in filtered_ids.tolist():
            oid = int(raw_id)
            if oid in seen:
                continue
            seen.add(oid)
            ids.append(oid)
    else:
        ids = [int(v) for v in present_ids_t.tolist()]

    rendered = frame_u8.clone()

    for oid in ids:
        fg = mask_t == oid
        if not torch.any(fg):
            continue

        ys, xs = torch.where(fg)
        x_min = int(xs.min().item())
        x_max = int(xs.max().item())
        y_min = int(ys.min().item())

        bbox_width = x_max - x_min + 1
        tri_width = max(12, min(48, int(round(bbox_width * 0.45))))
        tri_width = min(tri_width, max(1, int(rendered.shape[1]) - 1))
        tri_height = max(10, min(36, int(round(tri_width * 0.8))))
        tri_height = min(tri_height, max(1, int(rendered.shape[0]) - 1))
        gap = max(4, min(10, int(round(tri_height * 0.3))))

        centroid_x = int(torch.round(xs.to(torch.float32).mean()).item())
        half_width = max(1, (tri_width + 1) // 2)
        max_tip_x = max(half_width, int(rendered.shape[1]) - 1 - half_width)
        tip_x = max(half_width, min(centroid_x, max_tip_x))
        desired_tip_y = y_min - gap
        tip_y = max(tri_height, min(desired_tip_y, int(rendered.shape[0]) - 1))

        color_t = id_to_color(torch.tensor([oid], device=rendered.device, dtype=torch.int64))[0]
        color = tuple(int(channel) for channel in color_t.tolist())

        outline_thickness = 2 if tri_width >= 20 and tri_height >= 16 else 1
        draw_downward_triangle(
            rendered,
            tip_x=tip_x,
            tip_y=tip_y,
            width=tri_width,
            height=tri_height,
            color=color,
            outline_color=(0, 0, 0),
            outline_thickness=outline_thickness,
        )

    if frame_id is not None:
        fid = int(frame_id.reshape(-1)[0].item())
        draw_text(rendered, 8, 8, f"frame {fid}", (255, 255, 255), scale=2, bg=True)

    out = rendered.to(torch.float32) / 255.0
    return {"rgb_with_overlay": out.unsqueeze(0)}

BBoxesOverlayNode

BBoxesOverlayNode(
    line_thickness=2,
    draw_labels=False,
    draw_sparklines=False,
    sparkline_height=24,
    hide_untracked=False,
    **kwargs,
)

Bases: Node

Torch-only bounding-box overlay renderer for YOLO-style detections.

Source code in cuvis_ai/node/anomaly_visualization.py

def __init__(
    self,
    line_thickness: int = 2,
    draw_labels: bool = False,
    draw_sparklines: bool = False,
    sparkline_height: int = 24,
    hide_untracked: bool = False,
    **kwargs,
) -> None:
    self.line_thickness = int(line_thickness)
    self.draw_labels = bool(draw_labels)
    self.draw_sparklines = bool(draw_sparklines)
    self.sparkline_height = int(sparkline_height)
    self.hide_untracked = bool(hide_untracked)
    super().__init__(
        line_thickness=line_thickness,
        draw_labels=draw_labels,
        draw_sparklines=draw_sparklines,
        sparkline_height=sparkline_height,
        hide_untracked=hide_untracked,
        **kwargs,
    )

forward

forward(
    rgb_image,
    bboxes,
    category_ids,
    frame_id=None,
    confidences=None,
    spectral_signatures=None,
    **_,
)

Overlay bbox edges with deterministic per-class colors.

Source code in cuvis_ai/node/anomaly_visualization.py

@torch.no_grad()
def forward(
    self,
    rgb_image: torch.Tensor,
    bboxes: torch.Tensor,
    category_ids: torch.Tensor,
    frame_id: torch.Tensor | None = None,
    confidences: torch.Tensor | None = None,  # noqa: ARG002
    spectral_signatures: torch.Tensor | None = None,
    **_,
) -> dict[str, torch.Tensor]:
    """Overlay bbox edges with deterministic per-class colors."""
    # Optionally filter out untracked detections (category_id / track_id < 0).
    if self.hide_untracked and category_ids.numel() > 0:
        mask = category_ids[0] >= 0  # [N]
        bboxes = bboxes[:, mask]
        category_ids = category_ids[:, mask]
        if spectral_signatures is not None:
            spectral_signatures = spectral_signatures[:, mask]

    sigs = spectral_signatures if self.draw_sparklines else None
    return {
        "rgb_with_overlay": render_bboxes_overlay_torch(
            rgb_image=rgb_image,
            bboxes=bboxes,
            category_ids=category_ids,
            frame_id=frame_id,
            line_thickness=self.line_thickness,
            draw_labels=self.draw_labels,
            spectral_signatures=sigs,
            sparkline_height=self.sparkline_height,
        )
    }

ChannelWeightsViz

ChannelWeightsViz(
    max_samples=1,
    log_every_n_batches=1,
    cell_height=60,
    cell_width=12,
    **kwargs,
)

Bases: ImageArtifactVizBase

Visualize channel mixer weights as a heatmap.

Produces a [K, C] mixing matrix heatmap with output channels on the y-axis and input channels on the x-axis. Uses a diverging blue-white-red colormap centred at zero so positive/negative contributions are immediately visible.

Implemented in pure PyTorch (no matplotlib) so it adds negligible overhead to the training loop.

Parameters:

Name	Type	Description	Default
max_samples	int	Ignored (weights are per-model, not per-sample). Kept for base class compatibility. Default: 1.	1
log_every_n_batches	int	Log every N-th batch (default: 1).	1
cell_height	int	Pixel height per matrix row (default: 40).	60
cell_width	int	Pixel width per matrix column (default: 6).	12

Source code in cuvis_ai/node/anomaly_visualization.py

def __init__(
    self,
    max_samples: int = 1,
    log_every_n_batches: int = 1,
    cell_height: int = 60,
    cell_width: int = 12,
    **kwargs,
) -> None:
    super().__init__(
        max_samples=max_samples,
        log_every_n_batches=log_every_n_batches,
        **kwargs,
    )
    self.cell_height = cell_height
    self.cell_width = cell_width

forward

forward(weights, context, wavelengths=None)

Generate mixing matrix heatmap artifact.

Pure-torch rendering with R/G/B indicator bars, grid lines, and a diverging colorbar — no matplotlib for training-loop speed.

Parameters:

Name	Type	Description	Default
weights	Tensor	Mixing matrix [K, C].	required
context	Context	Execution context with stage, epoch, batch_idx.	required
wavelengths	ndarray	Wavelengths [C] in nm (reserved for future use).	None

Returns:

Type	Description
dict[str, list[Artifact]]	Dictionary with "artifacts" key.

Source code in cuvis_ai/node/anomaly_visualization.py

def forward(
    self,
    weights: torch.Tensor,
    context: Context,
    wavelengths: np.ndarray | None = None,
) -> dict[str, list[Artifact]]:
    """Generate mixing matrix heatmap artifact.

    Pure-torch rendering with R/G/B indicator bars, grid lines, and a
    diverging colorbar — no matplotlib for training-loop speed.

    Parameters
    ----------
    weights : Tensor
        Mixing matrix ``[K, C]``.
    context : Context
        Execution context with stage, epoch, batch_idx.
    wavelengths : ndarray, optional
        Wavelengths ``[C]`` in nm (reserved for future use).

    Returns
    -------
    dict[str, list[Artifact]]
        Dictionary with ``"artifacts"`` key.
    """
    if not self._should_log():
        return {"artifacts": []}

    w = weights.detach().float()
    if w.ndim == 1:
        w = w.unsqueeze(0)  # [1, C]

    K, C = w.shape
    vmax = w.abs().max().clamp_min(1e-8)
    t = (w / vmax + 1.0) * 0.5  # [K, C] in [0, 1], 0.5 = zero

    # Colormap the heatmap → [K, C, 3]
    heatmap = _diverging_colormap(t)

    # Build upscaled heatmap canvas with 1px black grid lines
    ch, cw = self.cell_height, self.cell_width
    grid_h = K * ch + (K + 1)
    grid_w = C * cw + (C + 1)
    canvas = torch.zeros(grid_h, grid_w, 3)  # black grid lines
    for r in range(K):
        y0 = 1 + r * (ch + 1)
        for c in range(C):
            x0 = 1 + c * (cw + 1)
            canvas[y0 : y0 + ch, x0 : x0 + cw] = heatmap[r, c]

    # Left margin: colored R/G/B indicator bars
    indicator_w = max(ch // 3, 8)
    _channel_colors = torch.tensor(
        [
            [1.0, 0.0, 0.0],  # R
            [0.0, 0.7, 0.0],  # G
            [0.0, 0.0, 1.0],  # B
        ]
    )
    indicator = torch.ones(grid_h, indicator_w, 3)  # white background
    for r in range(min(K, len(_channel_colors))):
        y0 = 1 + r * (ch + 1)
        indicator[y0 : y0 + ch, :] = _channel_colors[r]

    # Right margin: colorbar gradient (top=positive/red, bottom=negative/blue)
    cbar_w = max(ch // 3, 8)
    cbar_t = torch.linspace(1.0, 0.0, grid_h).unsqueeze(1)  # [grid_h, 1]
    cbar = _diverging_colormap(cbar_t).expand(grid_h, cbar_w, 3).contiguous()

    # Assemble: [indicator | gap | heatmap | gap | colorbar]
    gap = torch.ones(grid_h, 2, 3)  # 2px white gap
    full = torch.cat([indicator, gap, canvas, gap, cbar], dim=1)

    # Convert to uint8 numpy [H, W, 3]
    heatmap_np = (full * 255).clamp(0, 255).byte().cpu().numpy()

    # --- Text annotations via PIL (fast, no matplotlib) ---
    vmax_val = vmax.item()
    heat_h, heat_w = heatmap_np.shape[:2]
    left_margin, top_margin = 30, 5
    bottom_margin, right_margin = 45, 50

    canvas_img = Image.new(
        "RGB",
        (left_margin + heat_w + right_margin, top_margin + heat_h + bottom_margin),
        (255, 255, 255),
    )
    canvas_img.paste(Image.fromarray(heatmap_np), (left_margin, top_margin))
    draw = ImageDraw.Draw(canvas_img)
    font = ImageFont.load_default()

    # Y-axis: R / G / B labels (centered on each row)
    row_labels = ["R", "G", "B"] if K == 3 else [str(i) for i in range(K)]
    for r in range(K):
        y_center = top_margin + 1 + r * (ch + 1) + ch // 2
        draw.text((6, y_center - 5), row_labels[r], fill=(0, 0, 0), font=font)

    # X-axis: wavelength tick labels (every Nth to avoid overlap)
    # heatmap grid starts after: indicator_w + 2px gap + 1px border
    heatmap_x0 = left_margin + indicator_w + 2 + 1
    if wavelengths is not None and len(wavelengths) == C:
        step = max(1, C // 15)
        for i in range(0, C, step):
            x = heatmap_x0 + i * (cw + 1) + cw // 2
            y = top_margin + heat_h + 2
            label = str(int(wavelengths[i]))
            draw.text((x - len(label) * 3, y), label, fill=(0, 0, 0), font=font)

    # Colorbar scale: +vmax at top, 0 at middle, -vmax at bottom
    cbar_x = left_margin + heat_w + 4
    draw.text((cbar_x, top_margin), f"+{vmax_val:.2f}", fill=(0, 0, 0), font=font)
    draw.text((cbar_x, top_margin + heat_h // 2 - 5), "0", fill=(0, 0, 0), font=font)
    draw.text((cbar_x, top_margin + heat_h - 10), f"-{vmax_val:.2f}", fill=(0, 0, 0), font=font)

    img_np = np.array(canvas_img)

    stage = context.stage.name.lower() if context.stage else "unknown"
    return {
        "artifacts": [
            Artifact(
                name="mixing_matrix",
                value=img_np,
                el_id=0,
                desc=f"Mixing matrix — {stage} epoch {context.epoch}",
                type=ArtifactType.IMAGE,
                stage=context.stage,
                epoch=context.epoch,
                batch_idx=context.batch_idx,
            )
        ],
    }

render_bboxes_overlay_torch

render_bboxes_overlay_torch(
    rgb_image,
    bboxes,
    category_ids,
    frame_id=None,
    line_thickness=2,
    draw_labels=False,
    spectral_signatures=None,
    sparkline_height=24,
)

Render bbox edges on RGB frames using pure torch drawing primitives.

Source code in cuvis_ai/node/anomaly_visualization.py

def render_bboxes_overlay_torch(
    rgb_image: torch.Tensor,
    bboxes: torch.Tensor,
    category_ids: torch.Tensor,
    frame_id: torch.Tensor | None = None,
    line_thickness: int = 2,
    draw_labels: bool = False,
    spectral_signatures: torch.Tensor | None = None,
    sparkline_height: int = 24,
) -> torch.Tensor:
    """Render bbox edges on RGB frames using pure torch drawing primitives."""
    out = (rgb_image.clamp(0.0, 1.0) * 255.0).to(torch.uint8).clone()
    frame = out[0]  # [H, W, 3]

    num_boxes = int(bboxes.shape[1])
    if num_boxes == 0:
        return out.to(torch.float32) / 255.0

    classes = category_ids[0].to(torch.int64)
    colors = id_to_color(classes)
    n = min(num_boxes, int(colors.shape[0]))
    thickness = max(1, int(line_thickness))

    for i in range(n):
        x1, y1, x2, y2 = [int(v) for v in bboxes[0, i].round().tolist()]
        draw_box(frame, (x1, y1, x2, y2), colors[i], thickness=thickness)
        if draw_labels and int(classes[i].item()) >= 0:
            label = str(int(classes[i].item()))
            draw_text(frame, x1, max(0, y1 - 16), label, colors[i], scale=2, bg=True)

    # Draw spectral sparklines inside each bbox (bottom region)
    if spectral_signatures is not None:
        sigs = spectral_signatures[0]  # [N, C]
        sh = max(4, int(sparkline_height))
        n_sigs = min(n, int(sigs.shape[0]))
        for i in range(n_sigs):
            x1, y1, x2, y2 = [int(v) for v in bboxes[0, i].round().tolist()]
            bw = x2 - x1
            if bw < 4 or (y2 - y1) < sh + 4:
                continue
            spark_y = y2 - sh - thickness  # just inside the bottom edge
            draw_sparkline(
                frame,
                x1 + thickness,
                spark_y,
                bw - 2 * thickness,
                sh,
                sigs[i],
                colors[i],
            )

    if frame_id is not None:
        try:
            fid = int(frame_id.reshape(-1)[0].item())
            draw_text(frame, 8, 8, f"frame {fid}", (255, 255, 255), scale=2, bg=True)
        except Exception as exc:
            logger.debug("Failed to draw frame label on overlay: {}", exc)

    return out.to(torch.float32) / 255.0

Pipeline Visualization

pipeline_visualization

Pipeline and data visualization sink nodes for monitoring training progress.

CubeRGBVisualizer

CubeRGBVisualizer(name=None, up_to=5)

Bases: Node

Creates false-color RGB images from hyperspectral cube using channel weights.

Selects 3 channels with highest weights for R, G, B channels and creates a false-color visualization with wavelength annotations.

Source code in cuvis_ai/node/pipeline_visualization.py

def __init__(self, name: str | None = None, up_to: int = 5) -> None:
    super().__init__(name=name, execution_stages={ExecutionStage.INFERENCE, ExecutionStage.VAL})
    self.up_to = up_to

forward

forward(cube, weights, wavelengths, context)

Generate false-color RGB visualizations from hyperspectral cube.

Selects the 3 channels with highest weights and creates RGB images with wavelength annotations. Also generates a bar chart showing channel weights with the selected channels highlighted.

Parameters:

Name	Type	Description	Default
cube	Tensor	Hyperspectral cube [B, H, W, C].	required
weights	Tensor	Channel selection weights [C] indicating importance of each channel.	required
wavelengths	Tensor	Wavelengths for each channel [C] in nanometers.	required
context	Context	Execution context with stage, epoch, batch_idx information.	required

Returns:

Type	Description
dict[str, list[Artifact]]	Dictionary with "artifacts" key containing list of visualization artifacts.

Source code in cuvis_ai/node/pipeline_visualization.py

def forward(self, cube, weights, wavelengths, context) -> dict[str, list[Artifact]]:
    """Generate false-color RGB visualizations from hyperspectral cube.

    Selects the 3 channels with highest weights and creates RGB images
    with wavelength annotations. Also generates a bar chart showing
    channel weights with the selected channels highlighted.

    Parameters
    ----------
    cube : Tensor
        Hyperspectral cube [B, H, W, C].
    weights : Tensor
        Channel selection weights [C] indicating importance of each channel.
    wavelengths : Tensor
        Wavelengths for each channel [C] in nanometers.
    context : Context
        Execution context with stage, epoch, batch_idx information.

    Returns
    -------
    dict[str, list[Artifact]]
        Dictionary with "artifacts" key containing list of visualization artifacts.
    """
    top3_indices = torch.topk(weights, k=3).indices.cpu().numpy()
    top3_wavelengths = wavelengths[top3_indices]

    batch_size = min(cube.shape[0], self.up_to)
    artifacts = []

    for b in range(batch_size):
        rgb_channels = cube[b, :, :, top3_indices].cpu().numpy()

        rgb_img = (rgb_channels - rgb_channels.min()) / (
            rgb_channels.max() - rgb_channels.min() + 1e-8
        )

        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5))

        ax1.imshow(rgb_img)
        ax1.set_title(
            f"False RGB: R={top3_wavelengths[0]:.1f}nm, "
            f"G={top3_wavelengths[1]:.1f}nm, B={top3_wavelengths[2]:.1f}nm"
        )
        ax1.axis("off")

        ax2.bar(range(len(wavelengths)), weights.detach().cpu().numpy())
        ax2.scatter(
            top3_indices,
            weights[top3_indices].detach().cpu().numpy(),
            c="red",
            s=100,
            zorder=3,
        )
        ax2.set_xlabel("Channel Index")
        ax2.set_ylabel("Weight")
        ax2.set_title("Channel Selection Weights")
        ax2.grid(True, alpha=0.3)

        for idx in top3_indices:
            ax2.annotate(
                f"{wavelengths[idx]:.0f}nm",
                xy=(idx, weights[idx].item()),
                xytext=(0, 10),
                textcoords="offset points",
                ha="center",
                fontsize=8,
            )

        plt.tight_layout()

        img_array = fig_to_array(fig, dpi=150)

        artifact = Artifact(
            name=f"viz_rgb_sample_{b}",
            value=img_array,
            el_id=b,
            desc=f"False RGB visualization for sample {b}",
            type=ArtifactType.IMAGE,
        )
        artifacts.append(artifact)
        plt.close(fig)

    return {"artifacts": artifacts}

PCAVisualization

PCAVisualization(up_to=None, **kwargs)

Bases: Node

Visualize PCA-projected data with scatter and image plots.

Creates visualizations for each batch element showing:

1. Scatter plot of H*W points in 2D PC space (using first 2 PCs)
2. Image representation of the 2D projection reshaped to [H, W, 2]

Points in scatter plot are colored by spatial position. Returns artifacts for monitoring systems.

Executes only during validation stage.

Parameters:

Name	Type	Description	Default
up_to	int	Maximum number of batch elements to visualize. If None, visualizes all (default: None)	None

Examples:

>>> pca_viz = PCAVisualization(up_to=10)
>>> tensorboard_node = TensorBoardMonitorNode(output_dir="./runs")
>>> graph.connect(
...     (pca.projected, pca_viz.data),
...     (pca_viz.artifacts, tensorboard_node.artifacts),
... )

Source code in cuvis_ai/node/pipeline_visualization.py

def __init__(self, up_to: int | None = None, **kwargs) -> None:
    self.up_to = up_to

    super().__init__(execution_stages={ExecutionStage.VAL}, up_to=up_to, **kwargs)

forward

forward(data, context)

Create PCA projection visualizations as Artifact objects.

Parameters:

Name	Type	Description	Default
data	Tensor	PCA-projected data tensor [B, H, W, C] (uses first 2 components)	required
context	Context	Execution context with stage, epoch, batch_idx	required

Returns:

Type	Description
dict	Dictionary with "artifacts" key containing list of Artifact objects

Source code in cuvis_ai/node/pipeline_visualization.py

def forward(self, data: torch.Tensor, context: Context) -> dict:
    """Create PCA projection visualizations as Artifact objects.

    Parameters
    ----------
    data : torch.Tensor
        PCA-projected data tensor [B, H, W, C] (uses first 2 components)
    context : Context
        Execution context with stage, epoch, batch_idx

    Returns
    -------
    dict
        Dictionary with "artifacts" key containing list of Artifact objects
    """
    # Convert to numpy
    data_np = data.detach().cpu().numpy()

    # Handle input shape: [B, H, W, C]
    if data_np.ndim != 4:
        raise ValueError(f"Expected 4D input [B, H, W, C], got shape: {data_np.shape}")

    B, H, W, C = data_np.shape

    if C < 2:
        raise ValueError(f"Expected at least 2 components, got {C}")

    # Extract context information
    stage = context.stage.value
    epoch = context.epoch
    batch_idx = context.batch_idx

    # Determine how many images to visualize from this batch
    up_to_batch = B if self.up_to is None else min(B, self.up_to)

    # List to collect artifacts
    artifacts = []

    # Loop through each batch element
    for i in range(up_to_batch):
        # Get projection for this batch element: [H, W, C]
        projection = data_np[i]

        # Use only first 2 components
        projection_2d = projection[:, :, :2]  # [H, W, 2]

        # Flatten spatial dimensions for scatter plot
        projection_flat = projection_2d.reshape(-1, 2)  # [H*W, 2]

        # Create spatial position colors using 2D HSV encoding
        # x-coordinate maps to Hue (0-1)
        # y-coordinate maps to Saturation (0-1)
        # Value is constant at 1.0 for brightness
        y_coords, x_coords = np.meshgrid(np.arange(H), np.arange(W), indexing="ij")

        # Normalize coordinates to [0, 1]
        x_norm = x_coords / (W - 1) if W > 1 else np.zeros_like(x_coords)
        y_norm = y_coords / (H - 1) if H > 1 else np.zeros_like(y_coords)

        # Create HSV colors: H from x, S from y, V constant
        hsv_colors = np.stack(
            [
                x_norm.flatten(),  # Hue from x-coordinate
                y_norm.flatten(),  # Saturation from y-coordinate
                np.ones(H * W),  # Value constant at 1.0
            ],
            axis=-1,
        )

        # Convert HSV to RGB for matplotlib
        from matplotlib.colors import hsv_to_rgb

        rgb_colors = hsv_to_rgb(hsv_colors)

        # Create figure with 3 subplots
        fig, axes = plt.subplots(1, 3, figsize=(20, 6))

        # Subplot 1: Scatter plot colored by 2D spatial position
        axes[0].scatter(
            projection_flat[:, 0],
            projection_flat[:, 1],
            c=rgb_colors,
            alpha=0.6,
            s=20,
        )
        axes[0].set_xlabel("PC1 (1st component)")
        axes[0].set_ylabel("PC2 (2nd component)")
        axes[0].set_title(f"PCA Scatter - {stage} E{epoch} B{batch_idx} Img{i}")
        axes[0].grid(True, alpha=0.3)

        # Subplot 2: Spatial reference image
        # Create reference image showing the spatial color coding
        spatial_reference = hsv_to_rgb(
            np.stack([x_norm, y_norm, np.ones_like(x_norm)], axis=-1)
        )
        axes[1].imshow(spatial_reference, aspect="auto")
        axes[1].set_xlabel("Width (→ Hue)")
        axes[1].set_ylabel("Height (→ Saturation)")
        axes[1].set_title("Spatial Color Reference")

        # Subplot 3: Image representation
        # Normalize each channel to [0, 1] for visualization
        pc1_norm = (projection_2d[:, :, 0] - projection_2d[:, :, 0].min()) / (
            projection_2d[:, :, 0].max() - projection_2d[:, :, 0].min() + 1e-8
        )
        pc2_norm = (projection_2d[:, :, 1] - projection_2d[:, :, 1].min()) / (
            projection_2d[:, :, 1].max() - projection_2d[:, :, 1].min() + 1e-8
        )

        # Create RGB image: PC1 in red channel, PC2 in green channel, zeros in blue
        img_rgb = np.stack([pc1_norm, pc2_norm, np.zeros_like(pc1_norm)], axis=-1)

        axes[2].imshow(img_rgb, aspect="auto")
        axes[2].set_xlabel("Width")
        axes[2].set_ylabel("Height")
        axes[2].set_title("PCA Image (R=PC1, G=PC2)")

        # Add statistics text
        pc1_min = projection_2d[:, :, 0].min()
        pc1_max = projection_2d[:, :, 0].max()
        pc2_min = projection_2d[:, :, 1].min()
        pc2_max = projection_2d[:, :, 1].max()
        stats_text = (
            f"Shape: [{H}, {W}]\n"
            f"Points: {H * W}\n"
            f"PC1 range: [{pc1_min:.3f}, {pc1_max:.3f}]\n"
            f"PC2 range: [{pc2_min:.3f}, {pc2_max:.3f}]"
        )
        fig.text(
            0.98,
            0.5,
            stats_text,
            ha="left",
            va="center",
            bbox={
                "boxstyle": "round",
                "facecolor": "wheat",
                "alpha": 0.5,
            },
        )

        plt.tight_layout()

        # Convert figure to numpy array (RGB format)
        img_array = fig_to_array(fig, dpi=150)

        # Create Artifact object
        artifact = Artifact(
            name=f"pca_projection_img{i:02d}",
            stage=context.stage,
            epoch=context.epoch,
            batch_idx=context.batch_idx,
            value=img_array,
            el_id=i,
            desc=f"PCA projection for {stage} epoch {epoch}, batch {batch_idx}, image {i}",
            type=ArtifactType.IMAGE,
        )
        artifacts.append(artifact)

        progress_total = self.up_to if self.up_to else B
        description = (
            f"Created PCA projection artifact ({i + 1}/{progress_total}): {artifact.name}"
        )
        logger.info(description)

        plt.close(fig)

    # Return artifacts
    return {"artifacts": artifacts}

PipelineComparisonVisualizer

PipelineComparisonVisualizer(
    hsi_channels=None,
    max_samples=4,
    log_every_n_batches=1,
    **kwargs,
)

Bases: Node

TensorBoard visualization node for comparing pipeline stages.

Creates image artifacts for logging to TensorBoard:

• Input HSI cube visualization (false-color RGB from selected channels)
• Mixer output (3-channel RGB-like image that downstream model sees)
• Ground truth anomaly mask
• Anomaly scores (as heatmap)

Parameters:

Name	Type	Description	Default
hsi_channels	list[int]	Channel indices to use for false-color RGB visualization of HSI input (default: [0, 20, 40] for a simple false-color representation)	None
max_samples	int	Maximum number of samples to log per batch (default: 4)	4
log_every_n_batches	int	Log images every N batches to reduce TensorBoard size (default: 1, log every batch)	1

Source code in cuvis_ai/node/pipeline_visualization.py

def __init__(
    self,
    hsi_channels: list[int] | None = None,
    max_samples: int = 4,
    log_every_n_batches: int = 1,
    **kwargs,
) -> None:
    if hsi_channels is None:
        hsi_channels = [0, 20, 40]  # Default: use channels 0, 20, 40 for false-color RGB
    self.hsi_channels = hsi_channels
    self.max_samples = max_samples
    self.log_every_n_batches = log_every_n_batches
    self._batch_counter = 0

    super().__init__(
        execution_stages={ExecutionStage.TRAIN, ExecutionStage.VAL, ExecutionStage.TEST},
        hsi_channels=hsi_channels,
        max_samples=max_samples,
        log_every_n_batches=log_every_n_batches,
        **kwargs,
    )

forward

forward(
    hsi_cube,
    mixer_output,
    ground_truth_mask,
    adaclip_scores,
    context=None,
    **_,
)

Create image artifacts for TensorBoard logging.

Parameters:

Name	Type	Description	Default
hsi_cube	Tensor	Input HSI cube [B, H, W, C]	required
mixer_output	Tensor	Mixer output (RGB-like) [B, H, W, 3]	required
ground_truth_mask	Tensor	Ground truth anomaly mask [B, H, W, 1]	required
adaclip_scores	Tensor	Anomaly scores [B, H, W, 1]	required
context	Context	Execution context with stage, epoch, batch_idx info	None

Returns:

Type	Description
dict[str, list[Artifact]]	Dictionary with "artifacts" key containing list of Artifact objects

Source code in cuvis_ai/node/pipeline_visualization.py

def forward(
    self,
    hsi_cube: Tensor,
    mixer_output: Tensor,
    ground_truth_mask: Tensor,
    adaclip_scores: Tensor,
    context: Context | None = None,
    **_: Any,
) -> dict[str, list[Artifact]]:
    """Create image artifacts for TensorBoard logging.

    Parameters
    ----------
    hsi_cube : Tensor
        Input HSI cube [B, H, W, C]
    mixer_output : Tensor
        Mixer output (RGB-like) [B, H, W, 3]
    ground_truth_mask : Tensor
        Ground truth anomaly mask [B, H, W, 1]
    adaclip_scores : Tensor
        Anomaly scores [B, H, W, 1]
    context : Context, optional
        Execution context with stage, epoch, batch_idx info

    Returns
    -------
    dict[str, list[Artifact]]
        Dictionary with "artifacts" key containing list of Artifact objects
    """
    if context is None:
        context = Context()

    # Skip logging if not the right batch interval
    self._batch_counter += 1
    if (self._batch_counter - 1) % self.log_every_n_batches != 0:
        return {"artifacts": []}

    artifacts = []
    B = hsi_cube.shape[0]
    num_samples = min(B, self.max_samples)

    # Convert tensors to numpy for visualization
    hsi_np = hsi_cube.detach().cpu().numpy()
    mixer_np = mixer_output.detach().cpu().numpy()
    mask_np = ground_truth_mask.detach().cpu().numpy()
    scores_np = adaclip_scores.detach().cpu().numpy()

    for b in range(num_samples):
        # 1. HSI Input Visualization (false-color RGB)
        hsi_img = self._create_hsi_visualization(hsi_np[b])
        artifact = Artifact(
            name=f"hsi_input_sample_{b}",
            value=hsi_img,
            el_id=b,
            desc=f"HSI input (false-color RGB) for sample {b}",
            type=ArtifactType.IMAGE,
            stage=context.stage,
            epoch=context.epoch,
            batch_idx=context.batch_idx,
        )
        artifacts.append(artifact)

        # 2. Mixer Output (what downstream model sees as input)
        mixer_img = self._normalize_image(mixer_np[b])  # Already [H, W, 3]
        artifact = Artifact(
            name=f"mixer_output_adaclip_input_sample_{b}",
            value=mixer_img,
            el_id=b,
            desc=f"Mixer output (model input) for sample {b}",
            type=ArtifactType.IMAGE,
            stage=context.stage,
            epoch=context.epoch,
            batch_idx=context.batch_idx,
        )
        artifacts.append(artifact)

        # 3. Ground Truth Mask
        mask_img = self._create_mask_visualization(mask_np[b])  # [H, W, 1] -> [H, W, 3]
        artifact = Artifact(
            name=f"ground_truth_mask_sample_{b}",
            value=mask_img,
            el_id=b,
            desc=f"Ground truth anomaly mask for sample {b}",
            type=ArtifactType.IMAGE,
            stage=context.stage,
            epoch=context.epoch,
            batch_idx=context.batch_idx,
        )
        artifacts.append(artifact)

        # 4. Anomaly Scores (as heatmap)
        scores_img = self._create_scores_heatmap(scores_np[b])  # [H, W, 1] -> [H, W, 3]
        artifact = Artifact(
            name=f"adaclip_scores_heatmap_sample_{b}",
            value=scores_img,
            el_id=b,
            desc=f"Anomaly scores (heatmap) for sample {b}",
            type=ArtifactType.IMAGE,
            stage=context.stage,
            epoch=context.epoch,
            batch_idx=context.batch_idx,
        )
        artifacts.append(artifact)

    return {"artifacts": artifacts}