wangzhibo
2025-07-31 c00e9454085ef3e96b98a1ded81aad6a3b95d849
add video matting
9个文件已修改
4个文件已删除
12个文件已添加
1858 ■■■■■ 已修改文件
.gitignore 1 ●●●●● 补丁 | 查看 | 原始文档 | blame | 历史
auto_media_publisher/generators/docx_generator_gemini.py 4 ●●●●● 补丁 | 查看 | 原始文档 | blame | 历史
auto_media_publisher/modifiers/face_swap/face_analyser.py 10 ●●●●● 补丁 | 查看 | 原始文档 | blame | 历史
auto_media_publisher/modifiers/face_swap/face_enhancer.py 49 ●●●●● 补丁 | 查看 | 原始文档 | blame | 历史
auto_media_publisher/modifiers/face_swap/face_nsfw_predictor.py 5 ●●●●● 补丁 | 查看 | 原始文档 | blame | 历史
auto_media_publisher/modifiers/face_swap/face_swaper_avpy.py 166 ●●●●● 补丁 | 查看 | 原始文档 | blame | 历史
auto_media_publisher/modifiers/face_swap/face_swaper_ffmpeg.py 238 ●●●●● 补丁 | 查看 | 原始文档 | blame | 历史
auto_media_publisher/modifiers/face_swap/face_swaper_ffmpeg_stream.py 241 ●●●●● 补丁 | 查看 | 原始文档 | blame | 历史
auto_media_publisher/modifiers/face_swap/face_swapper_enhancer.py 248 ●●●●● 补丁 | 查看 | 原始文档 | blame | 历史
auto_media_publisher/modifiers/rewriter/docx_rewriter_gemini.py 2 ●●●●● 补丁 | 查看 | 原始文档 | blame | 历史
auto_media_publisher/modifiers/video_matting/__init__.py 补丁 | 查看 | 原始文档 | blame | 历史
auto_media_publisher/modifiers/video_matting/matting_model/__init__.py 补丁 | 查看 | 原始文档 | blame | 历史
auto_media_publisher/modifiers/video_matting/matting_model/decoder.py 210 ●●●●● 补丁 | 查看 | 原始文档 | blame | 历史
auto_media_publisher/modifiers/video_matting/matting_model/deep_guided_filter.py 61 ●●●●● 补丁 | 查看 | 原始文档 | blame | 历史
auto_media_publisher/modifiers/video_matting/matting_model/fast_guided_filter.py 76 ●●●●● 补丁 | 查看 | 原始文档 | blame | 历史
auto_media_publisher/modifiers/video_matting/matting_model/lraspp.py 29 ●●●●● 补丁 | 查看 | 原始文档 | blame | 历史
auto_media_publisher/modifiers/video_matting/matting_model/mobilenetv3.py 72 ●●●●● 补丁 | 查看 | 原始文档 | blame | 历史
auto_media_publisher/modifiers/video_matting/matting_model/model.py 79 ●●●●● 补丁 | 查看 | 原始文档 | blame | 历史
auto_media_publisher/modifiers/video_matting/matting_model/resnet.py 45 ●●●●● 补丁 | 查看 | 原始文档 | blame | 历史
auto_media_publisher/modifiers/video_matting/video_matting_converter.py 214 ●●●●● 补丁 | 查看 | 原始文档 | blame | 历史
auto_media_publisher/utils/media_input_output.py 58 ●●●●● 补丁 | 查看 | 原始文档 | blame | 历史
data/auto_media_publisher.sqlite3 补丁 | 查看 | 原始文档 | blame | 历史
data/cookies/18918594752/gongzhonghao.json 2 ●●●●● 补丁 | 查看 | 原始文档 | blame | 历史
test/test_face_swap.py 11 ●●●●● 补丁 | 查看 | 原始文档 | blame | 历史
test/test_video_matting.py 37 ●●●●● 补丁 | 查看 | 原始文档 | blame | 历史
.gitignore
@@ -31,3 +31,4 @@
# 其他不太想管关联的
下载19楼论坛/
data/cookies/18918594752/gongzhonghao.json
auto_media_publisher/generators/docx_generator_gemini.py
@@ -67,11 +67,11 @@
    def generate(self, output_path: str):
        self.logger.info(f"start to rewriting docx using gemini,using temperature:{self.temperature}. with prompt {self.prompt}.")
        self.logger.info(f"start to Generating docx using gemini, with prompt {self.prompt}.")
        self.chat = self.model.start_chat(history=[])
        response = self.safe_send_message(self.prompt)
        story_text = response.text.strip()
        self.logger.info(f"google gemini story end...")
        self.logger.info(f"[ google ] gemini story end...")
        paragraphs = [p.strip() for p in story_text.split('\n') if p.strip()]
        doc = Document()
auto_media_publisher/modifiers/face_swap/face_analyser.py
@@ -4,9 +4,7 @@
import numpy
from insightface.app import FaceAnalysis
from insightface.app.common import Face
from gfpgan.utils import GFPGANer
from .typing import Frame # 你已有定义
import numpy as np
SIMILAR_FACE_DISTANCE: float = 0.45     # < 0.4 同一人 0.4 -0.6 大概是同一人 >0.6 不是同一个
@@ -25,14 +23,14 @@
        else:
            return 'cpu'
        
    def get_many_faces(self,frame: Frame) -> Optional[List[Face]]:
    def get_many_faces(self,frame: np.ndarray) -> Optional[List[Face]]:
        try:
            return self._analyser.get(frame)
        except Exception as e:
            print(f"{e}")
            return None
    def get_one_face(self,frame: Frame, position: int = 0) -> Optional[Face]:
    def get_one_face(self,frame: np.ndarray, position: int = 0) -> Optional[Face]:
        many_faces = self.get_many_faces(frame)
        if many_faces:
            try:
@@ -42,7 +40,7 @@
        return None
    def find_similar_face(self,frame: Frame, reference_face: Face) -> Optional[Face]:
    def find_similar_face(self,frame: np.ndarray, reference_face: Face) -> Optional[Face]:
        many_faces = self.get_many_faces(frame)
        if many_faces:
            for face in many_faces:
auto_media_publisher/modifiers/face_swap/face_enhancer.py
File was deleted
auto_media_publisher/modifiers/face_swap/face_nsfw_predictor.py
@@ -1,10 +1,7 @@
import threading
import numpy as np
from PIL import Image
from keras import Model
import opennsfw2
from .typing import Frame  # 你已有定义
MAX_PROBABILITY = 0.85
@@ -12,7 +9,7 @@
    def __init__(self):
        self._predictor = opennsfw2.make_open_nsfw_model()
    def predict_frame(self, frame: Frame) -> bool:
    def predict_frame(self, frame: np.ndarray) -> bool:
        image = Image.fromarray(frame)
        image = opennsfw2.preprocess_image(image, opennsfw2.Preprocessing.YAHOO)
        views = np.expand_dims(image, axis=0)
