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 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 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):
|
return 'cpu'
|
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.")
|
