From c706a69abfe9bae1322adbd344dfdaae37cdf564 Mon Sep 17 00:00:00 2001
From: wangzhibo <wangzhibo@shsening.com>
Date: 星期三, 30 七月 2025 22:16:55 +0800
Subject: [PATCH] add face swaper
---
auto_media_publisher/modifiers/face_swap/__init__.py | 0
auto_media_publisher/modifiers/face_swap/face_nsfw_predictor.py | 28 +
auto_media_publisher/modifiers/face_swap/face_enhancer.py | 49 +++
auto_media_publisher/modifiers/face_swap/face_swaper_avpy.py | 166 +++++++++++
auto_media_publisher/modifiers/face_swap/typing.py | 7
auto_media_publisher/utils/media_input_output.py | 93 ++++++
auto_media_publisher/modifiers/face_swap/face_analyser.py | 53 +++
auto_media_publisher/modifiers/face_swap/face_swaper_ffmpeg.py | 238 +++++++++++++++
auto_media_publisher/modifiers/face_swap/face_swaper_ffmpeg_stream.py | 241 ++++++++++++++++
9 files changed, 875 insertions(+), 0 deletions(-)
diff --git a/auto_media_publisher/modifiers/face_swap/__init__.py b/auto_media_publisher/modifiers/face_swap/__init__.py
new file mode 100644
index 0000000..e69de29
--- /dev/null
+++ b/auto_media_publisher/modifiers/face_swap/__init__.py
diff --git a/auto_media_publisher/modifiers/face_swap/face_analyser.py b/auto_media_publisher/modifiers/face_swap/face_analyser.py
new file mode 100644
index 0000000..1c4538b
--- /dev/null
+++ b/auto_media_publisher/modifiers/face_swap/face_analyser.py
@@ -0,0 +1,53 @@
+
+import torch
+from typing import Optional, List
+import numpy
+from insightface.app import FaceAnalysis
+from insightface.app.common import Face
+
+from gfpgan.utils import GFPGANer
+from .typing import Frame # 浣犲凡鏈夊畾涔�
+
+SIMILAR_FACE_DISTANCE: float = 0.45 # < 0.4 鍚屼竴浜� 0.4 -0.6 澶ф鏄悓涓�浜� >0.6 涓嶆槸鍚屼竴涓�
+
+
+class FaceAnalyser:
+ def __init__(self):
+ self.device = self._select_device()
+ self._analyser = FaceAnalysis(name='buffalo_l', root='./', providers=['CUDAExecutionProvider' if self.device == 'cuda' else 'CPUExecutionProvider'])
+ self._analyser.prepare(ctx_id=0, det_size=(640, 640))
+
+ def _select_device(self):
+ if torch.cuda.is_available():
+ return 'cuda'
+ elif torch.backends.mps.is_available():
+ return 'mps'
+ else:
+ return 'cpu'
+
+ def get_many_faces(self,frame: Frame) -> 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]:
+ many_faces = self.get_many_faces(frame)
+ if many_faces:
+ try:
+ return many_faces[position]
+ except IndexError:
+ return many_faces[-1]
+ return None
+
+
+ def find_similar_face(self,frame: Frame, reference_face: Face) -> Optional[Face]:
+ many_faces = self.get_many_faces(frame)
+ if many_faces:
+ for face in many_faces:
+ if hasattr(face, 'normed_embedding') and hasattr(reference_face, 'normed_embedding'):
+ distance = numpy.sum(numpy.square(face.normed_embedding - reference_face.normed_embedding))
+ if distance < SIMILAR_FACE_DISTANCE:
+ return face
+ return None
diff --git a/auto_media_publisher/modifiers/face_swap/face_enhancer.py b/auto_media_publisher/modifiers/face_swap/face_enhancer.py
new file mode 100644
index 0000000..1b600fe
--- /dev/null
+++ b/auto_media_publisher/modifiers/face_swap/face_enhancer.py
@@ -0,0 +1,49 @@
+
+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
new file mode 100644
index 0000000..8cc5f96
--- /dev/null
+++ b/auto_media_publisher/modifiers/face_swap/face_nsfw_predictor.py
@@ -0,0 +1,28 @@
+import threading
+import numpy as np
+from PIL import Image
+from keras import Model
+import opennsfw2
+
+from .typing import Frame # 浣犲凡鏈夊畾涔�
+
+MAX_PROBABILITY = 0.85
+
+class NSFWPredictor:
+ def __init__(self):
+ self._predictor = opennsfw2.make_open_nsfw_model()
+
+ def predict_frame(self, frame: Frame) -> bool:
+ image = Image.fromarray(frame)
+ image = opennsfw2.preprocess_image(image, opennsfw2.Preprocessing.YAHOO)
+ views = np.expand_dims(image, axis=0)
+ _, probability = self._predictor.predict(views)[0]
+ return probability > MAX_PROBABILITY
+
+ def predict_image(self, image_path: str) -> bool:
+ probability = opennsfw2.predict_image(image_path)
+ return probability > MAX_PROBABILITY
+
+ def predict_video(self, video_path: str, frame_interval: int = 100) -> bool:
+ _, probabilities = opennsfw2.predict_video_frames(video_path, frame_interval=frame_interval)
+ return any(prob > MAX_PROBABILITY for prob in probabilities)
diff --git a/auto_media_publisher/modifiers/face_swap/face_swaper_avpy.py b/auto_media_publisher/modifiers/face_swap/face_swaper_avpy.py
new file mode 100644
index 0000000..52b0f41
--- /dev/null
+++ b/auto_media_publisher/modifiers/face_swap/face_swaper_avpy.py
@@ -0,0 +1,166 @@
+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
new file mode 100644
index 0000000..12f4d11
--- /dev/null
+++ b/auto_media_publisher/modifiers/face_swap/face_swaper_ffmpeg.py
@@ -0,0 +1,238 @@
+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
new file mode 100644
index 0000000..0f2452d
--- /dev/null
+++ b/auto_media_publisher/modifiers/face_swap/face_swaper_ffmpeg_stream.py
@@ -0,0 +1,241 @@
+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/typing.py b/auto_media_publisher/modifiers/face_swap/typing.py
new file mode 100644
index 0000000..bc8d6fc
--- /dev/null
+++ b/auto_media_publisher/modifiers/face_swap/typing.py
@@ -0,0 +1,7 @@
+
+from typing import Any
+from insightface.app.common import Face
+import numpy
+
+Face = Face
+Frame = numpy.ndarray[Any, Any]
\ No newline at end of file
diff --git a/auto_media_publisher/utils/media_input_output.py b/auto_media_publisher/utils/media_input_output.py
new file mode 100644
index 0000000..4e2e209
--- /dev/null
+++ b/auto_media_publisher/utils/media_input_output.py
@@ -0,0 +1,93 @@
+import av
+import os
+import pims
+import numpy as np
+from torch.utils.data import Dataset
+from torchvision.transforms.functional import to_pil_image
+from PIL import Image
+
+
+class VideoReader(Dataset):
+ def __init__(self, path, transform=None):
+ self.video = pims.PyAVVideoReader(path)
+ self.rate = self.video.frame_rate
+ self.transform = transform
+
+ @property
+ def frame_rate(self):
+ return self.rate
+
+ def __len__(self):
+ return len(self.video)
+
+ def __getitem__(self, idx):
+ frame = self.video[idx]
+ frame = Image.fromarray(np.asarray(frame))
+ if self.transform is not None:
+ frame = self.transform(frame)
+ return frame
+
+
+class VideoWriter:
+ def __init__(self, path, frame_rate, bit_rate=1000000):
+ 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 = {
+ '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',
+ }
+
+ def write(self, frames):
+ # frames: [T, C, H, W]
+ self.stream.width = frames.size(3)
+ self.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))
+
+ def close(self):
+ self.container.mux(self.stream.encode())
+ self.container.close()
+
+
+class ImageSequenceReader(Dataset):
+ def __init__(self, path, transform=None):
+ self.path = path
+ self.files = sorted(os.listdir(path))
+ self.transform = transform
+
+ def __len__(self):
+ return len(self.files)
+
+ def __getitem__(self, idx):
+ with Image.open(os.path.join(self.path, self.files[idx])) as img:
+ img.load()
+ if self.transform is not None:
+ return self.transform(img)
+ return img
+
+
+class ImageSequenceWriter:
+ def __init__(self, path, extension='jpg'):
+ self.path = path
+ self.extension = extension
+ self.counter = 0
+ os.makedirs(path, exist_ok=True)
+
+ def write(self, frames):
+ # frames: [T, C, H, W]
+ for t in range(frames.shape[0]):
+ to_pil_image(frames[t]).save(os.path.join(
+ self.path, str(self.counter).zfill(4) + '.' + self.extension))
+ self.counter += 1
+
+ def close(self):
+ pass
+
\ No newline at end of file
--
Gitblit v1.9.1