From c00e9454085ef3e96b98a1ded81aad6a3b95d849 Mon Sep 17 00:00:00 2001
From: wangzhibo <wangzhibo@shsening.com>
Date: 星期四, 31 七月 2025 22:20:08 +0800
Subject: [PATCH] add video matting
---
auto_media_publisher/utils/media_input_output.py | 58 +++++++++++++++++++++++++++++++++++++++++++++++++---------
1 files changed, 49 insertions(+), 9 deletions(-)
diff --git a/auto_media_publisher/utils/media_input_output.py b/auto_media_publisher/utils/media_input_output.py
index 4e2e209..0adf1bd 100644
--- a/auto_media_publisher/utils/media_input_output.py
+++ b/auto_media_publisher/utils/media_input_output.py
@@ -29,32 +29,72 @@
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', # 褰卞搷缂栫爜閫熷害锛宮edium 鏄�1x锛宖ast鏄�2x faster鏄�3x veryfast鏄�4x superfast鏄�6x ultrafast鏄�10x锛屾枃浠朵綋绉ぇ绾︿細澧炲姞(浠edium涓哄熀鍑�) 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):
--
Gitblit v1.9.1