tt06-urish-charge-pump-flythrough.py

rendervid.py

import mitsuba as mi
import argparse
from datetime import datetime
import time
import subprocess
import os

RENDER_WIDTH = 1920
RENDER_HEIGHT = 1080
RENDER_SPP = 256
VARIANT = "cuda_ad_rgb"
OUTPUT_NAME = f"scene_tinytapeout_{datetime.now().strftime('%Y%m%d_%H%M%S')}.png"
SENSOR_INDEX = 0

#########################
# 1) Parse Arguments
#########################
argparser = argparse.ArgumentParser(description="Render scene_tinytapeout.xml with optional camera flythrough.")

argparser.add_argument(
    "-width", "--output_width", required=False, type=int,
    help="Output resolution width"
)
argparser.add_argument(
    "-height", "--output_height", required=False, type=int,
    help="Output resolution height"
)
argparser.add_argument(
    "-spp", "--samples_per_pixel", required=False, type=int,
    help="Render samples per pixel"
)
argparser.add_argument(
    "-v", "--variant", required=False, type=str,
    choices=["scalar_rgb", "cuda_ad_rgb", "llvm_ad_rgb", "scalar_spectral"],
    help="Mitsuba variant"
)
argparser.add_argument(
    "-o", "--output", required=False, type=str,
    help="Output filename"
)

# New arguments for animation:
argparser.add_argument(
    "--duration", required=False, type=float, default=10.0,
    help="Animation duration in seconds (default: 10)"
)
argparser.add_argument(
    "--fps", required=False, type=int, default=60,
    help="Frames per second (default: 60)"
)
argparser.add_argument(
    "--fly", action="store_true",
    help="If set, perform the camera flythrough instead of a single frame."
)

args = vars(argparser.parse_args())

#########################
# 2) Set Defaults if Provided
#########################
if args["output_width"] is not None:
    RENDER_WIDTH = args["output_width"]
if args["output_height"] is not None:
    RENDER_HEIGHT = args["output_height"]
if args["samples_per_pixel"] is not None:
    RENDER_SPP = args["samples_per_pixel"]
if args["variant"] is not None:
    VARIANT = args["variant"]
if args["output"] is not None:
    OUTPUT_NAME = args["output"]

mi.set_variant(VARIANT)

#########################
# 3) Load Scene
#########################
scene = mi.load_file("scene_tinytapeout.xml")
sensor = scene.sensors()[SENSOR_INDEX]

# We can adjust film resolution on the chosen sensor:
film_params = mi.traverse(sensor)
film_params["film.size"] = mi.ScalarVector2u(RENDER_WIDTH, RENDER_HEIGHT)
film_params.update()

print(f"Rendering {RENDER_WIDTH}x{RENDER_HEIGHT} image with {RENDER_SPP} spp.")
print(f"Variant: {mi.variant()}")

#########################
# Single-Frame Render (if --fly is NOT used)
#########################
if not args["fly"]:
    process_start_time = time.time()
    img = mi.render(scene, spp=RENDER_SPP, sensor=SENSOR_INDEX)
    mi.util.write_bitmap(OUTPUT_NAME, img, False)
    print(f"Wrote image to {OUTPUT_NAME}")
    print(f"Elapsed time: {time.time() - process_start_time:.2f} s")
    exit(0)

#########################
# 4) Flythrough (if --fly is used)
#########################

# Start extra high, far away:
start_origin = [40.0, 40.0, 150.0]
start_target = [50.0, 80.0,  30.0]

# End closer to the ground, more forward:
end_origin   = [80.0, 80.0, 20.0]
end_target   = [80.0, 160.0, 5.0]

up_vector    = [0.0, 1.0, 0.0]

fps = args["fps"]
duration = args["duration"]
num_frames = int(fps * duration)

print(f"Performing flythrough from {start_origin} to {end_origin} over {duration}s at {fps} fps.")
print(f"Total frames: {num_frames}")

# Create an output directory for frames
frames_dir = "flythrough_frames"
os.makedirs(frames_dir, exist_ok=True)

def smootherstep(t):
    return 6*t**5 - 15*t**4 + 10*t**3

def lerp(a, b, t):
    return (1 - t) * a + t * b

def smooth_lerp_vec3(a, b, t):
    te = smootherstep(t)  # eased version of t
    return [
        lerp(a[0], b[0], te),
        lerp(a[1], b[1], te),
        lerp(a[2], b[2], te),
    ]

process_start_time = time.time()

for frame_idx in range(num_frames):
    # progress fraction from 0 -> 1
    if num_frames > 1:
        raw_t = frame_idx / (num_frames - 1)
    else:
        raw_t = 0.0

    # Use smooth_lerp_vec3 instead of linear interpolation
    current_origin = smooth_lerp_vec3(start_origin, end_origin, raw_t)
    current_target = smooth_lerp_vec3(start_target, end_target, raw_t)

    # Build the lookat transform
    lookat_transform = mi.ScalarTransform4f().look_at(
        origin=mi.ScalarPoint3f(*current_origin),
        target=mi.ScalarPoint3f(*current_target),
        up=mi.ScalarPoint3f(*up_vector)
    )
    
    # Update sensor's to_world transform
    sensor_params = mi.traverse(sensor)
    sensor_params["to_world"] = lookat_transform
    sensor_params.update()

    # Render
    img = mi.render(scene, spp=RENDER_SPP, sensor=SENSOR_INDEX)

    # Save frame
    frame_file = os.path.join(frames_dir, f"frame_{frame_idx:04d}.png")
    mi.util.write_bitmap(frame_file, img, False)

    print(f"Rendered frame {frame_idx+1}/{num_frames} -> {frame_file}")

elapsed_render_time = time.time() - process_start_time
print(f"All frames rendered in {elapsed_render_time:.2f} s")

#########################
# 5) Invoke ffmpeg to combine frames
#########################
# Example ffmpeg command to create an .mp4 at the same fps:
#   ffmpeg -framerate 60 -i frame_%04d.png -pix_fmt yuv420p out.mp4

video_filename = "flythrough.mp4"
cmd = [
    "ffmpeg",
    "-y",  # overwrite output if it exists
    "-framerate", str(fps),
    "-i", os.path.join(frames_dir, "frame_%04d.png"),
    "-pix_fmt", "yuv420p",
    video_filename
]

print("Running ffmpeg to produce final video...")
try:
    subprocess.run(cmd, check=True)
    print(f"Video saved to {video_filename}")
except subprocess.CalledProcessError as e:
    print(f"Failed to run ffmpeg: {e}")