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()