add matanyone-better auqliaty than robustvideomatting

wangzhibo

2025-08-03 56d194a4300610a17cc3c6682aaf63f039e9c4ab

 auto_media_publisher/modifiers/video_matting/robust_video_matting_converter.py

File was renamed from auto_media_publisher/modifiers/video_matting/video_matting_converter.py
@@ -7,7 +7,7 @@
import time
from pathlib import Path
from auto_media_publisher.utils.media_input_output import VideoReader, VideoWriter, ImageSequenceReader, ImageSequenceWriter
from .matting_model.model import MattingNetwork
from .robustvideomatting.model import MattingNetwork
import uuid


@@ -18,7 +18,7 @@
CHECKPOINT_MOBILENETV3 ="./models/rvm_mobilenetv3.pth"


class VideoMattingConverter:
class RobustVideoMattingConverter:
    """
    视频抠像转换器
    封装原来的 convert_video 逻辑，支持 video / png_sequence 输出
@@ -43,7 +43,7 @@
    def video_matting(self,
                input: str,
                bg_video: Optional[str] = None,
                bg_jpg: Optional[str] = None,
                bg_image: Optional[str] = None,
                output_video: Optional[str] = None,
                output_alpha: Optional[str] = None,
                output_fore: Optional[str] = None):
@@ -52,10 +52,11 @@
        if not os.path.exists(input):
            print("input file is NOT exist.")
            return False
        
        if bg_video and not os.path.exists(bg_video):
            print("input background video is NOT exist.")
            return False
        if bg_jpg and not os.path.exists(bg_jpg):
        if bg_image and not os.path.exists(bg_image):
            print("input background image is NOT exist.")
            return False
       
@@ -65,7 +66,7 @@
            downsample_ratio = None,
            output_type = 'video',
            background_video = bg_video,
            background_jpg = bg_jpg,
            background_jpg = bg_image,
            with_audio = True,
            output_composition = output_video,
            output_alpha = output_alpha,
@@ -116,13 +117,13 @@

            if output_composition:
                if with_audio:
                    writer_com = VideoWriter(output_composition, frame_rate, int(output_video_mbps * 1_000_000),input_source)
                    writer_com = VideoWriter(output_composition, frame_rate, source.width,source.height,int(output_video_mbps * 1_000_000),input_source)
                else:
                    writer_com = VideoWriter(output_composition, frame_rate, int(output_video_mbps * 1_000_000))
                    writer_com = VideoWriter(output_composition, frame_rate, source.width,source.height,int(output_video_mbps * 1_000_000))
            if output_alpha:
                writer_pha = VideoWriter(output_alpha, frame_rate, int(output_video_mbps * 1_000_000))
                writer_pha = VideoWriter(output_alpha, frame_rate, source.width,source.height,int(output_video_mbps * 1_000_000))
            if output_foreground:
                writer_fgr = VideoWriter(output_foreground, frame_rate, int(output_video_mbps * 1_000_000))
                writer_fgr = VideoWriter(output_foreground, frame_rate,source.width,source.height, int(output_video_mbps * 1_000_000))

        else:  # png_sequence
            if output_composition:
@@ -141,7 +142,8 @@
                bgr_IMG = Image.open(background_jpg)
                bgr = transforms.ToTensor()(bgr_IMG).to(self.device)
            else:
                bgr = torch.tensor([155, 255, 155], device=self.device, dtype=dtype).div(255).view(1, 1, 3, 1, 1)
                # 使用 0,200,0 而不是 0,255,0 避免过度曝光
                bgr = torch.tensor([0, 200, 0], device=self.device, dtype=dtype).div(255).view(1, 1, 3, 1, 1)   

        # 开始推理
        self.model = self.model.eval()