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