vibe_ml_lab/utils.py

"""Shared helpers for visualization and array handling."""

from __future__ import annotations

import math
from io import BytesIO
from typing import Iterable, Sequence

import cv2
import matplotlib
import numpy as np
from matplotlib import pyplot as plt

matplotlib.use("Agg")


def to_rgb(image: np.ndarray) -> np.ndarray:
    """Convert a grayscale or BGR image into RGB uint8."""
    if image.ndim == 2:
        return cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
    if image.shape[2] == 4:
        return cv2.cvtColor(image, cv2.COLOR_RGBA2RGB)
    return cv2.cvtColor(image, cv2.COLOR_BGR2RGB)


def normalize_uint8(image: np.ndarray) -> np.ndarray:
    """Scale a float image to uint8 if needed."""
    if image.dtype == np.uint8:
        return image
    clipped = np.clip(image, 0.0, 1.0)
    return (clipped * 255).astype(np.uint8)


def figure_to_image(fig: plt.Figure) -> np.ndarray:
    """Render a Matplotlib figure into an RGB numpy image."""
    buffer = BytesIO()
    fig.savefig(buffer, format="png", bbox_inches="tight", dpi=160)
    plt.close(fig)
    buffer.seek(0)
    data = np.frombuffer(buffer.getvalue(), dtype=np.uint8)
    decoded = cv2.imdecode(data, cv2.IMREAD_COLOR)
    return cv2.cvtColor(decoded, cv2.COLOR_BGR2RGB)


def plot_confusion_matrix(
    matrix: np.ndarray,
    labels: Sequence[str],
    title: str,
    cmap: str = "YlOrRd",
) -> np.ndarray:
    """Create a labeled confusion-matrix heatmap."""
    fig, ax = plt.subplots(figsize=(6.2, 5.2))
    heat = ax.imshow(matrix, cmap=cmap)
    ax.set_title(title, fontsize=14, fontweight="bold")
    ax.set_xticks(range(len(labels)))
    ax.set_yticks(range(len(labels)))
    ax.set_xticklabels(labels, rotation=25, ha="right")
    ax.set_yticklabels(labels)
    ax.set_xlabel("Predicted")
    ax.set_ylabel("True")
    ax.figure.colorbar(heat, ax=ax, shrink=0.82)

    threshold = float(matrix.max()) * 0.55 if matrix.size else 0.0
    for row in range(matrix.shape[0]):
        for col in range(matrix.shape[1]):
            value = int(matrix[row, col])
            color = "white" if value >= threshold else "#2f1d0b"
            ax.text(col, row, value, ha="center", va="center", color=color, fontsize=11)

    fig.tight_layout()
    return figure_to_image(fig)


def plot_lines(
    series: Sequence[tuple[str, Sequence[float], str]],
    title: str,
    xlabel: str = "Epoch",
    ylabel: str = "Value",
) -> np.ndarray:
    """Plot multiple lines on a single axis."""
    fig, ax = plt.subplots(figsize=(7.2, 4.0))
    for name, values, color in series:
        xs = np.arange(1, len(values) + 1)
        ax.plot(xs, values, label=name, linewidth=2.2, color=color)
    ax.set_title(title, fontsize=14, fontweight="bold")
    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)
    ax.grid(alpha=0.25)
    ax.legend(frameon=False)
    fig.tight_layout()
    return figure_to_image(fig)


def plot_bar(
    labels: Sequence[str],
    values: Sequence[float],
    title: str,
    color: str = "#ef7d32",
    ylabel: str = "Value",
) -> np.ndarray:
    """Create a simple labeled bar chart."""
    fig, ax = plt.subplots(figsize=(7.0, 4.0))
    positions = np.arange(len(labels))
    ax.bar(positions, values, color=color, edgecolor="#7a3310")
    ax.set_title(title, fontsize=14, fontweight="bold")
    ax.set_xticks(positions)
    ax.set_xticklabels(labels, rotation=20, ha="right")
    ax.set_ylabel(ylabel)
    ax.grid(axis="y", alpha=0.25)

    upper = max(values) if values else 1.0
    for idx, value in enumerate(values):
        ax.text(idx, value + upper * 0.02, f"{value:.2f}", ha="center", va="bottom", fontsize=10)

    fig.tight_layout()
    return figure_to_image(fig)


def make_image_grid(
    images: Sequence[np.ndarray],
    captions: Sequence[str] | None = None,
    columns: int = 4,
    tile_size: tuple[int, int] | None = None,
    background: tuple[int, int, int] = (247, 241, 232),
) -> np.ndarray:
    """Lay out images on a canvas with optional captions."""
    if not images:
        blank = np.full((240, 320, 3), background, dtype=np.uint8)
        cv2.putText(
            blank,
            "No images",
            (92, 126),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.9,
            (74, 67, 61),
            2,
            cv2.LINE_AA,
        )
        return blank

    rgb_images = [normalize_uint8(to_rgb(image)) for image in images]
    if tile_size is None:
        tile_h = max(image.shape[0] for image in rgb_images)
        tile_w = max(image.shape[1] for image in rgb_images)
    else:
        tile_w, tile_h = tile_size

    columns = max(1, columns)
    rows = math.ceil(len(rgb_images) / columns)
    caption_h = 28 if captions else 0
    canvas = np.full(
        (rows * (tile_h + caption_h + 18) + 20, columns * (tile_w + 18) + 20, 3),
        background,
        dtype=np.uint8,
    )

    for index, image in enumerate(rgb_images):
        row = index // columns
        col = index % columns
        y = 12 + row * (tile_h + caption_h + 18)
        x = 12 + col * (tile_w + 18)
        resized = cv2.resize(image, (tile_w, tile_h), interpolation=cv2.INTER_AREA)
        canvas[y : y + tile_h, x : x + tile_w] = resized
        cv2.rectangle(canvas, (x, y), (x + tile_w, y + tile_h), (201, 165, 130), 1)

        if captions:
            text = captions[index]
            cv2.putText(
                canvas,
                text,
                (x, y + tile_h + 19),
                cv2.FONT_HERSHEY_SIMPLEX,
                0.52,
                (84, 60, 42),
                1,
                cv2.LINE_AA,
            )

    return canvas


def softmax(logits: np.ndarray) -> np.ndarray:
    """Numerically stable softmax."""
    shifted = logits - logits.max(axis=1, keepdims=True)
    exp_values = np.exp(shifted)
    return exp_values / exp_values.sum(axis=1, keepdims=True)