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