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.")