Video Add Video Overlay

1import os2import sys3import json4import subprocess5import traceback67INPUT_DIR = "/data/input"8OUTPUT_DIR = "/data/output"910FFMPEG_IMAGE = "jrottenberg/ffmpeg:7-ubuntu"111213def get_overlay_duration(host_input, overlay_filename):14    """Use ffprobe to get the duration of the overlay video."""15    cmd = [16        "docker", "run", "--rm",17        "--network", "none",18        "--memory", "512m",19        "--cpus", "1.0",20        "-v", f"{host_input}:/data/input:ro",21        FFMPEG_IMAGE,22        "-v", "error",23        "-show_entries", "format=duration",24        "-of", "default=noprint_wrappers=1:nokey=1",25        "/data/input/" + overlay_filename,26    ]27    # Replace the entrypoint with ffprobe28    cmd.insert(4, "--entrypoint")29    cmd.insert(5, "ffprobe")30    # Remove the FFMPEG_IMAGE from its current position and reinsert after entrypoint args31    # Rebuild command properly32    cmd = [33        "docker", "run", "--rm",34        "--network", "none",35        "--memory", "512m",36        "--cpus", "1.0",37        "--entrypoint", "ffprobe",38        "-v", f"{host_input}:/data/input:ro",39        FFMPEG_IMAGE,40        "-v", "error",41        "-show_entries", "format=duration",42        "-of", "default=noprint_wrappers=1:nokey=1",43        "/data/input/" + overlay_filename,44    ]45    result = subprocess.run(cmd, capture_output=True, text=True)46    if result.returncode != 0:47        raise RuntimeError(f"ffprobe failed: {result.stderr}")48    return float(result.stdout.strip())495051def main():52    execution_input = json.loads(sys.stdin.read())53    inputs = execution_input.get("inputs", {})5455    # Required inputs56    video = inputs["video"]57    overlay_video = inputs["overlay_video"]5859    # Position and scale60    x = inputs.get("x", 0)61    y = inputs.get("y", 0)62    video_scale = inputs.get("video_scale", 200)6364    # Duration and audio65    output_duration = inputs.get("output_duration", "video_1")66    audio_source = inputs.get("audio_source", "video_1")6768    # Crop options69    crop = inputs.get("crop", "no")70    crop_width = inputs.get("crop_width", 640)71    crop_height = inputs.get("crop_height", 360)7273    # Chromakey options74    chromakey = inputs.get("chromakey", "no")75    chromakey_color = inputs.get("chromakey_color", "00FF00")76    chromakey_similarity = inputs.get("chromakey_similarity", 0.3)77    chromakey_blend = inputs.get("chromakey_blend", 0.15)7879    os.makedirs(OUTPUT_DIR, exist_ok=True)8081    # Host staging paths for sibling Docker container82    host_input = os.environ.get("HOST_STAGING_INPUT", INPUT_DIR)83    host_output = os.environ.get("HOST_STAGING_OUTPUT", OUTPUT_DIR)8485    # Build filter chain for overlay video [1:v]86    overlay_filters = []8788    if chromakey == "yes":89        overlay_filters.append(90            f"colorkey=0x{chromakey_color}:{chromakey_similarity}:{chromakey_blend}"91        )9293    if crop == "yes":94        overlay_filters.append(95            f"crop={crop_width}:{crop_height}:(in_w-{crop_width})/2:(in_h-{crop_height})/2"96        )9798    overlay_filters.append(f"scale=trunc({video_scale}/2)*2:-2")99100    # Two-stage shell script approach to avoid heap corruption in ffmpeg 7's101    # overlay filter with simultaneous dual video decoders in Docker.102    # Stage 1 normalizes the overlay, Stage 2 composites onto the base.103    video_path = f"/data/input/{video}"104    overlay_path = f"/data/input/{overlay_video}"105    scale_filter = ",".join(overlay_filters)106107    # Stage 1: normalize overlay video (apply crop/chromakey/scale to temp file)108    stage1 = (109        f"ffmpeg -y -threads 4 -i '{overlay_path}' "110        f"-vf '{scale_filter}' "111        f"-c:v libx264 -preset veryfast -crf 18 -pix_fmt yuv420p "112        f"-c:a aac -ar 44100 -ac 2 "113        f"/tmp/_overlay_scaled.mp4"114    )115116    # Stage 2: overlay scaled video onto base117    if audio_source == "mix":118        overlay_fc = f"[0:v][1:v]overlay={x}:{y}[vout];[0:a][1:a]amix=inputs=2:duration=longest[aout]"119        audio_cmd = "-map '[vout]' -map '[aout]'"120    elif audio_source == "video_2":121        overlay_fc = f"[0:v][1:v]overlay={x}:{y}[vout]"122        audio_cmd = "-map '[vout]' -map 1:a -codec:a copy"123    else:124        overlay_fc = f"[0:v][1:v]overlay={x}:{y}[vout]"125        audio_cmd = "-map '[vout]' -map 0:a -codec:a copy"126127    duration_cmd = ""128    if output_duration == "video_2":129        dur = get_overlay_duration(host_input, overlay_video)130        duration_cmd = f"-t {dur}"131132    stage2 = (133        f"ffmpeg -y -threads 4 -i '{video_path}' -i /tmp/_overlay_scaled.mp4 "134        f"-filter_complex '{overlay_fc}' "135        f"{audio_cmd} "136        f"-c:v libx264 -preset veryfast -crf 23 -pix_fmt yuv420p "137        f"-movflags +faststart "138        f"{duration_cmd} "139        f"/data/output/output.mp4"140    )141142    shell_script = f"set -e\n{stage1}\n{stage2}\n"143144    cmd = [145        "docker", "run", "--rm",146        "--network", "none",147        "--memory", "4g",148        "--cpus", "4.0",149        "--cpuset-cpus", "0-3",150        "-v", f"{host_input}:/data/input:ro",151        "-v", f"{host_output}:/data/output:rw",152        "--entrypoint", "sh",153        FFMPEG_IMAGE,154        "-c", shell_script,155    ]156157    result = subprocess.run(cmd, capture_output=True, text=True, timeout=1800)158    if result.returncode != 0:159        raise RuntimeError(f"ffmpeg failed (exit {result.returncode}): {result.stderr[-2000:]}")160161    print(json.dumps({"video": "output.mp4"}, indent=2))162163164if __name__ == "__main__":165    try:166        main()167    except Exception as e:168        print(json.dumps({169            "error": str(e),170            "errorType": type(e).__name__,171            "traceback": traceback.format_exc(),172        }), file=sys.stderr)173        sys.exit(1)