File size: 5,273 Bytes
592dffb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 | 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
|