add matanyone-better auqliaty than robustvideomatting

wangzhibo

2025-08-03 56d194a4300610a17cc3c6682aaf63f039e9c4ab

 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（H, 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',     # 影响编码速度，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
@@ -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]
                # 直接用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)

        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 数组，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):