"""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)