| | |
| | | *.swp |
| | | |
| | | # 其他不太想管关联的 |
| | | 下载19楼论坛/ |
| | | 下载19楼论坛/ |
| | | data/cookies/18918594752/gongzhonghao.json |
| | |
| | | |
| | | |
| | | 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() |
| | |
| | | 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 不是同一个 |
| | | |
| | |
| | | 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: |
| | |
| | | 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: |
| | |
| | | import threading |
| | | import numpy as np |
| | | from PIL import Image |
| | | from keras import Model |
| | | import opennsfw2 |
| | | |
| | | from .typing import Frame # 你已有定义 |
| | | |
| | | MAX_PROBABILITY = 0.85 |
| | | |
| | |
| | | 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) |
| New file |
| | |
| | | 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 |
| | | |
| | | |
| | | |
| | |
| | | |
| | | 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)) |
| New file |
| | |
| | | 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) |
| | | |
| New file |
| | |
| | | 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) |
| New file |
| | |
| | | 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 |
| New file |
| | |
| | | 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) |
| New file |
| | |
| | | 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) |
| New file |
| | |
| | | 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 |
| New file |
| | |
| | | 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) |
| New file |
| | |
| | | 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 |
| | | |
| | | |
| | |
| | | |
| | | |
| | | 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): |
| | |
| | | {"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]"}]}]} |
| | |
| | | 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") |
| | | |
| | |
| | | |
| | | |
| | | |
| | | 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) |
| | | |
| New file |
| | |
| | | |
| | | |
| | | 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 |