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
|
|
|
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
|
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 _decode_video(self):
|
# === Step 1: 用 ffprobe 获取视频信息 ===
|
probe_cmd = [
|
"ffprobe",
|
"-v", "error",
|
"-select_streams", "v:0",
|
"-show_entries", "stream=width,height,r_frame_rate",
|
"-of", "json",
|
self.input_video_path
|
]
|
probe_result = subprocess.run(probe_cmd, capture_output=True, text=True, check=True)
|
probe_data = json.loads(probe_result.stdout)
|
video_stream = probe_data['streams'][0]
|
self.width = int(video_stream['width'])
|
self.height = int(video_stream['height'])
|
self.fps = eval(video_stream['r_frame_rate']) # e.g. "30/1"
|
|
print(f"Video info: {self.width}x{self.height} @ {self.fps}fps")
|
|
# === Step 2: 用 ffmpeg 解码输出 rawvideo 到 stdout pipe ===
|
decode_cmd = [
|
"ffmpeg",
|
"-loglevel", "info",
|
"-i", self.input_video_path,
|
"-f", "rawvideo",
|
"-pix_fmt", "rgb24",
|
"pipe:1"
|
]
|
process = subprocess.Popen(decode_cmd, stdout=subprocess.PIPE, bufsize=10**8)
|
|
frame_size = self.width * self.height * 3
|
idx = 0
|
while True:
|
in_bytes = process.stdout.read(frame_size)
|
if not in_bytes:
|
break
|
frame = np.frombuffer(in_bytes, dtype=np.uint8).reshape((self.height, self.width, 3))
|
self.frame_queue.put((idx, frame)) # 改为 tuple 形式
|
idx += 1
|
|
self.total_frames = idx
|
print(f"_decode_video Total frames: {self.total_frames}")
|
|
process.stdout.close()
|
process.wait()
|
|
# 通知处理线程结束
|
for _ in range(self.num_processors):
|
self.frame_queue.put(None)
|
|
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、hevc_nvenc、h264_amf、h264_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()
|