auto_media_publisher/modifiers/face_swap/face_swaper_avpy.py
File was deleted
auto_media_publisher/modifiers/face_swap/face_swaper_ffmpeg.py
File was deleted
auto_media_publisher/modifiers/face_swap/face_swaper_ffmpeg_stream.py
File was deleted
auto_media_publisher/modifiers/face_swap/face_swapper_enhancer.py
New file
@@ -0,0 +1,248 @@
import av
import torch
import threading, queue
from threading import current_thread
import numpy as np
import subprocess
import json
import cv2
import time
import insightface
from insightface.app.common import Face
from insightface.app import FaceAnalysis
from gfpgan.utils import GFPGANer
from .typing import Frame
from auto_media_publisher.config.conf_base import FACE_ENHANCER_MODEL_PATH,FACE_SWAPPER_MODEL_PATH
from auto_media_publisher.utils.logger import get_logger
import os
from fractions import Fraction
SIMILAR_FACE_DISTANCE: float = 0.45     # < 0.4 同一人 0.4 -0.6 大概是同一人 >0.6 不是同一个
class FaceSwapperEnhancer:
    def __init__(self):
        self.logger = get_logger("FaceSwapperEnhancer","admin")
        self.frame_q = None  #queue.Queue(maxsize=10)
        self.processed_q = None  #queue.Queue(maxsize=10)
        self.device =self._select_device()
        self.with_swapper =False
        self.with_enhancer =False
        self.face_analyser = FaceAnalysis(name='buffalo_l', root='./', providers=['CUDAExecutionProvider'])
        self.face_analyser.prepare(ctx_id=0, det_size=(640, 640))
        self.face_swapper = insightface.model_zoo.get_model(FACE_SWAPPER_MODEL_PATH, providers=['CUDAExecutionProvider'])
        self.face_enhancer = GFPGANer(model_path=FACE_ENHANCER_MODEL_PATH, upscale=1, device=self.device)
    def _select_device(self):
        if torch.cuda.is_available():
            return 'cuda:0'
        elif torch.backends.mps.is_available():
            return 'mps'
        else:
            return 'cpu'
    def enhance_face(self,target_face, temp_frame:np.ndarray) -> np.ndarray:
        start_x, start_y, end_x, end_y = map(int, target_face['bbox'])
        padding_x = int((end_x - start_x) * 0.5)
        padding_y = int((end_y - start_y) * 0.5)
        start_x = max(0, start_x - padding_x)
        start_y = max(0, start_y - padding_y)
        end_x = max(0, end_x + padding_x)
        end_y = max(0, end_y + padding_y)
        temp_face = temp_frame[start_y:end_y, start_x:end_x]
        if temp_face.size:
            _, _, temp_face = self.face_enhancer.enhance(
                    temp_face,
                    paste_back=True
            )
            temp_frame[start_y:end_y, start_x:end_x] = temp_face
        return temp_frame
    def decode_video(self,video_path):
        container = av.open(video_path)
        for idx, frame in enumerate(container.decode(video=0)):
            img = frame.to_ndarray(format='rgb24')
            self.frame_q.put((idx, img))
        self.frame_q.put((None, None))
    def process_frame(self,new_face=None):
        start_time = time.time()
        while True:
            idx, frame = self.frame_q.get()
            if idx is None:
                self.processed_q.put((None, None))
                break
            try:
                many_faces = self.face_analyser.get(frame)
                if many_faces:
                    if self.with_swapper:
                        old_face = many_faces[0]
                        frame = self.face_swapper.get(frame, old_face, new_face, paste_back=True)
                    if self.with_enhancer:
                        for sinlge_face in many_faces:
                            frame = self.enhance_face(sinlge_face, frame)
                self.processed_q.put((idx, frame))
                print(f"{current_thread().name} has proceesed frame :{idx}")
            except Exception as e:
                print(f"[ERROR] Frame {idx} failed: {e}")
        print(f"totol time is : {time.time()- start_time}")
    def probe_video(self,old_video_path):
        cmd = [
            "ffprobe",
            "-v", "error",
            "-select_streams", "v:0",
            "-show_entries", "stream=width,height,r_frame_rate",
            "-of", "json",
            old_video_path,
        ]
        proc = subprocess.run(cmd, capture_output=True, text=True, check=True)
        info = json.loads(proc.stdout)
        stream = info["streams"][0]
        fps = eval(stream["r_frame_rate"])
        fps_fraction = Fraction(fps).limit_denominator()
        return fps_fraction,int(stream["width"]),int(stream["height"])
    def encode_video(self,output_path, fps, width, height,audio_source_path):
        audio_in_container = av.open(audio_source_path)
        audio_stream_in = None
        if audio_in_container.streams.audio:
            audio_stream_in = audio_in_container.streams.audio[0]
        out_container = av.open(output_path, mode='w')
        v_stream = out_container.add_stream('libx264', rate=fps)
        v_stream.width = width
        v_stream.height = height
        v_stream.pix_fmt = 'yuv420p'
        v_stream.options = {'crf': '18', 'preset': 'veryfast'}
        # 添加音频流,基于输入音频流模板
        audio_out_stream = None
        if audio_stream_in:
            audio_out_stream = out_container.add_stream(audio_stream_in.codec.name)
            '''
            audio_codec_name = audio_stream_in.codec.name
            audio_out_stream = out_container.add_stream(
                audio_codec_name,
                rate=audio_stream_in.rate or audio_stream_in.sample_rate
            )
            # 设置一些可写的 metadata(layout)
            try:
                audio_out_stream.layout = audio_stream_in.layout
            except Exception:
                pass
            '''
        else:
            audio_out_stream = None
        next_index = 0
        buffer = {}
        while True:
            idx, frame = self.processed_q.get()
            if idx is None:
                    break
            buffer[idx] = frame
            # 只要有当前要写入的帧,就输出
            while next_index in buffer:
                frame = buffer.pop(next_index)
                video_frame = av.VideoFrame.from_ndarray(frame, format='rgb24')
                for packet in v_stream.encode(video_frame):
                    out_container.mux(packet)
                next_index += 1
        # 清空编码缓存
        for packet in v_stream.encode():
            out_container.mux(packet)
        # 复制原音频流(直接拷贝,不解码)
        # 拷贝音频包,修改流指向
        if audio_stream_in and audio_out_stream:
            for packet in audio_in_container.demux(audio_stream_in):
                if packet.dts is None:
                    continue
                packet.stream = audio_out_stream  # 直接改,不用 copy()
                out_container.mux(packet)
        out_container.close()
        audio_in_container.close()
    def enhance(self,old_video_path,output_video_path):
        try:
            if not os.path.exists(old_video_path):
                self.logger.error(f"[face swap] input video file is not exist. {old_video_path}")
                return False
            self.frame_q = queue.Queue(maxsize=10)
            self.processed_q = queue.Queue(maxsize=10)
            self.with_swapper = False
            self.with_enhancer = True
            self.start_threads(old_video_path,output_video_path,None)
            return True
        except Exception as e:
            self.logger.error(f"[ error ]. {e}")
            return False
    def start_threads(self,old_video_path,output_video_path,new_face=None):
        threads = []
        # 解码线程
        t_decode = threading.Thread(target=self.decode_video,args=(old_video_path,))
        threads.append(t_decode)
        t_decode.start()
        # 处理线程  多处理线程 anaylse  enhancer 可能会出问题,暂时还是单线程
        t_process = threading.Thread(target=self.process_frame, args=(new_face,))
        threads.append(t_process)
        t_process.start()
        fps,width,height =self.probe_video(old_video_path)
        # 编码线程
        t_encode = threading.Thread(target=self.encode_video,args=(output_video_path, fps,width,height,old_video_path))
        threads.append(t_encode)
        t_encode.start()
        for t in threads:
            t.join()
    def swap(self,old_video_path,output_video_path,new_face_image_path,with_enhancer=True):
        try:
            if not os.path.exists(old_video_path):
                self.logger.error(f"[face swap] input video file is not exist. {old_video_path}")
                return False
            if not os.path.exists(new_face_image_path):
                self.logger.error(f"[face swap] input new face image file is not exist. {new_face_image_path}")
                return False
            faces = self.face_analyser.get(cv2.imread(new_face_image_path))
            if not faces:
                self.logger.error(f"[face swap] input new face image file DOES NOT HAVE FACE. {new_face_image_path}")
                return False
            new_face = sorted(faces, key=lambda x: x.bbox[0])[0]
            self.frame_q = queue.Queue(maxsize=10)
            self.processed_q = queue.Queue(maxsize=10)
            self.with_swapper = True
            self.with_enhancer = with_enhancer
            self.start_threads(old_video_path,output_video_path,new_face)
            return True
        except Exception as e:
            self.logger.error(f"[ error ]. {e}")
            return False
