webapp/render.py

import os
import shutil
import subprocess

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import cv2

OUT_W, OUT_H = 1920, 1080
VID_W, VID_H = OUT_W // 2, OUT_H   # 960 x 1080
CHART_W, CHART_H = OUT_W // 2, OUT_H  # 960 x 1080


def build_chart(prediction, top_k_chart=3):
    """Render the saliency curve + top predicted windows to a 960x1080 BGR image.

    Returns (chart_bgr, x_left_px, x_right_px) where x_left_px/x_right_px are
    the exact pixel x-coordinates of t=0 and t=video_duration on the chart,
    used later to draw a moving time-cursor in sync with video playback.
    """
    saliency = np.array(prediction["pred_saliency_scores"], dtype=np.float32)
    clip_len = prediction["clip_len"]
    video_duration = prediction["video_duration"]
    time_axis = np.arange(len(saliency)) * clip_len

    fig, ax = plt.subplots(figsize=(CHART_W / 100, CHART_H / 100), dpi=100)
    ax.plot(
        time_axis, saliency, marker="o", linestyle="-", linewidth=2,
        markersize=5, color="#1f77b4", label="Saliency Score",
    )

    windows = prediction["pred_relevant_windows"][:top_k_chart]
    colors = plt.cm.Set1(np.linspace(0, 1, max(len(windows), 1)))
    for i, (st, ed, sc) in enumerate(windows):
        ax.axvspan(
            st, ed, color=colors[i], alpha=0.3,
            label=f"Top-{i + 1} [{st:.1f}s-{ed:.1f}s] score={sc:.2f}",
        )

    ax.set_xlim(0, max(video_duration, 1e-6))
    ax.set_xlabel("Time (seconds)", fontsize=12, fontweight="bold")
    ax.set_ylabel("Saliency Score", fontsize=12, fontweight="bold")
    ax.set_title(f"Query: {prediction['query']}", fontsize=13, fontweight="bold")
    ax.grid(True, linestyle="--", alpha=0.6)
    ax.legend(loc="upper right", fontsize=8, framealpha=0.9)
    fig.tight_layout()
    fig.canvas.draw()

    y0 = ax.get_ylim()[0]
    x_left_px, _ = ax.transData.transform((0, y0))
    x_right_px, _ = ax.transData.transform((video_duration, y0))

    w, h = fig.canvas.get_width_height()
    buf = np.frombuffer(fig.canvas.buffer_rgba(), dtype=np.uint8).reshape(h, w, 4)
    chart_bgr = cv2.cvtColor(buf, cv2.COLOR_RGBA2BGR)
    plt.close(fig)

    return chart_bgr, int(round(x_left_px)), int(round(x_right_px))


def _resize_pad(img, tw, th, bg=(20, 20, 20)):
    h, w = img.shape[:2]
    scale = min(tw / w, th / h)
    nw, nh = max(int(w * scale), 1), max(int(h * scale), 1)
    small = cv2.resize(img, (nw, nh), interpolation=cv2.INTER_AREA)
    canvas = np.full((th, tw, 3), bg, dtype=np.uint8)
    y0 = (th - nh) // 2
    x0 = (tw - nw) // 2
    canvas[y0:y0 + nh, x0:x0 + nw] = small
    return canvas


def _put_label(img, text, pos=(12, 40), color=(220, 220, 220)):
    cv2.putText(img, text, pos, cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 0, 0), 4, cv2.LINE_AA)
    cv2.putText(img, text, pos, cv2.FONT_HERSHEY_SIMPLEX, 0.9, color, 2, cv2.LINE_AA)


def _reencode_h264(src_path, dst_path):
    cmd = [
        "ffmpeg", "-y", "-i", src_path,
        "-c:v", "libx264", "-pix_fmt", "yuv420p", "-movflags", "+faststart",
        dst_path,
    ]
    try:
        subprocess.run(cmd, check=True, capture_output=True)
    except (subprocess.CalledProcessError, FileNotFoundError):
        shutil.copyfile(src_path, dst_path)


def render_visualization_video(video_path, prediction, chart_bgr, x_left_px, x_right_px, out_path):
    """Write a side-by-side (video | chart-with-moving-cursor) mp4 to out_path.

    Only renders up to prediction['video_duration'] seconds of the source
    video, so the chart's time axis always matches what's shown on screen
    (relevant when the source video is longer than the 150s the model
    actually analyzed).
    """
    duration = prediction["video_duration"]

    raw_path = out_path + ".raw.mp4"

    try:
        cap = cv2.VideoCapture(video_path)
        in_fps = cap.get(cv2.CAP_PROP_FPS)
        if not in_fps or in_fps <= 0:
            in_fps = 25.0
        out_fps = min(in_fps, 30.0)

        fourcc = cv2.VideoWriter_fourcc(*"mp4v")
        writer = cv2.VideoWriter(raw_path, fourcc, out_fps, (OUT_W, OUT_H))

        frame_idx = 0
        out_frame_idx = 0
        while True:
            ret, frame = cap.read()
            if not ret:
                break
            current_time = frame_idx / in_fps
            if current_time > duration:
                break
            progress = min(current_time / duration, 1.0) if duration > 0 else 0.0
            target_out = int(current_time * out_fps)

            if target_out >= out_frame_idx:
                vid_panel = _resize_pad(frame, VID_W, VID_H)
                ts = f"{int(current_time // 60):02d}:{current_time % 60:05.2f}"
                _put_label(vid_panel, ts)

                chart_panel = chart_bgr.copy()
                xp = int(x_left_px + progress * (x_right_px - x_left_px))
                cv2.line(chart_panel, (xp, 0), (xp, CHART_H), (0, 60, 255), 3)

                combined = np.hstack([vid_panel, chart_panel])
                while out_frame_idx <= target_out:
                    writer.write(combined)
                    out_frame_idx += 1
            frame_idx += 1
    finally:
        cap.release()
        writer.release()

    _reencode_h264(raw_path, out_path)
    os.remove(raw_path)
    return out_path