import av
|
import torch
|
import threading, queue
|
from threading import current_thread
|
import numpy as np
|
import subprocess
|
import json
|
import cv2
|
import time
|
import insightface
|
from insightface.app.common import Face
|
from insightface.app import FaceAnalysis
|
from gfpgan.utils import GFPGANer
|
from .typing import Frame
|
from auto_media_publisher.config.conf_base import FACE_ENHANCER_MODEL_PATH,FACE_SWAPPER_MODEL_PATH
|
from auto_media_publisher.utils.logger import get_logger
|
import os
|
from fractions import Fraction
|
|
SIMILAR_FACE_DISTANCE: float = 0.45 # < 0.4 同一人 0.4 -0.6 大概是同一人 >0.6 不是同一个
|
|
class FaceSwapperEnhancer:
|
|
def __init__(self):
|
self.logger = get_logger("FaceSwapperEnhancer","admin")
|
self.frame_q = None #queue.Queue(maxsize=10)
|
self.processed_q = None #queue.Queue(maxsize=10)
|
self.device =self._select_device()
|
self.with_swapper =False
|
self.with_enhancer =False
|
self.face_analyser = FaceAnalysis(name='buffalo_l', root='./', providers=['CUDAExecutionProvider'])
|
self.face_analyser.prepare(ctx_id=0, det_size=(640, 640))
|
self.face_swapper = insightface.model_zoo.get_model(FACE_SWAPPER_MODEL_PATH, providers=['CUDAExecutionProvider'])
|
self.face_enhancer = GFPGANer(model_path=FACE_ENHANCER_MODEL_PATH, upscale=1, device=self.device)
|
|
|
def _select_device(self):
|
if torch.cuda.is_available():
|
return 'cuda:0'
|
elif torch.backends.mps.is_available():
|
return 'mps'
|
else:
|
return 'cpu'
|
|
|
|
def enhance_face(self,target_face, temp_frame:np.ndarray) -> np.ndarray:
|
start_x, start_y, end_x, end_y = map(int, target_face['bbox'])
|
padding_x = int((end_x - start_x) * 0.5)
|
padding_y = int((end_y - start_y) * 0.5)
|
start_x = max(0, start_x - padding_x)
|
start_y = max(0, start_y - padding_y)
|
end_x = max(0, end_x + padding_x)
|
end_y = max(0, end_y + padding_y)
|
temp_face = temp_frame[start_y:end_y, start_x:end_x]
|
|
if temp_face.size:
|
_, _, temp_face = self.face_enhancer.enhance(
|
temp_face,
|
paste_back=True
|
)
|
temp_frame[start_y:end_y, start_x:end_x] = temp_face
|
return temp_frame
|
|
|
def decode_video(self,video_path):
|
container = av.open(video_path)
|
for idx, frame in enumerate(container.decode(video=0)):
|
img = frame.to_ndarray(format='rgb24')
|
self.frame_q.put((idx, img))
|
self.frame_q.put((None, None))
|
|
def process_frame(self,new_face=None):
|
start_time = time.time()
|
while True:
|
idx, frame = self.frame_q.get()
|
if idx is None:
|
self.processed_q.put((None, None))
|
break
|
try:
|
many_faces = self.face_analyser.get(frame)
|
if many_faces:
|
if self.with_swapper:
|
old_face = many_faces[0]
|
frame = self.face_swapper.get(frame, old_face, new_face, paste_back=True)
|
|
if self.with_enhancer:
|
for sinlge_face in many_faces:
|
frame = self.enhance_face(sinlge_face, frame)
|
|
self.processed_q.put((idx, frame))
|
print(f"{current_thread().name} has proceesed frame :{idx}")
|
|
except Exception as e:
|
print(f"[ERROR] Frame {idx} failed: {e}")
|
|
print(f"totol time is : {time.time()- start_time}")
|
|
def probe_video(self,old_video_path):
|
cmd = [
|
"ffprobe",
|
"-v", "error",
|
"-select_streams", "v:0",
|
"-show_entries", "stream=width,height,r_frame_rate",
|
"-of", "json",
|
old_video_path,
|
]
|
proc = subprocess.run(cmd, capture_output=True, text=True, check=True)
|
info = json.loads(proc.stdout)
|
stream = info["streams"][0]
|
fps = eval(stream["r_frame_rate"])
|
fps_fraction = Fraction(fps).limit_denominator()
|
return fps_fraction,int(stream["width"]),int(stream["height"])
|
|
def encode_video(self,output_path, fps, width, height,audio_source_path):
|
|
audio_in_container = av.open(audio_source_path)
|
audio_stream_in = None
|
if audio_in_container.streams.audio:
|
audio_stream_in = audio_in_container.streams.audio[0]
|
|
out_container = av.open(output_path, mode='w')
|
|
v_stream = out_container.add_stream('libx264', rate=fps)
|
v_stream.width = width
|
v_stream.height = height
|
v_stream.pix_fmt = 'yuv420p'
|
v_stream.options = {'crf': '18', 'preset': 'veryfast'}
|
|
# 添加音频流，基于输入音频流模板
|
audio_out_stream = None
|
if audio_stream_in:
|
audio_out_stream = out_container.add_stream(audio_stream_in.codec.name)
|
'''
|
audio_codec_name = audio_stream_in.codec.name
|
audio_out_stream = out_container.add_stream(
|
audio_codec_name,
|
rate=audio_stream_in.rate or audio_stream_in.sample_rate
|
)
|
# 设置一些可写的 metadata（layout）
|
try:
|
audio_out_stream.layout = audio_stream_in.layout
|
except Exception:
|
pass
|
'''
|
else:
|
audio_out_stream = None
|
|
next_index = 0
|
buffer = {}
|
|
while True:
|
idx, frame = self.processed_q.get()
|
if idx is None:
|
break
|
buffer[idx] = frame
|
|
# 只要有当前要写入的帧，就输出
|
while next_index in buffer:
|
frame = buffer.pop(next_index)
|
video_frame = av.VideoFrame.from_ndarray(frame, format='rgb24')
|
for packet in v_stream.encode(video_frame):
|
out_container.mux(packet)
|
next_index += 1
|
|
# 清空编码缓存
|
for packet in v_stream.encode():
|
out_container.mux(packet)
|
|
# 复制原音频流（直接拷贝，不解码）
|
# 拷贝音频包，修改流指向
|
if audio_stream_in and audio_out_stream:
|
for packet in audio_in_container.demux(audio_stream_in):
|
if packet.dts is None:
|
continue
|
packet.stream = audio_out_stream # 直接改，不用 copy()
|
out_container.mux(packet)
|
|
out_container.close()
|
audio_in_container.close()
|
|
def enhance(self,old_video_path,output_video_path):
|
try:
|
if not os.path.exists(old_video_path):
|
self.logger.error(f"[face swap] input video file is not exist. {old_video_path}")
|
return False
|
|
self.frame_q = queue.Queue(maxsize=10)
|
self.processed_q = queue.Queue(maxsize=10)
|
self.with_swapper = False
|
self.with_enhancer = True
|
|
self.start_threads(old_video_path,output_video_path,None)
|
|
return True
|
except Exception as e:
|
self.logger.error(f"[ error ]. {e}")
|
return False
|
|
def start_threads(self,old_video_path,output_video_path,new_face=None):
|
threads = []
|
# 解码线程
|
t_decode = threading.Thread(target=self.decode_video,args=(old_video_path,))
|
threads.append(t_decode)
|
t_decode.start()
|
|
# 处理线程 多处理线程 anaylse enhancer 可能会出问题，暂时还是单线程
|
t_process = threading.Thread(target=self.process_frame, args=(new_face,))
|
threads.append(t_process)
|
t_process.start()
|
|
fps,width,height =self.probe_video(old_video_path)
|
|
# 编码线程
|
t_encode = threading.Thread(target=self.encode_video,args=(output_video_path, fps,width,height,old_video_path))
|
threads.append(t_encode)
|
t_encode.start()
|
|
for t in threads:
|
t.join()
|
|
def swap(self,old_video_path,output_video_path,new_face_image_path,with_enhancer=True):
|
try:
|
if not os.path.exists(old_video_path):
|
self.logger.error(f"[face swap] input video file is not exist. {old_video_path}")
|
return False
|
if not os.path.exists(new_face_image_path):
|
self.logger.error(f"[face swap] input new face image file is not exist. {new_face_image_path}")
|
return False
|
faces = self.face_analyser.get(cv2.imread(new_face_image_path))
|
if not faces:
|
self.logger.error(f"[face swap] input new face image file DOES NOT HAVE FACE. {new_face_image_path}")
|
return False
|
new_face = sorted(faces, key=lambda x: x.bbox[0])[0]
|
self.frame_q = queue.Queue(maxsize=10)
|
self.processed_q = queue.Queue(maxsize=10)
|
self.with_swapper = True
|
self.with_enhancer = with_enhancer
|
|
self.start_threads(old_video_path,output_video_path,new_face)
|
|
return True
|
except Exception as e:
|
self.logger.error(f"[ error ]. {e}")
|
return False
|
|
