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