Head_CT_Morph/generate_head_extension_video.py

import os
import argparse
from pathlib import Path

import imageio.v2 as imageio
import numpy as np
from PIL import Image, ImageDraw, ImageFont
from scipy.ndimage import map_coordinates

from head_extension_app import (
    crop_head_neck,
    fit_image,
    load_dicom_volume,
    sagittal_mip,
)


OUTPUT_DIR = Path("ppt_video")

FPS = 30
DURATION_SECONDS = 6
END_HOLD_SECONDS = 1


def video_soft_bend_2d(image, angle_degrees):
    """Video-only 2D deformation with a broad neck transition.

    The app's fast preview uses a compact transition, which is useful for
    interactive feedback but can visibly split vertebrae in animation. This
    broad ramp keeps the head and upper cervical spine moving together and
    blends gradually into the lower neck/shoulder region.
    """
    arr = np.asarray(image.convert("L")).astype(np.float32)
    height, width = arr.shape
    yy, xx = np.mgrid[0:height, 0:width]

    pivot_x = int(width * 0.55)
    pivot_y = int(height * 0.62)

    # Broad transition: nearly full motion for head/upper C-spine, then a long
    # smooth blend through the lower C-spine to avoid splitting bone structures.
    full_motion_y = height * 0.50
    fixed_y = height * 0.92
    t = np.clip((yy - full_motion_y) / (fixed_y - full_motion_y), 0, 1)
    weight = 1 - (t * t * (3 - 2 * t))

    # A small x-dependent term keeps the posterior contour from looking like a
    # straight sliced plane while remaining deterministic and smooth.
    x_soft = np.clip((xx - width * 0.15) / (width * 0.75), 0, 1)
    x_soft = x_soft * x_soft * (3 - 2 * x_soft)
    weight = np.clip(weight * (0.90 + 0.10 * x_soft), 0, 1)

    theta = np.deg2rad(angle_degrees) * weight
    cos_t = np.cos(theta)
    sin_t = np.sin(theta)
    dx = xx - pivot_x
    dy = yy - pivot_y

    src_x = pivot_x + cos_t * dx + sin_t * dy
    src_y = pivot_y - sin_t * dx + cos_t * dy

    warped = map_coordinates(arr, [src_y, src_x], order=1, mode="constant", cval=0)
    return Image.fromarray(np.clip(warped, 0, 255).astype(np.uint8)).convert("RGB")


def get_font(size):
    for path in [r"C:\Windows\Fonts\arial.ttf", r"C:\Windows\Fonts\calibri.ttf"]:
        if os.path.exists(path):
            return ImageFont.truetype(path, size)
    return ImageFont.load_default()


def smoothstep(t):
    return t * t * (3 - 2 * t)


def make_frame(before_image, angle, max_angle):
    after_image = video_soft_bend_2d(before_image, angle)

    frame = Image.new("RGB", (1920, 1080), (0, 0, 0))
    draw = ImageDraw.Draw(frame)

    left = fit_image(before_image, 760, 720)
    right = fit_image(after_image, 760, 720)
    frame.paste(left, (120, 245))
    frame.paste(right, (1040, 245))

    title_font = get_font(48)
    angle_font = get_font(56)
    draw.text((135, 120), "Original: 0 deg", font=title_font, fill=(255, 255, 255))
    draw.text(
        (1055, 120),
        f"Head extension: {angle:04.1f} deg",
        font=title_font,
        fill=(255, 255, 255),
    )

    # Yellow direction arrow on the animated side.
    arrow = (255, 210, 60)
    x0, y0, x1, y1 = 1390, 405, 1515, 335
    draw.line((x0, y0, x1, y1), fill=arrow, width=8)
    draw.polygon([(x1, y1), (x1 - 34, y1 + 3), (x1 - 12, y1 + 29)], fill=arrow)

    # Minimal progress bar.
    bar_x, bar_y, bar_w, bar_h = 1040, 980, 760, 12
    draw.rounded_rectangle(
        (bar_x, bar_y, bar_x + bar_w, bar_y + bar_h),
        radius=6,
        fill=(70, 70, 70),
    )
    fill_w = int(bar_w * angle / max_angle) if max_angle else 0
    draw.rounded_rectangle(
        (bar_x, bar_y, bar_x + fill_w, bar_y + bar_h),
        radius=6,
        fill=arrow,
    )
    draw.text((1040, 1000), "0 deg", font=get_font(24), fill=(210, 210, 210))
    draw.text((1745, 1000), f"{max_angle:g} deg", font=get_font(24), fill=(210, 210, 210))

    # Invisible-looking spacer: keeps the font imported above alive for some PIL builds.
    _ = angle_font
    return frame


def parse_args():
    parser = argparse.ArgumentParser(
        description="Generate a 0-degree to target-angle head-extension MP4 video."
    )
    parser.add_argument(
        "--input",
        default="input_ct_2F",
        help="Input DICOM folder. Default: input_ct_2F",
    )
    parser.add_argument(
        "--output",
        default=str(OUTPUT_DIR / "head_extension_0_to_20deg.mp4"),
        help="Output MP4 file path.",
    )
    parser.add_argument(
        "--max-angle",
        type=float,
        default=20.0,
        help="Target head-extension angle. Default: 20",
    )
    parser.add_argument(
        "--duration",
        type=float,
        default=DURATION_SECONDS,
        help="Animation duration in seconds before final hold. Default: 6",
    )
    return parser.parse_args()


def generate_video(input_dir, output_path, max_angle=20.0, duration_seconds=6.0):
    output_file = Path(output_path)
    output_file.parent.mkdir(parents=True, exist_ok=True)

    volume = load_dicom_volume(input_dir)
    before_image = crop_head_neck(sagittal_mip(volume))

    moving_frames = int(FPS * duration_seconds)
    hold_frames = FPS * END_HOLD_SECONDS

    with imageio.get_writer(
        output_file,
        format="FFMPEG",
        fps=FPS,
        codec="libx264",
        quality=8,
        macro_block_size=1,
    ) as writer:
        for index in range(moving_frames):
            t = index / (moving_frames - 1)
            angle = max_angle * smoothstep(t)
            writer.append_data(np.asarray(make_frame(before_image, angle, max_angle)))

        for _ in range(hold_frames):
            writer.append_data(np.asarray(make_frame(before_image, max_angle, max_angle)))

    return output_file.resolve()


def main():
    args = parse_args()
    output_file = generate_video(args.input, args.output, args.max_angle, args.duration)
    print(output_file)


if __name__ == "__main__":
    main()