auto_media_publisher/modifiers/rewriter/docx_rewriter_gemini.py
@@ -55,7 +55,7 @@
    def rewrite(self, input_path: str, output_path: str):
        blocks = self.extract_docx_blocks_by_order(input_path)
        self.logger.info(f"start to rewriting docx using gemini,using temperature:{self.temperature}. with style {self.rewrite_style}.  {input_path}")
        self.logger.info(f"start to rewriting docx using gemini,. with style {self.rewrite_style}.  {input_path}")
        if len(blocks) >0:
            self.chat = self.model.start_chat(history=[])
            self.safe_send_message(GOOGLE_REWRITE_STYLE_MAP.get(self.rewrite_style, GEMINI_PROMPT_INIT_DEFAULT))
auto_media_publisher/modifiers/video_matting/__init__.py
auto_media_publisher/modifiers/video_matting/matting_model/__init__.py
auto_media_publisher/modifiers/video_matting/matting_model/decoder.py
New file
@@ -0,0 +1,210 @@
import torch
from torch import Tensor
from torch import nn
from torch.nn import functional as F
from typing import Tuple, Optional
class RecurrentDecoder(nn.Module):
    def __init__(self, feature_channels, decoder_channels):
        super().__init__()
        self.avgpool = AvgPool()
        self.decode4 = BottleneckBlock(feature_channels[3])
        self.decode3 = UpsamplingBlock(feature_channels[3], feature_channels[2], 3, decoder_channels[0])
        self.decode2 = UpsamplingBlock(decoder_channels[0], feature_channels[1], 3, decoder_channels[1])
        self.decode1 = UpsamplingBlock(decoder_channels[1], feature_channels[0], 3, decoder_channels[2])
        self.decode0 = OutputBlock(decoder_channels[2], 3, decoder_channels[3])
    def forward(self,
                s0: Tensor, f1: Tensor, f2: Tensor, f3: Tensor, f4: Tensor,
                r1: Optional[Tensor], r2: Optional[Tensor],
                r3: Optional[Tensor], r4: Optional[Tensor]):
        s1, s2, s3 = self.avgpool(s0)
        x4, r4 = self.decode4(f4, r4)
        x3, r3 = self.decode3(x4, f3, s3, r3)
        x2, r2 = self.decode2(x3, f2, s2, r2)
        x1, r1 = self.decode1(x2, f1, s1, r1)
        x0 = self.decode0(x1, s0)
        return x0, r1, r2, r3, r4
class AvgPool(nn.Module):
    def __init__(self):
        super().__init__()
        self.avgpool = nn.AvgPool2d(2, 2, count_include_pad=False, ceil_mode=True)
    def forward_single_frame(self, s0):
        s1 = self.avgpool(s0)
        s2 = self.avgpool(s1)
        s3 = self.avgpool(s2)
        return s1, s2, s3
    def forward_time_series(self, s0):
        B, T = s0.shape[:2]
        s0 = s0.flatten(0, 1)
        s1, s2, s3 = self.forward_single_frame(s0)
        s1 = s1.unflatten(0, (B, T))
        s2 = s2.unflatten(0, (B, T))
        s3 = s3.unflatten(0, (B, T))
        return s1, s2, s3
    def forward(self, s0):
        if s0.ndim == 5:
            return self.forward_time_series(s0)
        else:
            return self.forward_single_frame(s0)
