utils/visualization.py

"""
Drawing helpers for pipeline visual outputs.

All functions operate on RGB numpy arrays (uint8) and return RGB numpy arrays.
They are stateless - no model or pipeline state is needed.
"""

import cv2
import numpy as np


# Detection overlay

def make_detection_overlay(
    image_rgb  : np.ndarray,
    boxes      : list[tuple],
    label      : str | None = None,
) -> np.ndarray:
    """
    Draw YOLO bounding boxes on the image.

    The box colour reflects the final classification label:
      - MALIGNANT -> red   (#DC3232)
      - BENIGN    -> green (#32B450)
      - Unknown   -> grey  (#B4B4B4)

    Parameters
    image_rgb : np.ndarray (H, W, 3) uint8, RGB
    boxes     : list of (x1, y1, x2, y2, conf) tuples from NoduleDetector
    label     : "MALIGNANT", "BENIGN", or None (classification not yet done)

    Returns
    -------
    np.ndarray (H, W, 3) uint8, RGB - image with bounding boxes drawn
    """
    overlay = image_rgb.copy()

    if label == "MALIGNANT":
        color = (220, 50, 50)
    elif label == "BENIGN":
        color = (50, 180, 80)
    else:
        color = (180, 180, 180)

    for (x1, y1, x2, y2, conf) in boxes:
        cv2.rectangle(overlay, (x1, y1), (x2, y2), color, 2)
        text  = f"Nodule  {conf:.0%}"
        org   = (x1, max(y1 - 8, 12))
        cv2.putText(
            overlay, text, org,
            cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2,
            cv2.LINE_AA,
        )

    return overlay


# Segmentation overlay

def make_segmentation_overlay(
    image_rgb : np.ndarray,
    mask      : np.ndarray,
    alpha     : float = 0.40,
) -> np.ndarray:
    """
    Blend the UNet++ binary mask as a semi-transparent orange overlay.

    A solid contour is drawn on top for clearer boundary visualisation.

    Parameters
    image_rgb : np.ndarray (H, W, 3) uint8, RGB
    mask      : np.ndarray (H, W)    uint8, 0/255
    alpha     : float — opacity of the colour overlay (default 0.40)

    Returns
    np.ndarray (H, W, 3) uint8, RGB - blended image
    """
    color_mask = np.zeros_like(image_rgb)
    color_mask[mask > 0] = [255, 140, 0]   # orange

    blended = cv2.addWeighted(image_rgb, 1.0,
                              color_mask, alpha, 0)

    # Draw contour outline for precise boundary visualisation
    contours, _ = cv2.findContours(
        mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
    )
    cv2.drawContours(blended, contours, -1, (255, 140, 0), 2)

    return blended