From c00e9454085ef3e96b98a1ded81aad6a3b95d849 Mon Sep 17 00:00:00 2001
From: wangzhibo <wangzhibo@shsening.com>
Date: 星期四, 31 七月 2025 22:20:08 +0800
Subject: [PATCH] add video matting
---
auto_media_publisher/modifiers/video_matting/__init__.py | 0
.gitignore | 3
auto_media_publisher/utils/media_input_output.py | 58 ++
auto_media_publisher/modifiers/face_swap/face_analyser.py | 10
auto_media_publisher/modifiers/video_matting/matting_model/fast_guided_filter.py | 76 +++
auto_media_publisher/modifiers/video_matting/matting_model/model.py | 79 +++
auto_media_publisher/modifiers/rewriter/docx_rewriter_gemini.py | 2
auto_media_publisher/modifiers/video_matting/matting_model/decoder.py | 210 ++++++++
auto_media_publisher/modifiers/video_matting/video_matting_converter.py | 214 ++++++++
data/auto_media_publisher.sqlite3 | 0
/dev/null | 241 ----------
data/cookies/18918594752/gongzhonghao.json | 2
test/test_face_swap.py | 11
auto_media_publisher/modifiers/video_matting/matting_model/resnet.py | 45 +
auto_media_publisher/modifiers/face_swap/face_nsfw_predictor.py | 5
auto_media_publisher/modifiers/video_matting/matting_model/lraspp.py | 29 +
test/test_video_matting.py | 37 +
auto_media_publisher/modifiers/video_matting/matting_model/mobilenetv3.py | 72 +++
auto_media_publisher/modifiers/video_matting/matting_model/__init__.py | 0
auto_media_publisher/modifiers/video_matting/matting_model/deep_guided_filter.py | 61 ++
auto_media_publisher/modifiers/face_swap/face_swapper_enhancer.py | 248 ++++++++++
auto_media_publisher/generators/docx_generator_gemini.py | 4
22 files changed, 1,135 insertions(+), 272 deletions(-)
diff --git a/.gitignore b/.gitignore
index 027c8c1..202bfc8 100644
--- a/.gitignore
+++ b/.gitignore
@@ -30,4 +30,5 @@
*.swp
# 鍏朵粬涓嶅お鎯崇鍏宠仈鐨�
-涓嬭浇19妤艰鍧�/
\ No newline at end of file
+涓嬭浇19妤艰鍧�/
+data/cookies/18918594752/gongzhonghao.json
diff --git a/auto_media_publisher/generators/docx_generator_gemini.py b/auto_media_publisher/generators/docx_generator_gemini.py
index 353f5ee..292036b 100644
--- a/auto_media_publisher/generators/docx_generator_gemini.py
+++ b/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()
diff --git a/auto_media_publisher/modifiers/face_swap/face_analyser.py b/auto_media_publisher/modifiers/face_swap/face_analyser.py
index 1c4538b..f6556f3 100644
--- a/auto_media_publisher/modifiers/face_swap/face_analyser.py
+++ b/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:
diff --git a/auto_media_publisher/modifiers/face_swap/face_enhancer.py b/auto_media_publisher/modifiers/face_swap/face_enhancer.py
deleted file mode 100644
index 1b600fe..0000000
--- a/auto_media_publisher/modifiers/face_swap/face_enhancer.py
+++ /dev/null
@@ -1,49 +0,0 @@
-
-import torch
-from typing import List
-from insightface.app.common import Face
-
-from gfpgan.utils import GFPGANer
-from .typing import Frame # 浣犲凡鏈夊畾涔�
-from auto_media_publisher.config.conf_base import FACE_ENHANCER_MODEL_PATH
-
-SIMILAR_FACE_DISTANCE: float = 0.45 # < 0.4 鍚屼竴浜� 0.4 -0.6 澶ф鏄悓涓�浜� >0.6 涓嶆槸鍚屼竴涓�
-
-
-class FaceEnhancer:
- def __init__(self):
- self.device = self._select_device()
- self._enhancer = GFPGANer(model_path=FACE_ENHANCER_MODEL_PATH, upscale=1, device=self.device)
-
- def _select_device(self):
- if torch.cuda.is_available():
- return 'cuda'
- elif torch.backends.mps.is_available():
- return 'mps'
- else:
- return 'cpu'
-
- def _enhance_face(self,target_face: Face, temp_frame: Frame) -> Frame:
- 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._enhancer.enhance(
- temp_face,
- paste_back=True
- )
- temp_frame[start_y:end_y, start_x:end_x] = temp_face
- return temp_frame
-
-
- def process_frame(self,temp_frame: Frame,many_faces:List) -> Frame:
- #many_faces = self._get_many_faces(temp_frame) check face_analyser.py
- if many_faces:
- for target_face in many_faces:
- temp_frame = self._enhance_face(target_face, temp_frame)
- return temp_frame
diff --git a/auto_media_publisher/modifiers/face_swap/face_nsfw_predictor.py b/auto_media_publisher/modifiers/face_swap/face_nsfw_predictor.py
index 8cc5f96..8a2eb7a 100644
--- a/auto_media_publisher/modifiers/face_swap/face_nsfw_predictor.py
+++ b/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)
diff --git a/auto_media_publisher/modifiers/face_swap/face_swaper_avpy.py b/auto_media_publisher/modifiers/face_swap/face_swaper_avpy.py
deleted file mode 100644
index 52b0f41..0000000
--- a/auto_media_publisher/modifiers/face_swap/face_swaper_avpy.py
+++ /dev/null
@@ -1,166 +0,0 @@
-import av
-import torch
-import threading, queue
-from threading import current_thread
-import numpy as np
-import insightface
-import cv2
-import time
-from insightface.app.common import Face
-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 insightface.app import FaceAnalysis
-
-SIMILAR_FACE_DISTANCE: float = 0.45 # < 0.4 鍚屼竴浜� 0.4 -0.6 澶ф鏄悓涓�浜� >0.6 涓嶆槸鍚屼竴涓�
-
-def _select_device():
- if torch.cuda.is_available():
- return 'cuda:0'
- elif torch.backends.mps.is_available():
- return 'mps'
- else:
- return 'cpu'
-
-NUM_WORKERS = 1
-
-#face_swapper = insightface.model_zoo.get_model(FACE_SWAPPER_MODEL_PATH, providers=['CUDAExecutionProvider'])
-#face_enhancer = GFPGANer(model_path=FACE_ENHANCER_MODEL_PATH, upscale=1, device=_select_device())
-
-frame_q = queue.Queue(maxsize=10)
-processed_q = queue.Queue(maxsize=10)
-
-THREAD_SEMAPHORE =threading.Semaphore(4) # 涓嶉檺鍒朵簡銆�
-
-def enhance_face(thread_face_enhancer,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]
- #print(temp_face.size)
- if temp_face.size:
- _, _, temp_face = thread_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(video_path):
- container = av.open(video_path)
- for idx, frame in enumerate(container.decode(video=0)):
- img = frame.to_ndarray(format='rgb24')
- frame_q.put((idx, img))
- for _ in range(NUM_WORKERS):
- frame_q.put((None, None))
-
-def process_frame_loop(new_face):
- start_time = time.time()
- thread_face_analyser = FaceAnalysis(name='buffalo_l', root='./', providers=['CUDAExecutionProvider'])
- thread_face_analyser.prepare(ctx_id=3, det_size=(640, 640))
- #thread_face_enhancer = GFPGANer(model_path=FACE_ENHANCER_MODEL_PATH, upscale=1, device=_select_device())
- thread_face_enhancer = GFPGANer(model_path=FACE_ENHANCER_MODEL_PATH, upscale=1, device=_select_device())
- thread_face_swapper = insightface.model_zoo.get_model(FACE_SWAPPER_MODEL_PATH, providers=['CUDAExecutionProvider'])
-
- while True:
- idx, frame = frame_q.get()
- if idx is None:
- processed_q.put((None, None))
- break
- try:
- many_faces = thread_face_analyser.get(frame)
- if many_faces:
- old_face = many_faces[0]
- frame = thread_face_swapper.get(frame, old_face, new_face, paste_back=True)
-
- for sinlge_face in many_faces:
- frame_ = enhance_face(thread_face_enhancer,sinlge_face, frame)
-
- 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 encode_video(output_path, fps, width, height):
- out = av.open(output_path, mode='w')
-
- #stream = out.add_stream('h264_nvenc', rate=fps)
- stream = out.add_stream('libx264', rate=fps)
- stream.width = width
- stream.height = height
- stream.pix_fmt = 'yuv420p'
-
- stream.options = {'crf': '18', 'preset': 'veryfast'}
- #stream.options = {'rc': 'constqp', 'qp': '18','preset': 'p7','profile': 'high'}
-
- next_index = 0
- buffer = {}
- end_count = 0
-
- while True:
- idx, frame = processed_q.get()
- if idx is None:
- end_count += 1
- if end_count >= NUM_WORKERS:
- break
- continue
- 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 stream.encode(video_frame):
- out.mux(packet)
- next_index += 1
-
- # 娓呯┖缂栫爜缂撳瓨
- for packet in stream.encode():
- out.mux(packet)
- out.close()
-# 鍚姩浠诲姟
-video_path = './temp/buhui.mp4'
-output_path = './temp/buhui-wangrong-enhancer.mp4'
-new_face_image_path = './temp/wangrong.jpg'
-face_analyser = FaceAnalysis(name='buffalo_l', root='./', providers=['CUDAExecutionProvider'])
-face_analyser.prepare(ctx_id=0, det_size=(640, 640))
-
-faces = face_analyser.get(cv2.imread(new_face_image_path))
-if faces:
- new_face = sorted(faces, key=lambda x: x.bbox[0])[0]
-
- #threading.Thread(target=decode_video, args=(video_path,)).start()
- #threading.Thread(target=process_frame_loop, args=(new_face,)).start()
- #threading.Thread(target=encode_video, args=(output_path, 25,576,1080)).start()
- threads = []
-
- # 瑙ｇ爜绾跨▼
- t_decode = threading.Thread(target=decode_video,args=(video_path,))
- threads.append(t_decode)
- t_decode.start()
-
- # 澶氬鐞嗙嚎绋�
- for i in range(NUM_WORKERS):
- t = threading.Thread(target=process_frame_loop, args=(new_face,),name=f"process_worker_{i}")
- threads.append(t)
- t.start()
-
- # 缂栫爜绾跨▼
- t_encode = threading.Thread(target=encode_video,args=(output_path, 25,576,1080))
- threads.append(t_encode)
- t_encode.start()
-
- for t in threads:
- t.join()
-
- print("Processing complete.")
-
-
diff --git a/auto_media_publisher/modifiers/face_swap/face_swaper_ffmpeg.py b/auto_media_publisher/modifiers/face_swap/face_swaper_ffmpeg.py
deleted file mode 100644
index 12f4d11..0000000
--- a/auto_media_publisher/modifiers/face_swap/face_swaper_ffmpeg.py
+++ /dev/null
@@ -1,238 +0,0 @@
-import cv2
-import os
-import torch
-import queue
-import threading
-import subprocess
-import json
-import numpy as np
-import time
-from typing import Union
-import insightface
-from insightface.app.common import Face
-from .face_nsfw_predictor import NSFWPredictor
-from .face_enhancer import FaceEnhancer
-from .face_analyser import FaceAnalyser
-from auto_media_publisher.utils.logger import get_logger
-from auto_media_publisher.config.conf_base import FACE_SWAPPER_MODEL_PATH
-from auto_media_publisher.utils.media_input_output import VideoReader,VideoWriter
-from torch.utils.data import DataLoader
-from torchvision import transforms
-
-class FaceSwapWorker:
- def __init__(self, enable_enhancer:bool = False,enable_nsfw_predictor:bool=False):
- self.logger = get_logger("FaceSwapWorker", "admin")
- self.device = self._select_device()
- self.face_swaper = insightface.model_zoo.get_model(FACE_SWAPPER_MODEL_PATH, providers=['CUDAExecutionProvider' if self.device == 'cuda' else 'CPUExecutionProvider'])
- self.main_analyser = FaceAnalyser()
- self.enable_enhancer = enable_enhancer
- self.enable_nsfw_predictor = enable_nsfw_predictor
- self.batch_size = 4
- if self.enable_nsfw_predictor:
- self.nsfw_predictor = NSFWPredictor()
- self._init()
-
- def _init(self):
- self.frame_queue = queue.Queue()
- self.processed_queue = queue.Queue()
- self.buffered_frames = {}
- self.expected_index = 0
- self.fps = 0
- self.width = 0
- self.height = 0
- self.total_frames = 0
- self.process_threads = []
- self.num_processors = 4
-
- def _select_device(self):
- if torch.cuda.is_available():
- return 'cuda'
- elif torch.backends.mps.is_available():
- return 'mps'
- else:
- return 'cpu'
-
- def auto_downsample_ratio(self,h, w):
- """
- Automatically find a downsample ratio so that the largest side of the resolution be 512px.
- """
- return min(512 / max(h, w), 1)
-
- def _decode_video(self):
- transform = transforms.ToTensor()
- source = VideoReader(self.input_video_path, transform)
-
- reader = DataLoader(source, batch_size=self.batch_size, pin_memory=True, num_workers=self.num_processors)
-
- writer = VideoWriter(
- path=self.output_video_path,
- frame_rate=30,
- bit_rate=int(1 * 1000000))
-
- for src in reader:
-
- if downsample_ratio is None:
- downsample_ratio = auto_downsample_ratio(*src.shape[2:])
-
- src = src.to(device, dtype, non_blocking=True).unsqueeze(0) # [B, T, C, H, W]
- fgr, pha, *rec = model(src, *rec, downsample_ratio)
-
- if output_foreground is not None:
- writer_fgr.write(fgr[0])
- if output_alpha is not None:
- writer_pha.write(pha[0])
- if output_composition is not None:
- if output_type == 'video':
- com = fgr * pha + bgr * (1 - pha)
- else:
- fgr = fgr * pha.gt(0)
- com = torch.cat([fgr, pha], dim=-3)
- writer_com.write(com[0])
-
-
-
-
- def _process_frame(self):
- thread_face_analyser = FaceAnalyser()
- if self.enable_enhancer:
- thread_face_enhancer = FaceEnhancer()
- while True:
- item = self.frame_queue.get()
- if item is None:
- self.processed_queue.put(None) # 鎶婄粨鏉熶俊鍙风户缁紶閫�
- break
-
- idx, frame = item
-
- print(f"[Thread ID: {threading.get_ident()}] is processing frame : {idx} of video file 锛歿self.input_video_path}")
-
- old_face = thread_face_analyser.get_one_face(frame)
- if old_face:
- frame = self.face_swaper.get(frame,old_face,self.new_face)
- else:
- print("face not detected")
-
- print(self.enable_enhancer)
-
- if self.enable_enhancer:
- many_faces= thread_face_analyser.get_many_faces(frame)
- if many_faces:
- frame = thread_face_enhancer.process_frame(frame,many_faces)
-
- self.processed_queue.put((idx, frame))
-
- def _encode_video(self):
- while not (self.width and self.height and self.fps):
- time.sleep(0.01)
- encode_cmd = [
- "ffmpeg",
- "-loglevel", "info",
- "-y", # 瑕嗙洊杈撳嚭鏂囦欢
- "-f", "rawvideo",
- "-pixel_format", "rgb24",
- "-video_size", f"{self.width}x{self.height}",
- "-framerate", str(self.fps),
- "-i", "pipe:0", # 浠� stdin 璇诲彇鍘熷瑙嗛甯�
- #"-i", self.input_video_path if self.new_audio_path is None else self.new_audio_path, # 鎻愪緵闊抽
- "-c:v", "libx264", # 浣犲彲浠ユ浛鎹㈡垚 h264_nvenc銆乭evc_nvenc銆乭264_amf銆乭264_qsv 绛夌‖浠剁紪鐮佸櫒,璐ㄩ噺鍙兘绋嶅井閫婅壊
- "-crf", "18",
- "-preset", "veryfast",
- "-pix_fmt", "yuv420p",
- #"-c:a", "copy",
- #"-map", "0:v:0",
- #"-map", "1:a:0",
- #"-shortest",
- self.output_video_path
- ]
-
- process = subprocess.Popen(encode_cmd, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE,bufsize=10**8)
-
- end_signals_received = 0
- finished = 0
-
- while True:
- item = self.processed_queue.get()
- if item is None:
- end_signals_received += 1
- if end_signals_received == self.num_processors:
- break
- else:
- continue
- idx,frame = item
- self.buffered_frames[idx] = frame
-
- while self.expected_index in self.buffered_frames:
- next_frame = self.buffered_frames.pop(self.expected_index)
- print(f"_encode frame next_frame: {self.expected_index}")
- try:
- process.stdin.write(next_frame.astype(np.uint8).tobytes())
- except BrokenPipeError:
- print("BrokenPipeError: ffmpeg process closed stdin.")
- break
- self.expected_index += 1
- finished += 1
- print(f"_encode frame finished: {finished}")
-
-
- print(f"Total frames encoded: {finished}")
-
- process.stdin.close()
- return_code = process.wait()
- print("FFmpeg exited with code:", return_code)
-
- def swap_face(self,new_face_image_path:str,input_video_path: str, output_video_path: str,num_processors=4):
- """
- 鎵ц鎹㈣劯鎿嶄綔锛屽弬鏁板涓嬶細
- 1銆� 鏂扮殑鑴哥収鐗囷紝璺緞瑕佹纭紝杩樿鍖呭惈涓�寮犱汉鑴革紱
- 2銆� 鑰佺殑瑙嗛锛岃矾寰勮姝ｇ‘
- 3銆� 鏂拌棰� 鏈�缁堣緭鍑虹殑瑙嗛璺緞
- 4銆� 骞跺彂澶勭悊杩涚▼鏁伴噺锛岄粯璁ゆ槸4锛屼緷鎹‖浠舵�ц兘璁剧疆
- 5銆� 鏂拌棰� 鐨勯煶棰戯紝濡傛灉涓嶄负绌猴紝灏变娇鐢紱 鍚﹀垯 浣跨敤鑰佽棰戠殑闊抽
- """
- self._init()
- self.num_processors = num_processors
-
- if not os.path.exists(new_face_image_path):
- self.logger.error(f"face swap image_path: {new_face_image_path} does not exists.")
- return
- self.new_face = self.main_analyser.get_one_face(cv2.imread(new_face_image_path))
- if not self.new_face :
- self.logger.error(f"target image DO NOT has face : {new_face_image_path} does not exists.")
- return
-
- if not os.path.exists(input_video_path):
- self.logger.error(f"face swap video_path: {input_video_path} does not exists.")
- return
-
- self.input_video_path = input_video_path
- self.output_video_path = output_video_path
- #self.new_audio_path =new_audio_path
-
- if self.enable_nsfw_predictor and self.nsfw_predictor:
- if self.nsfw_predictor.predict_video(input_video_path):
- self.logger.error(f" 瑙嗛杩濊锛� {self.input_video_path}")
- return
-
- # 1. 鍚姩 encode 绾跨▼锛堥渶瑕佺瓑寰� width/height/fps 璁剧疆锛屼絾鍙互鏀惧湪鍓嶉潰锛�
- encode_thread = threading.Thread(target=self._encode_video)
- encode_thread.start()
-
- # 2. 鍚姩 decode 绾跨▼锛堣礋璐ｄ粠瑙嗛涓鍙栧抚锛�
- decode_thread = threading.Thread(target=self._decode_video)
- decode_thread.start()
-
- # 3. 鍚姩澶勭悊绾跨▼
- for i in range(self.num_processors):
- t = threading.Thread(target=self._process_frame)
- t.start()
- self.process_threads.append(t)
-
- # 4. 绛夊緟 decode 瀹屾垚
- decode_thread.join()
-
- # 5. 绛夊緟鎵�鏈夊鐞嗙嚎绋嬪畬鎴�
- for t in self.process_threads:
- t.join()
-
- # 6. 绛夊緟 encode 瀹屾垚
- encode_thread.join()
diff --git a/auto_media_publisher/modifiers/face_swap/face_swaper_ffmpeg_stream.py b/auto_media_publisher/modifiers/face_swap/face_swaper_ffmpeg_stream.py
deleted file mode 100644
index 0f2452d..0000000
--- a/auto_media_publisher/modifiers/face_swap/face_swaper_ffmpeg_stream.py
+++ /dev/null
@@ -1,241 +0,0 @@
-import os
-import cv2
-import subprocess
-import threading
-from threading import Thread, current_thread
-import queue
-import numpy as np
-import torch
-import json
-import insightface
-from .face_nsfw_predictor import NSFWPredictor
-from .face_enhancer import FaceEnhancer
-from .face_analyser import FaceAnalyser
-from auto_media_publisher.utils.logger import get_logger
-from auto_media_publisher.config.conf_base import FACE_SWAPPER_MODEL_PATH
-
-
-class FFmpegVideoProcessor:
- def __init__(self,enable_enhancer:bool=False,enable_nsfw_predictor:bool=False):
- self.logger = get_logger("FaceSwapWorker", "admin")
- self.device = self._select_device()
- print(f"using device {self.device}")
- self.face_swaper = insightface.model_zoo.get_model(FACE_SWAPPER_MODEL_PATH, providers=['CUDAExecutionProvider' if self.device == 'cuda' else 'CPUExecutionProvider'])
- self.main_analyser = FaceAnalyser()
- self.enable_enhancer = enable_enhancer
- self.enable_nsfw_predictor = enable_nsfw_predictor
- if self.enable_nsfw_predictor:
- self.nsfw_predictor = NSFWPredictor()
- self._init()
-
- def _init(self):
- self.width = 0
- self.height = 0
- self.fps = 0
- self.frame_queue = queue.Queue(maxsize=2000) # 瑙ｇ爜绾跨▼鏀惧抚
- self.processed_queue = queue.Queue(maxsize=2000) # 澶勭悊绾跨▼鏀惧抚
- self.buffered_frames = {}
- self.expected_index = 0
-
- self.stop_event = threading.Event()
-
- def _select_device(self):
- if torch.cuda.is_available():
- return 'cuda'
- elif torch.backends.mps.is_available():
- return 'mps'
- else:
- return 'cpu'
-
- def probe_video(self):
- cmd = [
- "ffprobe",
- "-v", "error",
- "-select_streams", "v:0",
- "-show_entries", "stream=width,height,r_frame_rate",
- "-of", "json",
- self.input_video_path,
- ]
- proc = subprocess.run(cmd, capture_output=True, text=True, check=True)
- info = json.loads(proc.stdout)
- stream = info["streams"][0]
- self.width = int(stream["width"])
- self.height = int(stream["height"])
- self.fps = eval(stream["r_frame_rate"])
- print(f"Video info: {self.width}x{self.height} @ {self.fps}fps")
-
- def decode_worker(self):
- decode_cmd = [
- "ffmpeg",
- "-loglevel", "error",
- "-i", self.input_video_path,
- "-f", "rawvideo",
- "-pix_fmt", "rgb24",
- "pipe:1",
- ]
- self.decode_proc = subprocess.Popen(decode_cmd, stdout=subprocess.PIPE, bufsize=10**8)
-
- frame_size = self.width * self.height * 3
- idx = 0
- while not self.stop_event.is_set():
- raw_frame = self.decode_proc.stdout.read(frame_size)
- if len(raw_frame) < frame_size:
- break
- frame = np.frombuffer(raw_frame, np.uint8).reshape((self.height, self.width, 3))
- self.frame_queue.put((idx, frame))
- idx += 1
-
- # 缁撴潫淇″彿
- for _ in range(self.num_workers):
- self.frame_queue.put(None)
- self.decode_proc.stdout.close()
- self.decode_proc.wait()
-
- def process_frame(self, frame: np.ndarray,thread_face_analyser,thread_face_enhancer) -> np.ndarray:
- # 杩欓噷鏀句綘鐨勬崲鑴搞�佸彉閫熺瓑澶勭悊閫昏緫
- # 渚嬪锛�
- # frame = your_face_swap_function(frame)
- # 鎴栬�呯畝鍗曠ず鑼冿細鍙嶈浆棰滆壊
- #return 255 - frame
- if self.enable_enhancer:
- many_faces= thread_face_analyser.get_many_faces(frame)
- if many_faces:
- frame = self.face_swaper.get(frame,many_faces[0],self.new_face)
- frame = thread_face_enhancer.process_frame(frame,many_faces)
- else:
- old_face = thread_face_analyser.get_one_face(frame)
- if old_face:
- frame = self.face_swaper.get(frame,old_face,self.new_face)
- else:
- print("face not detected")
-
- return frame
-
- def process_worker(self):
- thread_face_analyser = FaceAnalyser()
- if self.enable_enhancer:
- thread_face_enhancer = FaceEnhancer()
- else:
- thread_face_enhancer = None
- while True:
- item = self.frame_queue.get()
- if item is None:
- self.processed_queue.put(None)
- break
- idx, frame = item
- try:
- processed = self.process_frame(frame,thread_face_analyser,thread_face_enhancer)
- print(f"{current_thread().name} has proceesed frame :{idx}")
- except:
- print(f"{current_thread().name} has error frame :{idx}")
- processed=frame
- self.processed_queue.put((idx, processed))
-
- def encode_worker(self):
- input_audio_path = self.input_video_path
- if self.new_audio_path and os.path.exists(self.new_audio_path):
- input_audio_path = self.new_audio_path
- encode_cmd = [
- "ffmpeg",
- "-y",
- "-loglevel", "error",
- "-f", "rawvideo",
- "-pixel_format", "rgb24",
- "-video_size", f"{self.width}x{self.height}",
- "-framerate", str(self.fps),
- "-i", "pipe:0",
- "-i", input_audio_path,
- "-c:v", "libx264",
- "-crf", "18",
- "-preset", "veryfast",
- "-c:a", "copy",
- "-map", "0:v:0",
- "-map", "1:a:0",
- "-shortest",
- self.output_video_path,
- ]
- self.encode_proc = subprocess.Popen(encode_cmd, stdin=subprocess.PIPE, bufsize=10**8)
-
- end_signals = 0
- finished = 0
- while True:
- item = self.processed_queue.get()
- if item is None:
- end_signals += 1
- if end_signals == self.num_workers:
- break
- else:
- continue
-
- idx, frame = item
- self.buffered_frames[idx] = frame
-
- while self.expected_index in self.buffered_frames:
- f = self.buffered_frames.pop(self.expected_index)
- try:
- self.encode_proc.stdin.write(f.astype(np.uint8).tobytes())
- print(f"{self.expected_index} has been writed into stdin")
- except BrokenPipeError:
- print("FFmpeg encoding pipe broken.")
- self.stop_event.set()
- break
- self.expected_index += 1
- finished += 1
-
- self.encode_proc.stdin.close()
- self.encode_proc.wait()
- print(f"Encoding finished, total frames: {finished}")
-
- def run(self,new_face_image_path,input_video_path, output_video_path, num_workers=4,new_audio_path=None):
- """
- 鎵ц鎹㈣劯鎿嶄綔锛屽弬鏁板涓嬶細
- 1銆� 鏂扮殑鑴哥収鐗囷紝璺緞瑕佹纭紝杩樿鍖呭惈涓�寮犱汉鑴革紱
- 2銆� 鑰佺殑瑙嗛锛岃矾寰勮姝ｇ‘
- 3銆� 鏂拌棰� 鏈�缁堣緭鍑虹殑瑙嗛璺緞
- 4銆� 骞跺彂澶勭悊杩涚▼鏁伴噺锛岄粯璁ゆ槸4锛屼緷鎹‖浠舵�ц兘璁剧疆
- 5銆� 鏂拌棰� 鐨勯煶棰戯紝濡傛灉涓嶄负绌猴紝灏变娇鐢紱 鍚﹀垯 浣跨敤鑰佽棰戠殑闊抽
- """
- if not os.path.exists(new_face_image_path):
- self.logger.error(f"face swap image_path: {new_face_image_path} does not exists.")
- return
- self.new_face = self.main_analyser.get_one_face(cv2.imread(new_face_image_path))
-
- if not self.new_face :
- self.logger.error(f"target image DO NOT has face : {new_face_image_path} does not exists.")
- return
-
- if not os.path.exists(input_video_path):
- self.logger.error(f"face swap video_path: {input_video_path} does not exists.")
- return
-
- self._init()
- self.input_video_path = input_video_path
- self.output_video_path = output_video_path
- self.num_workers = num_workers
- self.new_audio_path =new_audio_path
-
- self.probe_video()
-
- threads = []
-
- # 瑙ｇ爜绾跨▼
- t_decode = threading.Thread(target=self.decode_worker)
- threads.append(t_decode)
- t_decode.start()
-
- # 澶氬鐞嗙嚎绋�
- for i in range(self.num_workers):
- t = threading.Thread(target=self.process_worker, name=f"process_worker_{i}")
- threads.append(t)
- t.start()
-
- # 缂栫爜绾跨▼
- t_encode = threading.Thread(target=self.encode_worker)
- threads.append(t_encode)
- t_encode.start()
-
- for t in threads:
- t.join()
-
- print("Processing complete.")
-
diff --git a/auto_media_publisher/modifiers/face_swap/face_swapper_enhancer.py b/auto_media_publisher/modifiers/face_swap/face_swapper_enhancer.py
new file mode 100644
index 0000000..5e04040
--- /dev/null
+++ b/auto_media_publisher/modifiers/face_swap/face_swapper_enhancer.py
@@ -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锛坙ayout锛�
+ 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
+
+
+
diff --git a/auto_media_publisher/modifiers/rewriter/docx_rewriter_gemini.py b/auto_media_publisher/modifiers/rewriter/docx_rewriter_gemini.py
index a607199..09c91b9 100644
--- a/auto_media_publisher/modifiers/rewriter/docx_rewriter_gemini.py
+++ b/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))
diff --git a/auto_media_publisher/modifiers/video_matting/__init__.py b/auto_media_publisher/modifiers/video_matting/__init__.py
new file mode 100644
index 0000000..e69de29
--- /dev/null
+++ b/auto_media_publisher/modifiers/video_matting/__init__.py
diff --git a/auto_media_publisher/modifiers/video_matting/matting_model/__init__.py b/auto_media_publisher/modifiers/video_matting/matting_model/__init__.py
new file mode 100644
index 0000000..e69de29
--- /dev/null
+++ b/auto_media_publisher/modifiers/video_matting/matting_model/__init__.py
diff --git a/auto_media_publisher/modifiers/video_matting/matting_model/decoder.py b/auto_media_publisher/modifiers/video_matting/matting_model/decoder.py
new file mode 100644
index 0000000..7307435
--- /dev/null
+++ b/auto_media_publisher/modifiers/video_matting/matting_model/decoder.py
@@ -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)
+
\ No newline at end of file
diff --git a/auto_media_publisher/modifiers/video_matting/matting_model/deep_guided_filter.py b/auto_media_publisher/modifiers/video_matting/matting_model/deep_guided_filter.py
new file mode 100644
index 0000000..a24b8c5
--- /dev/null
+++ b/auto_media_publisher/modifiers/video_matting/matting_model/deep_guided_filter.py
@@ -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)
diff --git a/auto_media_publisher/modifiers/video_matting/matting_model/fast_guided_filter.py b/auto_media_publisher/modifiers/video_matting/matting_model/fast_guided_filter.py
new file mode 100644
index 0000000..df9b4b2
--- /dev/null
+++ b/auto_media_publisher/modifiers/video_matting/matting_model/fast_guided_filter.py
@@ -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
\ No newline at end of file
diff --git a/auto_media_publisher/modifiers/video_matting/matting_model/lraspp.py b/auto_media_publisher/modifiers/video_matting/matting_model/lraspp.py
new file mode 100644
index 0000000..5fc7079
--- /dev/null
+++ b/auto_media_publisher/modifiers/video_matting/matting_model/lraspp.py
@@ -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)
\ No newline at end of file
diff --git a/auto_media_publisher/modifiers/video_matting/matting_model/mobilenetv3.py b/auto_media_publisher/modifiers/video_matting/matting_model/mobilenetv3.py
new file mode 100644
index 0000000..712a298
--- /dev/null
+++ b/auto_media_publisher/modifiers/video_matting/matting_model/mobilenetv3.py
@@ -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)
diff --git a/auto_media_publisher/modifiers/video_matting/matting_model/model.py b/auto_media_publisher/modifiers/video_matting/matting_model/model.py
new file mode 100644
index 0000000..71fc684
--- /dev/null
+++ b/auto_media_publisher/modifiers/video_matting/matting_model/model.py
@@ -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
diff --git a/auto_media_publisher/modifiers/video_matting/matting_model/resnet.py b/auto_media_publisher/modifiers/video_matting/matting_model/resnet.py
new file mode 100644
index 0000000..7eeff1c
--- /dev/null
+++ b/auto_media_publisher/modifiers/video_matting/matting_model/resnet.py
@@ -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)
diff --git a/auto_media_publisher/modifiers/video_matting/video_matting_converter.py b/auto_media_publisher/modifiers/video_matting/video_matting_converter.py
new file mode 100644
index 0000000..f8662d6
--- /dev/null
+++ b/auto_media_publisher/modifiers/video_matting/video_matting_converter.py
@@ -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) # 纭繚鏂囦欢閲婃斁锛學indows 涓婇潪甯稿繀瑕�
+ 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
+
+
diff --git a/auto_media_publisher/utils/media_input_output.py b/auto_media_publisher/utils/media_input_output.py
index 4e2e209..0adf1bd 100644
--- a/auto_media_publisher/utils/media_input_output.py
+++ b/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', # 褰卞搷缂栫爜閫熷害锛宮edium 鏄�1x锛宖ast鏄�2x faster鏄�3x veryfast鏄�4x superfast鏄�6x ultrafast鏄�10x锛屾枃浠朵綋绉ぇ绾︿細澧炲姞(浠edium涓哄熀鍑�) 10%锛�15%锛�20% 30% 50%
'tune': 'film',
}
+ self.audio_source_path =audio_source_path
+ self.audio_in_container = None
+ self.audio_stream_in = None
+ self.audio_out_stream = None
+
+ 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):
diff --git a/data/auto_media_publisher.sqlite3 b/data/auto_media_publisher.sqlite3
index d32698e..1364c4f 100644
--- a/data/auto_media_publisher.sqlite3
+++ b/data/auto_media_publisher.sqlite3
Binary files differ
diff --git a/data/cookies/18918594752/gongzhonghao.json b/data/cookies/18918594752/gongzhonghao.json
index 4cfe59d..588d31f 100644
--- a/data/cookies/18918594752/gongzhonghao.json
+++ b/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]"}]}]}
\ No newline at end of file
+{"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]"}]}]}
\ No newline at end of file
diff --git a/test/test_face_swap.py b/test/test_face_swap.py
index 7462699..675a059 100644
--- a/test/test_face_swap.py
+++ b/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)
diff --git a/test/test_video_matting.py b/test/test_video_matting.py
new file mode 100644
index 0000000..7d9e8f4
--- /dev/null
+++ b/test/test_video_matting.py
@@ -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
--
Gitblit v1.9.1