deepivus/visualization/video.py

"""Video rendering utilities."""

import numpy as np
from PIL import Image, ImageDraw

from ..processing.preprocessing import ensure_uint8


def draw_contour(draw: ImageDraw.ImageDraw, x_points, y_points, color):
    """Draw a closed contour if points are available."""
    if not x_points or not y_points:
        return
    points = [(float(x_points[i]), float(y_points[i])) for i in range(len(x_points))]
    if len(points) < 2:
        return
    points.append(points[0])
    draw.line(points, fill=color, width=2)


def _require_imageio():
    try:
        import imageio.v2 as imageio
    except Exception as exc:
        raise RuntimeError(
            "Video export requires imageio and imageio-ffmpeg. "
            "Install with: python -m pip install imageio imageio-ffmpeg"
        ) from exc
    return imageio


def write_overlay_video(images: np.ndarray, lumen, plaque, video_path: str, fps: float) -> None:
    """Write contour overlay video with lumen/plaque channels."""
    imageio = _require_imageio()
    images_u8 = ensure_uint8(images)
    with imageio.get_writer(video_path, fps=fps) as writer:
        for frame_idx in range(images_u8.shape[0]):
            frame = Image.fromarray(images_u8[frame_idx], mode="L").convert("RGB")
            draw = ImageDraw.Draw(frame)
            draw_contour(draw, lumen[0][frame_idx], lumen[1][frame_idx], (0, 255, 0))
            draw_contour(draw, plaque[0][frame_idx], plaque[1][frame_idx], (255, 64, 64))
            writer.append_data(np.asarray(frame))


def write_overlay_video_with_bifurcation_flags(
    images: np.ndarray,
    lumen,
    plaque,
    video_path: str,
    fps: float,
    bifurcation_probabilities: np.ndarray,
    bifurcation_labels: np.ndarray,
    threshold: float,
) -> None:
    """Write contour overlay video with per-frame bifurcation text labels."""
    imageio = _require_imageio()
    images_u8 = ensure_uint8(images)
    probs = np.asarray(bifurcation_probabilities, dtype=np.float32)
    labels = np.asarray(bifurcation_labels, dtype=np.int32)
    with imageio.get_writer(video_path, fps=fps) as writer:
        for frame_idx in range(images_u8.shape[0]):
            frame = Image.fromarray(images_u8[frame_idx], mode="L").convert("RGB")
            draw = ImageDraw.Draw(frame)
            draw_contour(draw, lumen[0][frame_idx], lumen[1][frame_idx], (0, 255, 0))
            draw_contour(draw, plaque[0][frame_idx], plaque[1][frame_idx], (255, 64, 64))

            if frame_idx < probs.shape[0] and frame_idx < labels.shape[0]:
                prob = float(probs[frame_idx])
                is_bif = bool(labels[frame_idx])
                tag = "Branch" if is_bif else "Non-branch"
                text = f"{tag}  p={prob:.2f}  t={float(threshold):.2f}"
                text_color = (255, 90, 90) if is_bif else (100, 220, 255)
                draw.rectangle([(10, 10), (290, 36)], fill=(0, 0, 0))
                draw.text((16, 16), text, fill=text_color)

            writer.append_data(np.asarray(frame))


def write_model_comparison_video(images: np.ndarray, tf_lumen, sam_lumen, video_path: str, fps: float) -> None:
    """Write a two-model comparison video (TF red, SAM blue)."""
    imageio = _require_imageio()
    images_u8 = ensure_uint8(images)
    with imageio.get_writer(video_path, fps=fps) as writer:
        for frame_idx in range(images_u8.shape[0]):
            frame = Image.fromarray(images_u8[frame_idx], mode="L").convert("RGB")
            draw = ImageDraw.Draw(frame)
            draw_contour(draw, tf_lumen[0][frame_idx], tf_lumen[1][frame_idx], (255, 0, 0))
            draw_contour(draw, sam_lumen[0][frame_idx], sam_lumen[1][frame_idx], (0, 0, 255))
            writer.append_data(np.asarray(frame))


def _draw_graph_panel(frame_rgb: np.ndarray, scores: np.ndarray, sustained_flags: np.ndarray, frame_idx: int) -> np.ndarray:
    panel_h = 150
    h, w = frame_rgb.shape[:2]
    canvas = Image.new("RGB", (w, h + panel_h), color=(0, 0, 0))
    canvas.paste(Image.fromarray(frame_rgb), (0, 0))
    draw = ImageDraw.Draw(canvas)

    top = h + 12
    left = 16
    right = w - 16
    bottom = h + panel_h - 18
    draw.rectangle([(left, top), (right, bottom)], fill=(20, 20, 20), outline=(90, 90, 90), width=1)

    n = len(scores)
    if n > 1:
        in_run = False
        run_start = 0
        for i in range(n):
            if sustained_flags[i] and not in_run:
                run_start = i
                in_run = True
            if (not sustained_flags[i] or i == n - 1) and in_run:
                run_end = i if not sustained_flags[i] else i + 1
                x0 = left + int((run_start / (n - 1)) * (right - left))
                x1 = left + int(((run_end - 1) / (n - 1)) * (right - left))
                draw.rectangle([(x0, top), (x1, bottom)], fill=(70, 20, 20))
                in_run = False

        def px(i):
            return left + int((i / (n - 1)) * (right - left))

        def py(v):
            v = float(np.clip(v, 0.0, 1.0))
            return bottom - int(v * (bottom - top))

        pts = [(px(i), py(scores[i])) for i in range(n)]
        if len(pts) > 1:
            draw.line(pts, fill=(120, 255, 120), width=2)

        cur = min(max(frame_idx, 0), n - 1)
        cx = px(cur)
        draw.line([(cx, top), (cx, bottom)], fill=(255, 255, 0), width=2)
        flag_txt = "SUSTAINED BRANCH SIGNAL" if sustained_flags[cur] else "normal"
        draw.text((left + 6, top + 4), f"Oblongness: {float(scores[cur]):.3f}  [{flag_txt}]", fill=(230, 230, 230))
    else:
        draw.text((left + 6, top + 4), "Oblongness: insufficient data", fill=(230, 230, 230))

    return np.asarray(canvas)


def write_overlay_video_with_graph(
    images: np.ndarray,
    lumen,
    plaque,
    video_path: str,
    fps: float,
    oblong_scores: np.ndarray,
    sustained_flags: np.ndarray,
) -> None:
    """Write the fused overlay video with an animated bottom graph panel."""
    imageio = _require_imageio()
    images_u8 = ensure_uint8(images)
    with imageio.get_writer(video_path, fps=fps) as writer:
        for frame_idx in range(images_u8.shape[0]):
            frame = Image.fromarray(images_u8[frame_idx], mode="L").convert("RGB")
            draw = ImageDraw.Draw(frame)
            draw_contour(draw, lumen[0][frame_idx], lumen[1][frame_idx], (0, 255, 0))
            draw_contour(draw, plaque[0][frame_idx], plaque[1][frame_idx], (255, 64, 64))
            with_panel = _draw_graph_panel(np.asarray(frame), oblong_scores, sustained_flags, frame_idx)
            writer.append_data(with_panel)