Files
Head_CT_Morph/generate_head_extension_video.py

154 lines
4.3 KiB
Python

import os
import argparse
from pathlib import Path
import imageio.v2 as imageio
import numpy as np
from PIL import Image, ImageDraw, ImageFont
from head_extension_app import (
ct_window,
fit_image,
load_dicom_volume,
)
OUTPUT_DIR = Path("ppt_video")
FPS = 30
DURATION_SECONDS = 6
END_HOLD_SECONDS = 1
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 make_frame(slice_image, slice_index, slice_count, source_label, show_arrow=True):
frame = Image.new("RGB", (1920, 1080), (0, 0, 0))
draw = ImageDraw.Draw(frame)
panel = fit_image(slice_image, 1160, 860)
frame.paste(panel, (380, 145))
title_font = get_font(48)
small_font = get_font(28)
draw.text((410, 70), f"{source_label} DICOM z-axis sequence", font=title_font, fill=(255, 255, 255))
draw.text(
(410, 1015),
f"Slice {slice_index + 1} / {slice_count}",
font=small_font,
fill=(210, 210, 210),
)
arrow = (255, 210, 60)
if show_arrow:
x0, y0, x1, y1 = 1320, 275, 1450, 205
draw.line((x0, y0, x1, y1), fill=arrow, width=8)
draw.polygon([(x1, y1), (x1 - 34, y1 + 3), (x1 - 12, y1 + 29)], fill=arrow)
bar_x, bar_y, bar_w, bar_h = 650, 1018, 870, 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 * (slice_index + 1) / slice_count) if slice_count else 0
draw.rounded_rectangle(
(bar_x, bar_y, bar_x + fill_w, bar_y + bar_h),
radius=6,
fill=arrow,
)
return frame
def parse_args():
parser = argparse.ArgumentParser(
description="Generate a z-axis DICOM sequence 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="Kept for compatibility; z-axis sequence videos do not use this value.",
)
parser.add_argument(
"--duration",
type=float,
default=DURATION_SECONDS,
help="Animation duration in seconds before final hold. Default: 6",
)
parser.add_argument(
"--no-arrow",
action="store_true",
help="Hide the yellow direction arrow.",
)
parser.add_argument(
"--source-label",
default="DICOM",
help="Label shown in the video title.",
)
return parser.parse_args()
def generate_video(input_dir, output_path, max_angle=20.0, duration_seconds=6.0, show_arrow=True, source_label="DICOM"):
output_file = Path(output_path)
output_file.parent.mkdir(parents=True, exist_ok=True)
volume = load_dicom_volume(input_dir)
slice_count = volume.shape[0]
moving_frames = max(1, 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) if moving_frames > 1 else 0
slice_index = min(slice_count - 1, int(round(t * (slice_count - 1))))
slice_image = Image.fromarray(ct_window(volume[slice_index])).convert("RGB")
writer.append_data(np.asarray(make_frame(slice_image, slice_index, slice_count, source_label, show_arrow)))
for _ in range(hold_frames):
slice_image = Image.fromarray(ct_window(volume[-1])).convert("RGB")
writer.append_data(np.asarray(make_frame(slice_image, slice_count - 1, slice_count, source_label, show_arrow)))
return output_file.resolve()
def main():
args = parse_args()
output_file = generate_video(
args.input,
args.output,
args.max_angle,
args.duration,
not args.no_arrow,
args.source_label,
)
print(output_file)
if __name__ == "__main__":
main()