import av
|
import os
|
import pims
|
import numpy as np
|
from torch.utils.data import Dataset
|
from torchvision.transforms.functional import to_pil_image
|
from PIL import Image
|
import imageio.v3 as iio
|
import uuid
|
import traceback
|
import subprocess
|
import time
|
|
class VideoReader(Dataset):
|
def __init__(self, path, transform=None):
|
self.video = pims.PyAVVideoReader(path)
|
self.rate = self.video.frame_rate
|
|
# 直接拿第一帧的 shape（H, W, C）
|
first_frame = np.asarray(self.video[0])
|
self.height, self.width = first_frame.shape[:2]
|
|
self.transform = transform
|
|
@property
|
def frame_rate(self):
|
return self.rate
|
|
def __len__(self):
|
return len(self.video)
|
|
def __getitem__(self, idx):
|
frame = self.video[idx]
|
frame = Image.fromarray(np.asarray(frame))
|
if self.transform is not None:
|
frame = self.transform(frame)
|
return frame
|
|
|
class VideoWriter:
|
def __init__(self, output_path, frame_rate, width,height, bit_rate=1000000,audio_source_path=None):
|
"""
|
:param path: 输出文件路径
|
:param frame_rate: 视频帧率
|
:param bit_rate: 视频码率
|
:param audio_source_path: 原视频文件路径（用于拷贝音频）
|
"""
|
print(f"ready to create video file : {output_path}")
|
self.output_path = output_path
|
self.audio_source_path =audio_source_path
|
if self.audio_source_path and os.path.exists(self.audio_source_path):
|
# using temp video file without audio
|
self.need_audio_action = True
|
self.temp_path = f"{uuid.uuid4()}.mp4"
|
print(self.temp_path)
|
self.container = av.open(self.temp_path, mode='w')
|
else:
|
self.container = av.open(self.output_path, mode='w')
|
self.need_audio_action = False
|
|
self.v_stream = self.container.add_stream('h264', rate=round(frame_rate))
|
self.v_stream.pix_fmt = 'yuv420p'
|
self.v_stream.bit_rate = bit_rate
|
self.v_stream.width = width
|
self.v_stream.height = height
|
self.v_stream.codec_context.options = {
|
'crf': '18', # 控制画质，越小越清晰（18是高质量）
|
'preset': 'veryfast', # 影响编码速度，medium 是1x，fast是2x faster是3x veryfast是4x superfast是6x ultrafast是10x，文件体积大约会增加(以medium为基准) 10%，15%，20% 30% 50%
|
'tune': 'film',
|
}
|
'''
|
self.audio_in_container = None
|
self.audio_stream_in = None
|
self.audio_out_stream = None
|
|
if self.audio_source_path:
|
try:
|
self.audio_in_container = av.open(audio_source_path)
|
if self.audio_in_container.streams.audio:
|
self.audio_stream_in = self.audio_in_container.streams.audio[0]
|
except:
|
pass
|
|
if self.audio_stream_in:
|
self.audio_out_stream = self.container.add_stream(self.audio_stream_in.codec.name)
|
else:
|
self.audio_out_stream = None
|
'''
|
|
|
def write(self, frames):
|
|
if frames.size(1) == 1:
|
frames = frames.repeat(1, 3, 1, 1)
|
|
frames = frames.mul(255).byte().cpu().permute(0, 2, 3, 1).numpy()
|
|
for t in range(frames.shape[0]):
|
frame = frames[t]
|
video_frame = av.VideoFrame.from_ndarray(frame, format='rgb24')
|
for packet in self.v_stream.encode(video_frame):
|
self.container.mux(packet)
|
|
def write_nparray(self,frames):
|
|
try:
|
if frames.ndim == 3: # 单帧
|
frames = np.expand_dims(frames, axis=0)
|
|
if frames.shape[-1] == 1: # 灰度图
|
frames = np.repeat(frames, 3, axis=-1) # HWC → RGB
|
|
for t in range(frames.shape[0]):
|
frame = frames[t]
|
# 直接用numpy写入
|
video_frame = av.VideoFrame.from_ndarray(frame, format='rgb24')
|
for packet in self.v_stream.encode(video_frame):
|
self.container.mux(packet)
|
except Exception as e:
|
print(f"{e}")
|
traceback.format_exc() # 打印堆栈信息（如果 logger 支持 debug）
|
|
|
|
def merge_video_audio(self,video_path, audio_path, output_path):
|
"""
|
用 FFmpeg 把 video_path 的画面 和 audio_path 的音频 合并成 output_path
|
最短的一个结束 (-shortest)，合并后删除 video_path
|
"""
|
try:
|
# 构造 FFmpeg 命令
|
cmd = [
|
"ffmpeg", "-y", # -y 覆盖输出
|
"-loglevel", "error", # 只显示错误
|
"-i", video_path, # 画面
|
"-i", audio_path, # 音频
|
"-c:v", "copy", # 视频不重新编码
|
"-c:a", "copy", # 音频不重新编码
|
"-map", "0:v:0", # 第一个输入的视频
|
"-map", "1:a:0", # 第二个输入的音频
|
"-shortest", # 以最短的输入结束
|
output_path
|
]
|
|
print(f"[INFO] 开始合并视频和音频...\n{' '.join(cmd)}")
|
subprocess.run(cmd, check=True)
|
|
print("[INFO] 合并完成，等待 1 秒后删除原画面视频...")
|
time.sleep(1)
|
|
# 删除原视频
|
if os.path.exists(video_path):
|
os.remove(video_path)
|
print(f"[INFO] 已删除原画面视频: {video_path}")
|
|
except subprocess.CalledProcessError as e:
|
print(f"[ERROR] FFmpeg 合并失败，错误码: {e.returncode}")
|
except Exception as e:
|
print(f"[ERROR] 发生未知错误: {e}")
|
|
|
|
def close(self):
|
'''
|
if self.audio_stream_in and self.audio_out_stream:
|
for packet in self.audio_in_container.demux(self.audio_stream_in):
|
if packet.dts is None:
|
continue
|
packet.stream = self.audio_out_stream # 直接改，不用 copy()
|
self.container.mux(packet)
|
if self.audio_in_container:
|
self.audio_in_container.close()
|
'''
|
self.container.close()
|
|
if self.need_audio_action:
|
self.merge_video_audio(self.temp_path,self.audio_source_path,self.output_path)
|
|
|
|
class WebMWriter:
|
"""
|
支持透明通道的中间视频格式
|
# 示例使用：
|
writer = WebMWriter("output_alpha.webm", fps=30)
|
for i in range(frames):
|
writer.write(frame)
|
writer.close()
|
"""
|
def __init__(self, filename, fps=30):
|
self.filename = filename
|
self.fps = fps
|
self.frames = []
|
|
def write(self, frame: np.ndarray):
|
"""
|
传入的 frame 应该是含 alpha 通道的 numpy 数组，shape=(H,W,4)，dtype=uint8
|
RGBA 格式，alpha 通道决定透明度
|
"""
|
assert frame.shape[2] == 4, "Frame must have 4 channels (RGBA)"
|
self.frames.append(frame)
|
|
def close(self):
|
iio.imwrite(
|
self.filename,
|
self.frames,
|
fps=self.fps,
|
codec="vp8", # vp8支持alpha通道
|
ffmpeg_params=[
|
"-auto-alt-ref", "0", # 关闭参考帧，支持透明
|
"-pix_fmt", "yuva420p" # VP8透明格式
|
],
|
quality=10 # 质量，范围1-63，10比较高质量
|
)
|
print(f"WebM video saved to {self.filename}")
|
|
|
class ImageSequenceReader(Dataset):
|
def __init__(self, path, transform=None):
|
self.path = path
|
self.files = sorted(os.listdir(path))
|
self.transform = transform
|
|
def __len__(self):
|
return len(self.files)
|
|
def __getitem__(self, idx):
|
with Image.open(os.path.join(self.path, self.files[idx])) as img:
|
img.load()
|
if self.transform is not None:
|
return self.transform(img)
|
return img
|
|
|
class ImageSequenceWriter:
|
def __init__(self, path, extension='jpg'):
|
self.path = path
|
self.extension = extension
|
self.counter = 0
|
os.makedirs(path, exist_ok=True)
|
|
def write(self, frames):
|
# frames: [T, C, H, W]
|
for t in range(frames.shape[0]):
|
to_pil_image(frames[t]).save(os.path.join(
|
self.path, str(self.counter).zfill(4) + '.' + self.extension))
|
self.counter += 1
|
|
def close(self):
|
pass
|
|
