trace_synthesis/video_gen.py
2025-04-15 22:44:08 +08:00

228 lines
8.8 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

import os
import glob
import re
import subprocess
import sys
import zipfile
import shutil
def extract_timestamp(filename):
"""Extracts the float timestamp from the filename."""
# Use regex for robustness, matching the pattern before '.jpeg'
# It looks for one or more digits, optionally followed by a dot and more digits,
# right before the .jpeg extension, preceded by a hyphen.
match = re.search(r'-(\d+(\.\d+)?)\.jpeg$', filename)
if match:
try:
return float(match.group(1))
except ValueError:
return None
return None
def create_video_from_images(image_folder, output_video_file, default_last_frame_duration=0.1, max_duration=5.0):
"""
Creates a WebM video from timestamped JPEG images in a folder.
Args:
image_folder (str): Path to the folder containing JPEG images.
output_video_file (str): Path for the output WebM video file.
default_last_frame_duration (float): Duration (in seconds) to display the last frame.
max_duration (float): Maximum duration (in seconds) for a single output video.
If total duration exceeds this, multiple videos will be created.
"""
print(f"Scanning folder: {image_folder}")
search_pattern = os.path.join(image_folder, 'page@*.jpeg')
image_files = glob.glob(search_pattern)
if not image_files:
print(f"Error: No JPEG files matching pattern '{search_pattern}' found.")
return
print(f"Found {len(image_files)} matching image files.")
# Extract timestamps and store as (timestamp, full_path) tuples
timed_files = []
for img_path in image_files:
timestamp = extract_timestamp(os.path.basename(img_path))
if timestamp is not None:
timed_files.append((timestamp, img_path))
else:
print(f"Warning: Could not extract timestamp from {os.path.basename(img_path)}. Skipping.")
if not timed_files:
print("Error: No files with valid timestamps found.")
return
# Sort files chronologically based on timestamp
timed_files.sort()
print(f"Processing {len(timed_files)} files with valid timestamps.")
# Split into segments based on max_duration
segments = []
current_segment = []
current_segment_duration = 0.0
for i in range(len(timed_files)):
timestamp, img_path = timed_files[i]
# Calculate duration for this frame
if i < len(timed_files) - 1:
next_timestamp, _ = timed_files[i+1]
duration = next_timestamp - timestamp
# Prevent zero or negative durations
if duration <= 0:
duration = 0.01
else:
# Duration for the last frame
duration = default_last_frame_duration
# Check if adding this frame would exceed max_duration
if current_segment_duration + duration > max_duration and current_segment:
# Current segment is full, start a new one
segments.append(current_segment)
current_segment = [(timestamp, img_path, duration)]
current_segment_duration = duration
else:
# Add to current segment
current_segment.append((timestamp, img_path, duration))
current_segment_duration += duration
# Add the last segment if not empty
if current_segment:
segments.append(current_segment)
print(f"Split into {len(segments)} segments (max duration: {max_duration} seconds)")
# Process each segment
for segment_index, segment in enumerate(segments):
# Generate output filename
if len(segments) > 1:
# Extract base name and extension
base_name, extension = os.path.splitext(output_video_file)
segment_output_file = f"{base_name}_part{segment_index+1}{extension}"
else:
segment_output_file = output_video_file
print(f"\nProcessing segment {segment_index+1}/{len(segments)} -> {segment_output_file}")
# Create FFmpeg input file for this segment
ffmpeg_input_file = f"ffmpeg_input_segment_{segment_index+1}.txt"
try:
with open(ffmpeg_input_file, 'w', encoding='utf-8') as f:
f.write("ffconcat version 1.0\n")
for timestamp, img_path, duration in segment:
abs_img_path = os.path.abspath(img_path)
safe_img_path = abs_img_path.replace("'", "'\\''")
f.write(f"file '{safe_img_path}'\n")
f.write(f"duration {duration:.6f}\n")
# Add the last frame again for final duration to apply
_, last_img_path, _ = segment[-1]
abs_last_img_path = os.path.abspath(last_img_path)
safe_last_img_path = abs_last_img_path.replace("'", "'\\''")
f.write(f"file '{safe_last_img_path}'\n")
# Run FFmpeg command
ffmpeg_command = [
'ffmpeg',
'-f', 'concat',
'-safe', '0',
'-i', ffmpeg_input_file,
'-c:v', 'libvpx-vp9',
'-crf', '30',
'-b:v', '0',
'-pix_fmt', 'yuv420p',
'-y',
segment_output_file
]
print("\nRunning FFmpeg command:")
print(" ".join(f"'{arg}'" if " " in arg else arg for arg in ffmpeg_command))
try:
result = subprocess.run(ffmpeg_command, check=True, capture_output=True, text=True, encoding='utf-8')
print("\nFFmpeg output:")
print(result.stdout)
print(f"Segment {segment_index+1} finished successfully!")
print(f"Output video saved to: {segment_output_file}")
except subprocess.CalledProcessError as e:
print(f"\nError running FFmpeg for segment {segment_index+1}!")
print(f"Return code: {e.returncode}")
print("FFmpeg stdout:")
print(e.stdout)
print("FFmpeg stderr:")
print(e.stderr)
except FileNotFoundError:
print("\nError: 'ffmpeg' command not found. Make sure FFmpeg is installed and in your system's PATH.")
finally:
# Clean up the temporary input file
if os.path.exists(ffmpeg_input_file):
os.remove(ffmpeg_input_file)
print(f"Cleaned up temporary file: {ffmpeg_input_file}")
def process_zip_files():
"""
Processes all zip files in the current directory, extracts them,
and creates videos from the extracted images.
"""
# 获取当前目录下所有的zip文件
zip_files = glob.glob('*.trace.zip')
if not zip_files:
print("没有找到任何.trace.zip文件")
return
print(f"找到{len(zip_files)}个zip文件")
# 创建video目录如果不存在
video_base_dir = os.path.join(os.getcwd(), 'video')
if not os.path.exists(video_base_dir):
os.makedirs(video_base_dir)
# 处理每个zip文件
for zip_file in zip_files:
# 获取不带.zip扩展名的文件名
base_name = zip_file[:-4] # 移除.zip
extract_dir = os.path.join(os.getcwd(), base_name)
video_output_dir = os.path.join(video_base_dir, base_name)
print(f"\n处理zip文件: {zip_file}")
# 如果解压目录已存在,先删除
if os.path.exists(extract_dir):
print(f"删除已存在的目录: {extract_dir}")
shutil.rmtree(extract_dir)
# 创建解压目录
os.makedirs(extract_dir)
# 创建视频输出目录
if not os.path.exists(video_output_dir):
os.makedirs(video_output_dir)
# 解压文件
try:
print(f"解压文件到: {extract_dir}")
with zipfile.ZipFile(zip_file, 'r') as zip_ref:
zip_ref.extractall(extract_dir)
# 进入解压目录并处理图像
print(f"进入目录: {extract_dir}")
output_video = os.path.join(video_output_dir, f"{base_name}_recording.webm")
image_dir = extract_dir + "\\resources\\"
# 调用视频创建函数
create_video_from_images(image_dir, output_video, max_duration=30.0)
except Exception as e:
print(f"处理{zip_file}时出错: {str(e)}")
# --- 主程序入口 ---
if __name__ == "__main__":
process_zip_files()