From 56d194a4300610a17cc3c6682aaf63f039e9c4ab Mon Sep 17 00:00:00 2001
From: wangzhibo <wangzhibo@shsening.com>
Date: 星期日, 03 八月 2025 16:11:37 +0800
Subject: [PATCH] add matanyone-better auqliaty than robustvideomatting
---
auto_media_publisher/utils/media_input_output.py | 160 ++++++++++++++++++++++++++++++++++++++++++++++-------
1 files changed, 139 insertions(+), 21 deletions(-)
diff --git a/auto_media_publisher/utils/media_input_output.py b/auto_media_publisher/utils/media_input_output.py
index 0adf1bd..f4a9af6 100644
--- a/auto_media_publisher/utils/media_input_output.py
+++ b/auto_media_publisher/utils/media_input_output.py
@@ -5,12 +5,21 @@
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锛圚, W, C锛�
+ first_frame = np.asarray(self.video[0])
+ self.height, self.width = first_frame.shape[:2]
+
self.transform = transform
@property
@@ -29,23 +38,37 @@
class VideoWriter:
- def __init__(self, path, frame_rate, bit_rate=1000000,audio_source_path=None):
+ 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: 鍘熻棰戞枃浠惰矾寰勶紙鐢ㄤ簬鎷疯礉闊抽锛�
"""
- self.container = av.open(path, mode='w')
+ 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', # 褰卞搷缂栫爜閫熷害锛宮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
@@ -62,39 +85,134 @@
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.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.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]
- frame = av.VideoFrame.from_ndarray(frame, format='rgb24')
- # old version
- #self.container.mux(self.v_stream.encode(frame))
- for packet in self.v_stream.encode(frame):
+ video_frame = av.VideoFrame.from_ndarray(frame, format='rgb24')
+ for packet in self.v_stream.encode(video_frame):
self.container.mux(packet)
-
- def close(self):
- # old version
- #self.container.mux(self.v_stream.encode())
- for packet in self.v_stream.encode():
- 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]
+ # 鐩存帴鐢╪umpy鍐欏叆
+ 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)
-
- self.container.close()
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 鏁扮粍锛宻hape=(H,W,4)锛宒type=uint8
+ RGBA 鏍煎紡锛宎lpha 閫氶亾鍐冲畾閫忔槑搴�
+ """
+ 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):
--
Gitblit v1.9.1