class BottleneckBlock(nn.Module):
    def __init__(self, channels):
        super().__init__()
        self.channels = channels
        self.gru = ConvGRU(channels // 2)
    def forward(self, x, r: Optional[Tensor]):
        a, b = x.split(self.channels // 2, dim=-3)
        b, r = self.gru(b, r)
        x = torch.cat([a, b], dim=-3)
        return x, r
class UpsamplingBlock(nn.Module):
    def __init__(self, in_channels, skip_channels, src_channels, out_channels):
        super().__init__()
        self.out_channels = out_channels
        self.upsample = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False)
        self.conv = nn.Sequential(
            nn.Conv2d(in_channels + skip_channels + src_channels, out_channels, 3, 1, 1, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(True),
        )
        self.gru = ConvGRU(out_channels // 2)
    def forward_single_frame(self, x, f, s, r: Optional[Tensor]):
        x = self.upsample(x)
        x = x[:, :, :s.size(2), :s.size(3)]
        x = torch.cat([x, f, s], dim=1)
        x = self.conv(x)
        a, b = x.split(self.out_channels // 2, dim=1)
        b, r = self.gru(b, r)
        x = torch.cat([a, b], dim=1)
        return x, r
    def forward_time_series(self, x, f, s, r: Optional[Tensor]):
        B, T, _, H, W = s.shape
        x = x.flatten(0, 1)
        f = f.flatten(0, 1)
        s = s.flatten(0, 1)
        x = self.upsample(x)
        x = x[:, :, :H, :W]
        x = torch.cat([x, f, s], dim=1)
        x = self.conv(x)
        x = x.unflatten(0, (B, T))
        a, b = x.split(self.out_channels // 2, dim=2)
        b, r = self.gru(b, r)
        x = torch.cat([a, b], dim=2)
        return x, r
    def forward(self, x, f, s, r: Optional[Tensor]):
        if x.ndim == 5:
            return self.forward_time_series(x, f, s, r)
        else:
            return self.forward_single_frame(x, f, s, r)
class OutputBlock(nn.Module):
    def __init__(self, in_channels, src_channels, out_channels):
        super().__init__()
        self.upsample = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False)
        self.conv = nn.Sequential(
            nn.Conv2d(in_channels + src_channels, out_channels, 3, 1, 1, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(True),
            nn.Conv2d(out_channels, out_channels, 3, 1, 1, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(True),
        )
    def forward_single_frame(self, x, s):
        x = self.upsample(x)
        x = x[:, :, :s.size(2), :s.size(3)]
        x = torch.cat([x, s], dim=1)
        x = self.conv(x)
        return x
    def forward_time_series(self, x, s):
        B, T, _, H, W = s.shape
        x = x.flatten(0, 1)
        s = s.flatten(0, 1)
        x = self.upsample(x)
        x = x[:, :, :H, :W]
        x = torch.cat([x, s], dim=1)
        x = self.conv(x)
        x = x.unflatten(0, (B, T))
        return x
    def forward(self, x, s):
        if x.ndim == 5:
            return self.forward_time_series(x, s)
        else:
            return self.forward_single_frame(x, s)
class ConvGRU(nn.Module):
    def __init__(self,
                 channels: int,
                 kernel_size: int = 3,
                 padding: int = 1):
        super().__init__()
        self.channels = channels
        self.ih = nn.Sequential(
            nn.Conv2d(channels * 2, channels * 2, kernel_size, padding=padding),
            nn.Sigmoid()
        )
        self.hh = nn.Sequential(
            nn.Conv2d(channels * 2, channels, kernel_size, padding=padding),
            nn.Tanh()
        )
    def forward_single_frame(self, x, h):
        r, z = self.ih(torch.cat([x, h], dim=1)).split(self.channels, dim=1)
        c = self.hh(torch.cat([x, r * h], dim=1))
        h = (1 - z) * h + z * c
        return h, h
    def forward_time_series(self, x, h):
        o = []
        for xt in x.unbind(dim=1):
            ot, h = self.forward_single_frame(xt, h)
            o.append(ot)
        o = torch.stack(o, dim=1)
        return o, h
    def forward(self, x, h: Optional[Tensor]):
        if h is None:
            h = torch.zeros((x.size(0), x.size(-3), x.size(-2), x.size(-1)),
                            device=x.device, dtype=x.dtype)
        if x.ndim == 5:
            return self.forward_time_series(x, h)
        else:
            return self.forward_single_frame(x, h)
class Projection(nn.Module):
    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, 1)
    def forward_single_frame(self, x):
        return self.conv(x)
    def forward_time_series(self, x):
        B, T = x.shape[:2]
        return self.conv(x.flatten(0, 1)).unflatten(0, (B, T))
    def forward(self, x):
        if x.ndim == 5:
            return self.forward_time_series(x)
        else:
            return self.forward_single_frame(x)
auto_media_publisher/modifiers/video_matting/matting_model/deep_guided_filter.py
New file
@@ -0,0 +1,61 @@
import torch
from torch import nn
from torch.nn import functional as F
"""
Adopted from <https://github.com/wuhuikai/DeepGuidedFilter/>
"""
class DeepGuidedFilterRefiner(nn.Module):
    def __init__(self, hid_channels=16):
        super().__init__()
        self.box_filter = nn.Conv2d(4, 4, kernel_size=3, padding=1, bias=False, groups=4)
        self.box_filter.weight.data[...] = 1 / 9
        self.conv = nn.Sequential(
            nn.Conv2d(4 * 2 + hid_channels, hid_channels, kernel_size=1, bias=False),
            nn.BatchNorm2d(hid_channels),
            nn.ReLU(True),
            nn.Conv2d(hid_channels, hid_channels, kernel_size=1, bias=False),
            nn.BatchNorm2d(hid_channels),
            nn.ReLU(True),
            nn.Conv2d(hid_channels, 4, kernel_size=1, bias=True)
        )
    def forward_single_frame(self, fine_src, base_src, base_fgr, base_pha, base_hid):
        fine_x = torch.cat([fine_src, fine_src.mean(1, keepdim=True)], dim=1)
        base_x = torch.cat([base_src, base_src.mean(1, keepdim=True)], dim=1)
        base_y = torch.cat([base_fgr, base_pha], dim=1)
        mean_x = self.box_filter(base_x)
        mean_y = self.box_filter(base_y)
        cov_xy = self.box_filter(base_x * base_y) - mean_x * mean_y
        var_x  = self.box_filter(base_x * base_x) - mean_x * mean_x
        A = self.conv(torch.cat([cov_xy, var_x, base_hid], dim=1))
        b = mean_y - A * mean_x
        H, W = fine_src.shape[2:]
        A = F.interpolate(A, (H, W), mode='bilinear', align_corners=False)
        b = F.interpolate(b, (H, W), mode='bilinear', align_corners=False)
        out = A * fine_x + b
        fgr, pha = out.split([3, 1], dim=1)
        return fgr, pha
    def forward_time_series(self, fine_src, base_src, base_fgr, base_pha, base_hid):
        B, T = fine_src.shape[:2]
        fgr, pha = self.forward_single_frame(
            fine_src.flatten(0, 1),
            base_src.flatten(0, 1),
            base_fgr.flatten(0, 1),
            base_pha.flatten(0, 1),
            base_hid.flatten(0, 1))
        fgr = fgr.unflatten(0, (B, T))
        pha = pha.unflatten(0, (B, T))
        return fgr, pha
    def forward(self, fine_src, base_src, base_fgr, base_pha, base_hid):
        if fine_src.ndim == 5:
            return self.forward_time_series(fine_src, base_src, base_fgr, base_pha, base_hid)
        else:
            return self.forward_single_frame(fine_src, base_src, base_fgr, base_pha, base_hid)
auto_media_publisher/modifiers/video_matting/matting_model/fast_guided_filter.py
New file
@@ -0,0 +1,76 @@
import torch
from torch import nn
from torch.nn import functional as F
"""
Adopted from <https://github.com/wuhuikai/DeepGuidedFilter/>
"""
class FastGuidedFilterRefiner(nn.Module):
    def __init__(self, *args, **kwargs):
        super().__init__()
        self.guilded_filter = FastGuidedFilter(1)
    def forward_single_frame(self, fine_src, base_src, base_fgr, base_pha):
        fine_src_gray = fine_src.mean(1, keepdim=True)
        base_src_gray = base_src.mean(1, keepdim=True)
        fgr, pha = self.guilded_filter(
            torch.cat([base_src, base_src_gray], dim=1),
            torch.cat([base_fgr, base_pha], dim=1),
            torch.cat([fine_src, fine_src_gray], dim=1)).split([3, 1], dim=1)
        return fgr, pha
    def forward_time_series(self, fine_src, base_src, base_fgr, base_pha):
        B, T = fine_src.shape[:2]
        fgr, pha = self.forward_single_frame(
            fine_src.flatten(0, 1),
            base_src.flatten(0, 1),
            base_fgr.flatten(0, 1),
            base_pha.flatten(0, 1))
        fgr = fgr.unflatten(0, (B, T))
        pha = pha.unflatten(0, (B, T))
        return fgr, pha
    def forward(self, fine_src, base_src, base_fgr, base_pha, base_hid):
        if fine_src.ndim == 5:
            return self.forward_time_series(fine_src, base_src, base_fgr, base_pha)
        else:
            return self.forward_single_frame(fine_src, base_src, base_fgr, base_pha)
class FastGuidedFilter(nn.Module):
    def __init__(self, r: int, eps: float = 1e-5):
        super().__init__()
        self.r = r
        self.eps = eps
        self.boxfilter = BoxFilter(r)
    def forward(self, lr_x, lr_y, hr_x):
        mean_x = self.boxfilter(lr_x)
        mean_y = self.boxfilter(lr_y)
        cov_xy = self.boxfilter(lr_x * lr_y) - mean_x * mean_y
        var_x = self.boxfilter(lr_x * lr_x) - mean_x * mean_x
        A = cov_xy / (var_x + self.eps)
        b = mean_y - A * mean_x
        A = F.interpolate(A, hr_x.shape[2:], mode='bilinear', align_corners=False)
        b = F.interpolate(b, hr_x.shape[2:], mode='bilinear', align_corners=False)
        return A * hr_x + b
class BoxFilter(nn.Module):
    def __init__(self, r):
        super(BoxFilter, self).__init__()
        self.r = r
    def forward(self, x):
        # Note: The original implementation at <https://github.com/wuhuikai/DeepGuidedFilter/>
        #       uses faster box blur. However, it may not be friendly for ONNX export.
        #       We are switching to use simple convolution for box blur.
        kernel_size = 2 * self.r + 1
        kernel_x = torch.full((x.data.shape[1], 1, 1, kernel_size), 1 / kernel_size, device=x.device, dtype=x.dtype)
        kernel_y = torch.full((x.data.shape[1], 1, kernel_size, 1), 1 / kernel_size, device=x.device, dtype=x.dtype)
        x = F.conv2d(x, kernel_x, padding=(0, self.r), groups=x.data.shape[1])
        x = F.conv2d(x, kernel_y, padding=(self.r, 0), groups=x.data.shape[1])
        return x
auto_media_publisher/modifiers/video_matting/matting_model/lraspp.py
New file
@@ -0,0 +1,29 @@
from torch import nn
class LRASPP(nn.Module):
    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.aspp1 = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, 1, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(True)
        )
        self.aspp2 = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),
            nn.Conv2d(in_channels, out_channels, 1, bias=False),
            nn.Sigmoid()
        )
    def forward_single_frame(self, x):
        return self.aspp1(x) * self.aspp2(x)
    def forward_time_series(self, x):
        B, T = x.shape[:2]
        x = self.forward_single_frame(x.flatten(0, 1)).unflatten(0, (B, T))
        return x
    def forward(self, x):
        if x.ndim == 5:
            return self.forward_time_series(x)
        else:
            return self.forward_single_frame(x)
auto_media_publisher/modifiers/video_matting/matting_model/mobilenetv3.py
New file
@@ -0,0 +1,72 @@
import torch
from torch import nn
from torchvision.models.mobilenetv3 import MobileNetV3, InvertedResidualConfig
from torchvision.transforms.functional import normalize
class MobileNetV3LargeEncoder(MobileNetV3):
    def __init__(self, pretrained: bool = False):
        super().__init__(
            inverted_residual_setting=[
                InvertedResidualConfig( 16, 3,  16,  16, False, "RE", 1, 1, 1),
                InvertedResidualConfig( 16, 3,  64,  24, False, "RE", 2, 1, 1),  # C1
                InvertedResidualConfig( 24, 3,  72,  24, False, "RE", 1, 1, 1),
                InvertedResidualConfig( 24, 5,  72,  40,  True, "RE", 2, 1, 1),  # C2
                InvertedResidualConfig( 40, 5, 120,  40,  True, "RE", 1, 1, 1),
                InvertedResidualConfig( 40, 5, 120,  40,  True, "RE", 1, 1, 1),
                InvertedResidualConfig( 40, 3, 240,  80, False, "HS", 2, 1, 1),  # C3
                InvertedResidualConfig( 80, 3, 200,  80, False, "HS", 1, 1, 1),
                InvertedResidualConfig( 80, 3, 184,  80, False, "HS", 1, 1, 1),
                InvertedResidualConfig( 80, 3, 184,  80, False, "HS", 1, 1, 1),
                InvertedResidualConfig( 80, 3, 480, 112,  True, "HS", 1, 1, 1),
                InvertedResidualConfig(112, 3, 672, 112,  True, "HS", 1, 1, 1),
                InvertedResidualConfig(112, 5, 672, 160,  True, "HS", 2, 2, 1),  # C4
                InvertedResidualConfig(160, 5, 960, 160,  True, "HS", 1, 2, 1),
                InvertedResidualConfig(160, 5, 960, 160,  True, "HS", 1, 2, 1),
            ],
            last_channel=1280
        )
        if pretrained:
            self.load_state_dict(torch.hub.load_state_dict_from_url(
                'https://download.pytorch.org/models/mobilenet_v3_large-8738ca79.pth'))
        del self.avgpool
        del self.classifier
    def forward_single_frame(self, x):
        x = normalize(x, [0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        x = self.features[0](x)
        x = self.features[1](x)
        f1 = x
        x = self.features[2](x)
        x = self.features[3](x)
        f2 = x
        x = self.features[4](x)
        x = self.features[5](x)
        x = self.features[6](x)
        f3 = x
        x = self.features[7](x)
        x = self.features[8](x)
        x = self.features[9](x)
        x = self.features[10](x)
        x = self.features[11](x)
        x = self.features[12](x)
        x = self.features[13](x)
        x = self.features[14](x)
        x = self.features[15](x)
        x = self.features[16](x)
        f4 = x
        return [f1, f2, f3, f4]
    def forward_time_series(self, x):
        B, T = x.shape[:2]
        features = self.forward_single_frame(x.flatten(0, 1))
        features = [f.unflatten(0, (B, T)) for f in features]
        return features
    def forward(self, x):
        if x.ndim == 5:
            return self.forward_time_series(x)
        else:
            return self.forward_single_frame(x)
auto_media_publisher/modifiers/video_matting/matting_model/model.py
New file
@@ -0,0 +1,79 @@
import torch
from torch import Tensor
from torch import nn
from torch.nn import functional as F
from typing import Optional, List
from .mobilenetv3 import MobileNetV3LargeEncoder
from .resnet import ResNet50Encoder
from .lraspp import LRASPP
from .decoder import RecurrentDecoder, Projection
from .fast_guided_filter import FastGuidedFilterRefiner
from .deep_guided_filter import DeepGuidedFilterRefiner
class MattingNetwork(nn.Module):
    def __init__(self,
                 variant: str = 'mobilenetv3',
                 refiner: str = 'deep_guided_filter',
                 pretrained_backbone: bool = False):
        super().__init__()
        assert variant in ['mobilenetv3', 'resnet50']
        assert refiner in ['fast_guided_filter', 'deep_guided_filter']
        if variant == 'mobilenetv3':
            self.backbone = MobileNetV3LargeEncoder(pretrained_backbone)
            self.aspp = LRASPP(960, 128)
            self.decoder = RecurrentDecoder([16, 24, 40, 128], [80, 40, 32, 16])
        else:
            self.backbone = ResNet50Encoder(pretrained_backbone)
            self.aspp = LRASPP(2048, 256)
            self.decoder = RecurrentDecoder([64, 256, 512, 256], [128, 64, 32, 16])
        self.project_mat = Projection(16, 4)
        self.project_seg = Projection(16, 1)
        if refiner == 'deep_guided_filter':
            self.refiner = DeepGuidedFilterRefiner()
        else:
            self.refiner = FastGuidedFilterRefiner()
    def forward(self,
                src: Tensor,
                r1: Optional[Tensor] = None,
                r2: Optional[Tensor] = None,
                r3: Optional[Tensor] = None,
                r4: Optional[Tensor] = None,
                downsample_ratio: float = 1,
                segmentation_pass: bool = False):
        if downsample_ratio != 1:
            src_sm = self._interpolate(src, scale_factor=downsample_ratio)
        else:
            src_sm = src
        f1, f2, f3, f4 = self.backbone(src_sm)
        f4 = self.aspp(f4)
        hid, *rec = self.decoder(src_sm, f1, f2, f3, f4, r1, r2, r3, r4)
        if not segmentation_pass:
            fgr_residual, pha = self.project_mat(hid).split([3, 1], dim=-3)
            if downsample_ratio != 1:
                fgr_residual, pha = self.refiner(src, src_sm, fgr_residual, pha, hid)
            fgr = fgr_residual + src
            fgr = fgr.clamp(0., 1.)
            pha = pha.clamp(0., 1.)
            return [fgr, pha, *rec]
        else:
            seg = self.project_seg(hid)
            return [seg, *rec]
    def _interpolate(self, x: Tensor, scale_factor: float):
        if x.ndim == 5:
            B, T = x.shape[:2]
            x = F.interpolate(x.flatten(0, 1), scale_factor=scale_factor,
                mode='bilinear', align_corners=False, recompute_scale_factor=False)
            x = x.unflatten(0, (B, T))
        else:
            x = F.interpolate(x, scale_factor=scale_factor,
                mode='bilinear', align_corners=False, recompute_scale_factor=False)
        return x
auto_media_publisher/modifiers/video_matting/matting_model/resnet.py
New file
@@ -0,0 +1,45 @@
import torch
from torch import nn
from torchvision.models.resnet import ResNet, Bottleneck
class ResNet50Encoder(ResNet):
    def __init__(self, pretrained: bool = False):
        super().__init__(
            block=Bottleneck,
            layers=[3, 4, 6, 3],
            replace_stride_with_dilation=[False, False, True],
            norm_layer=None)
        if pretrained:
            self.load_state_dict(torch.hub.load_state_dict_from_url(
                'https://download.pytorch.org/models/resnet50-0676ba61.pth'))
        del self.avgpool
        del self.fc
    def forward_single_frame(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        f1 = x  # 1/2
        x = self.maxpool(x)
        x = self.layer1(x)
        f2 = x  # 1/4
        x = self.layer2(x)
        f3 = x  # 1/8
        x = self.layer3(x)
        x = self.layer4(x)
        f4 = x  # 1/16
        return [f1, f2, f3, f4]
    def forward_time_series(self, x):
        B, T = x.shape[:2]
        features = self.forward_single_frame(x.flatten(0, 1))
        features = [f.unflatten(0, (B, T)) for f in features]
        return features
    def forward(self, x):
        if x.ndim == 5:
            return self.forward_time_series(x)
        else:
            return self.forward_single_frame(x)
auto_media_publisher/modifiers/video_matting/video_matting_converter.py
New file
@@ -0,0 +1,214 @@
import os
import torch
from torch.utils.data import DataLoader
from torchvision import transforms
from typing import Optional, Tuple
from PIL import Image
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
import uuid
VARIANT_RESNET50 ="resnet50"
CHECKPOINT_RESNET50 ="./models/rvm_resnet50.pth"
VARIANT_MOBILENETV3 ="mobilenetv3"
CHECKPOINT_MOBILENETV3 ="./models/rvm_mobilenetv3.pth"
class VideoMattingConverter:
    """
    视频抠像转换器
    封装原来的 convert_video 逻辑,支持 video / png_sequence 输出
    """
    def __init__(self, precision:str = 'high', device: str = "cuda"):
        self.device = device
        if precision=='high':
            self.model = MattingNetwork(VARIANT_RESNET50).eval().to(device)
            self.model.load_state_dict(torch.load(CHECKPOINT_RESNET50, map_location=device))
        else:
            self.model = MattingNetwork(VARIANT_MOBILENETV3).eval().to(device)
            self.model.load_state_dict(torch.load(CHECKPOINT_MOBILENETV3, map_location=device))
        self.model = torch.jit.script(self.model)
        self.model = torch.jit.freeze(self.model)
    @staticmethod
    def auto_downsample_ratio(h: int, w: int) -> float:
        """自动计算 downsample_ratio"""
        return min(512 / max(h, w), 1)
    def video_matting(self,
                input: str,
                bg_video: Optional[str] = None,
                bg_jpg: Optional[str] = None,
                output_video: Optional[str] = None,
                output_alpha: Optional[str] = None,
                output_fore: Optional[str] = None):
        assert any([output_video, output_alpha, output_fore]), 'Must provide at least one output.'
        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):
            print("input background image is NOT exist.")
            return False
        return self._video_matting_core(
            input_source = input,
            input_resize = None,
            downsample_ratio = None,
            output_type = 'video',
            background_video = bg_video,
            background_jpg = bg_jpg,
            with_audio = True,
            output_composition = output_video,
            output_alpha = output_alpha,
            output_foreground = output_fore,
            output_video_mbps = None,
            dtype =torch.float32)
    def _video_matting_core(self,
                input_source: str,
                input_resize: Optional[Tuple[int, int]] = None,
                downsample_ratio: Optional[float] = None,
                output_type: str = 'video',
                background_video: Optional[str] = None,
                background_jpg: Optional[str] = None,
                with_audio: bool = True,
                output_composition: Optional[str] = None,
                output_alpha: Optional[str] = None,
                output_foreground: Optional[str] = None,
                output_video_mbps: Optional[float] = None,
                dtype: Optional[torch.dtype] = torch.float32):
        assert downsample_ratio is None or (0 < downsample_ratio <= 1), 'Downsample ratio must be between 0 and 1.'
        assert any([output_composition, output_alpha, output_foreground]), 'Must provide at least one output.'
        assert output_type in ['video', 'png_sequence'], 'Only support video or png_sequence.'
        success = False
        # 初始化 transform
        if input_resize is not None:
            transform = transforms.Compose([
                transforms.Resize(input_resize[::-1]),
                transforms.ToTensor()
            ])
        else:
            transform = transforms.ToTensor()
        # 读取输入源
        if os.path.isfile(input_source):
            source = VideoReader(input_source, transform)
        else:
            source = ImageSequenceReader(input_source, transform)
        reader = DataLoader(source, batch_size=1, pin_memory=True, num_workers=0)
        # 初始化输出
        if output_type == 'video':
            frame_rate = source.frame_rate if isinstance(source, VideoReader) else 30
            output_video_mbps = 1 if output_video_mbps is None else output_video_mbps
            if output_composition:
                if with_audio:
                    writer_com = VideoWriter(output_composition, frame_rate, int(output_video_mbps * 1_000_000),input_source)
                else:
                    writer_com = VideoWriter(output_composition, frame_rate, int(output_video_mbps * 1_000_000))
            if output_alpha:
                writer_pha = VideoWriter(output_alpha, frame_rate, int(output_video_mbps * 1_000_000))
            if output_foreground:
                writer_fgr = VideoWriter(output_foreground, frame_rate, int(output_video_mbps * 1_000_000))
        else:  # png_sequence
            if output_composition:
                writer_com = ImageSequenceWriter(output_composition, 'png')
            if output_alpha:
                writer_pha = ImageSequenceWriter(output_alpha, 'png')
            if output_foreground:
                writer_fgr = ImageSequenceWriter(output_foreground, 'png')
        # 背景设置
        background_video_reader = None
        if output_composition and output_type == 'video':
            if background_video:
                background_video_reader = VideoReader(background_video, transform=transforms.ToTensor())
            elif background_jpg:
                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)
        # 开始推理
        self.model = self.model.eval()
        try:
            with torch.no_grad():
                rec = [None] * 4
                current_bg_idx = 1
                bg_increasing = True
                for src in reader:
                    if downsample_ratio is None:
                        downsample_ratio = self.auto_downsample_ratio(*src.shape[2:])
                    src = src.to(self.device, dtype, non_blocking=True).unsqueeze(0)
                    fgr, pha, *rec = self.model(src, *rec, downsample_ratio)
                    # 前景
                    if output_foreground:
                        writer_fgr.write(fgr[0])
                    if output_alpha:
                        writer_pha.write(pha[0])
                    # 合成图
                    if output_composition:
                        if output_type == 'video':
                            if background_video:
                                tmp_bgr = background_video_reader[current_bg_idx].to(self.device)
                                if bg_increasing:
                                    current_bg_idx += 1
                                    if current_bg_idx >= len(background_video_reader):
                                        current_bg_idx = len(background_video_reader) - 2
                                        bg_increasing = False
                                else:
                                    current_bg_idx -= 1
                                    if current_bg_idx < 1:
                                        current_bg_idx = 1
                                        bg_increasing = True
                                com = fgr * pha + tmp_bgr * (1 - pha)
                            else:
                                com = fgr * pha + bgr * (1 - pha)
                        else:
                            fgr = fgr * pha.gt(0)
                            com = torch.cat([fgr, pha], dim=-3)
                        writer_com.write(com[0])
            # process finished.
            success = True
        finally:
            # Clean up
            if output_composition is not None:
                writer_com.close()
                '''
                time.sleep(0.5)  # 确保文件释放,Windows 上非常必要
                if success and with_audio and os.path.exists(output_composition) :
                    out_dir = os.path.dirname(output_composition)
                    base_name = os.path.basename(output_composition)
                    audio_output_path = os.path.join(out_dir, f"audio_{base_name}")
                    os.system(f'ffmpeg -i {output_composition} -i {input_source} -c copy '
                              f'-map 0:v:0 -map 1:a:0 {audio_output_path}')
                '''
            if output_alpha:
                writer_pha.close()
            if output_foreground:
                writer_fgr.close()
        return success
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',     # 影响编码速度,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):
data/auto_media_publisher.sqlite3
Binary files differ
data/cookies/18918594752/gongzhonghao.json
@@ -1 +1 @@
{"cookies": [{"name": "ua_id", "value": "Yp8GoVBn7Tj0PWR7AAAAADaCnYPYpvql9WPFzFkbL_A=", "domain": "mp.weixin.qq.com", "path": "/", "expires": 1788443924.066288, "httpOnly": true, "secure": true, "sameSite": "Lax"}, {"name": "wxuin", "value": "53845090997782", "domain": "mp.weixin.qq.com", "path": "/", "expires": 1788405090.13786, "httpOnly": false, "secure": false, "sameSite": "Lax"}, {"name": "_clck", "value": "1oxp806|1|fy1|0", "domain": ".qq.com", "path": "/", "expires": 1785381090, "httpOnly": false, "secure": false, "sameSite": "Lax"}, {"name": "uuid", "value": "8093d628ddf047d5a340ec437ca7fd08", "domain": "mp.weixin.qq.com", "path": "/", "expires": -1, "httpOnly": true, "secure": true, "sameSite": "Lax"}, {"name": "rand_info", "value": "CAESIEjp7aANF1MLCho4lIEi6Ju8oYLEexqCUxELmbqpfG3Y", "domain": "mp.weixin.qq.com", "path": "/", "expires": 1754229560, "httpOnly": true, "secure": true, "sameSite": "Lax"}, {"name": "slave_bizuin", "value": "3932822909", "domain": "mp.weixin.qq.com", "path": "/", "expires": 1754229560, "httpOnly": true, "secure": true, "sameSite": "Lax"}, {"name": "data_bizuin", "value": "3932822909", "domain": "mp.weixin.qq.com", "path": "/", "expires": 1754229560, "httpOnly": true, "secure": true, "sameSite": "Lax"}, {"name": "bizuin", "value": "3932822909", "domain": "mp.weixin.qq.com", "path": "/", "expires": 1754229560, "httpOnly": true, "secure": true, "sameSite": "Lax"}, {"name": "data_ticket", "value": "TXXMQ0chAAhF8D9olYk9AhfbDAjnS2jHUUO8GHs0e45nt6P/zjA38GlY5q12SD0E", "domain": "mp.weixin.qq.com", "path": "/", "expires": 1754229560, "httpOnly": true, "secure": true, "sameSite": "Lax"}, {"name": "slave_sid", "value": "OHl1NHRHZjd2N3poQTFDQ2d5WVBBWnJ5bm9obkZNVjZGank5S2dxSUMwblRqeFUwalBSNEpTbHlrMjZ6cUtxT2dRWk9NdElDcGRmUko1a1F4MENlUWFja0FvVjZDZnY0RXVkT0NiNUJnOVlITGQzOHRGYTBLUG10UEZzSzJhdWozakN1SUo5VXZRNmdkUGlz", "domain": "mp.weixin.qq.com", "path": "/", "expires": 1754229560, "httpOnly": true, "secure": true, "sameSite": "Lax"}, {"name": "slave_user", "value": "gh_5624b93dc545", "domain": "mp.weixin.qq.com", "path": "/", "expires": 1754229560, "httpOnly": true, "secure": true, "sameSite": "Lax"}, {"name": "xid", "value": "fedc4a4eeb3b124f5812d5b88b1265f4", "domain": "mp.weixin.qq.com", "path": "/", "expires": 1788405116.292443, "httpOnly": true, "secure": true, "sameSite": "Lax"}, {"name": "mm_lang", "value": "zh_CN", "domain": "mp.weixin.qq.com", "path": "/", "expires": 1788405116.292574, "httpOnly": false, "secure": true, "sameSite": "Lax"}, {"name": "_clsk", "value": "19x6m1p|1753883936581|2|1|mp.weixin.qq.com/weheat-agent/payload/record", "domain": ".qq.com", "path": "/", "expires": 1753970336, "httpOnly": false, "secure": false, "sameSite": "Lax"}], "origins": [{"origin": "https://mp.weixin.qq.com", "localStorage": [{"name": "USER_AGENT_BROWSER_MAJOR_VERSION", "value": "127"}, {"name": "loginMode", "value": "1"}, {"name": "adTransitionText", "value": "noAdText"}, {"name": "hasClick", "value": "[0]"}]}]}
{"cookies": [{"name": "ua_id", "value": "Yp8GoVBn7Tj0PWR7AAAAADaCnYPYpvql9WPFzFkbL_A=", "domain": "mp.weixin.qq.com", "path": "/", "expires": 1788489057.928054, "httpOnly": true, "secure": true, "sameSite": "Lax"}, {"name": "wxuin", "value": "53845090997782", "domain": "mp.weixin.qq.com", "path": "/", "expires": 1788405090.13786, "httpOnly": false, "secure": false, "sameSite": "Lax"}, {"name": "uuid", "value": "8093d628ddf047d5a340ec437ca7fd08", "domain": "mp.weixin.qq.com", "path": "/", "expires": -1, "httpOnly": true, "secure": true, "sameSite": "Lax"}, {"name": "rand_info", "value": "CAESIEjp7aANF1MLCho4lIEi6Ju8oYLEexqCUxELmbqpfG3Y", "domain": "mp.weixin.qq.com", "path": "/", "expires": 1754274687, "httpOnly": true, "secure": true, "sameSite": "Lax"}, {"name": "slave_bizuin", "value": "3932822909", "domain": "mp.weixin.qq.com", "path": "/", "expires": 1754274687, "httpOnly": true, "secure": true, "sameSite": "Lax"}, {"name": "data_bizuin", "value": "3932822909", "domain": "mp.weixin.qq.com", "path": "/", "expires": 1754274687, "httpOnly": true, "secure": true, "sameSite": "Lax"}, {"name": "bizuin", "value": "3932822909", "domain": "mp.weixin.qq.com", "path": "/", "expires": 1754274687, "httpOnly": true, "secure": true, "sameSite": "Lax"}, {"name": "data_ticket", "value": "TXXMQ0chAAhF8D9olYk9AhfbDAjnS2jHUUO8GHs0e45nt6P/zjA38GlY5q12SD0E", "domain": "mp.weixin.qq.com", "path": "/", "expires": 1754274687, "httpOnly": true, "secure": true, "sameSite": "Lax"}, {"name": "slave_sid", "value": "OHl1NHRHZjd2N3poQTFDQ2d5WVBBWnJ5bm9obkZNVjZGank5S2dxSUMwblRqeFUwalBSNEpTbHlrMjZ6cUtxT2dRWk9NdElDcGRmUko1a1F4MENlUWFja0FvVjZDZnY0RXVkT0NiNUJnOVlITGQzOHRGYTBLUG10UEZzSzJhdWozakN1SUo5VXZRNmdkUGlz", "domain": "mp.weixin.qq.com", "path": "/", "expires": 1754274687, "httpOnly": true, "secure": true, "sameSite": "Lax"}, {"name": "slave_user", "value": "gh_5624b93dc545", "domain": "mp.weixin.qq.com", "path": "/", "expires": 1754274687, "httpOnly": true, "secure": true, "sameSite": "Lax"}, {"name": "xid", "value": "fedc4a4eeb3b124f5812d5b88b1265f4", "domain": "mp.weixin.qq.com", "path": "/", "expires": 1788405116.292443, "httpOnly": true, "secure": true, "sameSite": "Lax"}, {"name": "mm_lang", "value": "zh_CN", "domain": "mp.weixin.qq.com", "path": "/", "expires": 1788405116.292574, "httpOnly": false, "secure": true, "sameSite": "Lax"}, {"name": "_clck", "value": "3932822909|1|fy2|0", "domain": ".qq.com", "path": "/", "expires": 1785464417, "httpOnly": false, "secure": false, "sameSite": "Lax"}, {"name": "_clsk", "value": "184wq4q|1753929064020|7|1|mp.weixin.qq.com/weheat-agent/payload/record", "domain": ".qq.com", "path": "/", "expires": 1754015464, "httpOnly": false, "secure": false, "sameSite": "Lax"}], "origins": [{"origin": "https://mp.weixin.qq.com", "localStorage": [{"name": "USER_AGENT_BROWSER_MAJOR_VERSION", "value": "127"}, {"name": "loginMode", "value": "1"}, {"name": "adTransitionText", "value": "noAdText"}, {"name": "hasClick", "value": "[0]"}]}]}
test/test_face_swap.py
@@ -1,7 +1,6 @@
from auto_media_publisher.utils.face_swap.face_swaper_ffmpeg import FaceSwapWorker
from auto_media_publisher.utils.face_swap.face_swaper_ffmpeg_stream import FFmpegVideoProcessor
from auto_media_publisher.modifiers.face_swap.face_swapper_enhancer import FaceSwapperEnhancer
import time
import os
import warnings
warnings.filterwarnings("ignore")
@@ -9,13 +8,11 @@
    processor2 = FFmpegVideoProcessor(enable_enhancer=True,enable_nsfw_predictor=False)
    processor = FaceSwapperEnhancer()
    start_time = time.time()
    processor2.run("./temp/lilizhen.jpg","./temp/buhui.mp4", "./temp/buhui-lilizhen-with-enhancer.mp4", num_workers=4)
    #worker.swap_face('./temp/lilizhen.jpg', './temp/g1.mp4', './temp/g1_no_enhancer.mp4',1)
    processor.swap("./temp/buhui.mp4", "./temp/buhui-lilizhen-with-new.mp4","./temp/lilizhen.jpg")
    print(time.time()-start_time)
test/test_video_matting.py
New file
@@ -0,0 +1,37 @@
from auto_media_publisher.modifiers.video_matting.video_matting_converter import VideoMattingConverter
import time
if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('--input', type=str, required=True)
    parser.add_argument('--bg_video', type=str)
    parser.add_argument('--bg_jpg', type=str)
    parser.add_argument('--output_video', type=str)
    parser.add_argument('--output_alpha', type=str)
    parser.add_argument('--output_fore', type=str)
    args = parser.parse_args()
    time_start = time.time()
    converter = VideoMattingConverter()
    print(f"model loaded .. {time.time()-time_start}")
    time_start = time.time()
    converter.video_matting(
        input=args.input,
        bg_video=args.bg_video,
        bg_jpg=args.bg_jpg,
        output_video=args.output_video,
        output_alpha=args.output_alpha,
        output_fore=args.output_fore,
        )
    print(f"video matting finished .. {time.time()-time_start}")
    # conda activate mediapublish
    # f:
    # cd automedia
    #python -m test.test_video_matting --input ./temp/buhui.mp4 --bg_video ./temp/ocean.mp4 --output_video ./temp/buhui_ocean-36.mp4