| | |
|---|
| | | |
|---|
| | | |
|---|
| | | class VideoWriter: |
|---|
| | | def __init__(self, path, frame_rate, bit_rate=1000000): |
|---|
| | | def __init__(self, path, frame_rate, bit_rate=1000000,audio_source_path=None): |
|---|
| | | """ |
|---|
| | | :param path: 输出文件路径 |
|---|
| | | :param frame_rate: 视频帧率 |
|---|
| | | :param bit_rate: 视频码率 |
|---|
| | | :param audio_source_path: 原视频文件路径（用于拷贝音频） |
|---|
| | | """ |
|---|
| | | self.container = av.open(path, mode='w') |
|---|
| | | self.stream = self.container.add_stream('h264', rate=round(frame_rate)) |
|---|
| | | self.stream.pix_fmt = 'yuv420p' |
|---|
| | | self.stream.bit_rate = bit_rate |
|---|
| | | self.stream.codec_context.options = { |
|---|
| | | 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.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_source_path =audio_source_path |
|---|
| | | 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): |
|---|
| | | # frames: [T, C, H, W] |
|---|
| | | self.stream.width = frames.size(3) |
|---|
| | | self.stream.height = frames.size(2) |
|---|
| | | self.v_stream.width = frames.size(3) |
|---|
| | | self.v_stream.height = frames.size(2) |
|---|
| | | if frames.size(1) == 1: |
|---|
| | | frames = frames.repeat(1, 3, 1, 1) # convert grayscale to RGB |
|---|
| | | frames = frames.mul(255).byte().cpu().permute(0, 2, 3, 1).numpy() |
|---|
| | | for t in range(frames.shape[0]): |
|---|
| | | frame = frames[t] |
|---|
| | | frame = av.VideoFrame.from_ndarray(frame, format='rgb24') |
|---|
| | | self.container.mux(self.stream.encode(frame)) |
|---|
| | | # old version |
|---|
| | | #self.container.mux(self.v_stream.encode(frame)) |
|---|
| | | for packet in self.v_stream.encode(frame): |
|---|
| | | self.container.mux(packet) |
|---|
| | | |
|---|
| | | def close(self): |
|---|
| | | self.container.mux(self.stream.encode()) |
|---|
| | | # old version |
|---|
| | | #self.container.mux(self.v_stream.encode()) |
|---|
| | | for packet in self.v_stream.encode(): |
|---|
| | | self.container.mux(packet) |
|---|
| | | |
|---|
| | | 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) |
|---|
| | | |
|---|
| | | self.container.close() |
|---|
| | | if self.audio_in_container: |
|---|
| | | self.audio_in_container.close() |
|---|
| | | |
|---|
| | | |
|---|
| | | class ImageSequenceReader(Dataset): |
|---|
