add face swaper

wangzhibo

2025-07-30 c706a69abfe9bae1322adbd344dfdaae37cdf564

add face swaper

 auto_media_publisher/modifiers/face_swap/__init__.py

 auto_media_publisher/modifiers/face_swap/face_analyser.py

New file
@@ -0,0 +1,53 @@

import torch
from typing import Optional, List
import numpy
from insightface.app import FaceAnalysis
from insightface.app.common import Face

from gfpgan.utils import GFPGANer
from .typing import Frame # 你已有定义

SIMILAR_FACE_DISTANCE: float = 0.45     # < 0.4 同一人 0.4 -0.6 大概是同一人 >0.6 不是同一个


class FaceAnalyser:
    def __init__(self):
        self.device = self._select_device()
        self._analyser = FaceAnalysis(name='buffalo_l', root='./', providers=['CUDAExecutionProvider' if self.device == 'cuda' else 'CPUExecutionProvider'])
        self._analyser.prepare(ctx_id=0, det_size=(640, 640))

    def _select_device(self):
        if torch.cuda.is_available():
            return 'cuda'
        elif torch.backends.mps.is_available():
            return 'mps'
        else:
            return 'cpu'
        
    def get_many_faces(self,frame: Frame) -> Optional[List[Face]]:
        try:
            return self._analyser.get(frame)
        except Exception as e:
            print(f"{e}")
            return None

    def get_one_face(self,frame: Frame, position: int = 0) -> Optional[Face]:
        many_faces = self.get_many_faces(frame)
        if many_faces:
            try:
                return many_faces[position]
            except IndexError:
                return many_faces[-1]
        return None


    def find_similar_face(self,frame: Frame, reference_face: Face) -> Optional[Face]:
        many_faces = self.get_many_faces(frame)
        if many_faces:
            for face in many_faces:
                if hasattr(face, 'normed_embedding') and hasattr(reference_face, 'normed_embedding'):
                    distance = numpy.sum(numpy.square(face.normed_embedding - reference_face.normed_embedding))
                    if distance < SIMILAR_FACE_DISTANCE:
                        return face
        return None

 auto_media_publisher/modifiers/face_swap/face_enhancer.py

New file
@@ -0,0 +1,49 @@

import torch
from typing import List
from insightface.app.common import Face

from gfpgan.utils import GFPGANer
from .typing import Frame # 你已有定义
from auto_media_publisher.config.conf_base import FACE_ENHANCER_MODEL_PATH

SIMILAR_FACE_DISTANCE: float = 0.45     # < 0.4 同一人 0.4 -0.6 大概是同一人 >0.6 不是同一个


class FaceEnhancer:
    def __init__(self):
        self.device = self._select_device()
        self._enhancer = GFPGANer(model_path=FACE_ENHANCER_MODEL_PATH, upscale=1, device=self.device)

    def _select_device(self):
        if torch.cuda.is_available():
            return 'cuda'
        elif torch.backends.mps.is_available():
            return 'mps'
        else:
            return 'cpu'
        
    def _enhance_face(self,target_face: Face, temp_frame: Frame) -> Frame:
        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._enhancer.enhance(
                temp_face,
                paste_back=True
            )
            temp_frame[start_y:end_y, start_x:end_x] = temp_face
        return temp_frame


    def process_frame(self,temp_frame: Frame,many_faces:List) -> Frame:
        #many_faces = self._get_many_faces(temp_frame)  check face_analyser.py
        if many_faces:
            for target_face in many_faces:
                temp_frame = self._enhance_face(target_face, temp_frame)
        return temp_frame

 auto_media_publisher/modifiers/face_swap/face_nsfw_predictor.py

New file
@@ -0,0 +1,28 @@
import threading
import numpy as np
from PIL import Image
from keras import Model
import opennsfw2

from .typing import Frame  # 你已有定义

MAX_PROBABILITY = 0.85

class NSFWPredictor:
    def __init__(self):
        self._predictor = opennsfw2.make_open_nsfw_model()

    def predict_frame(self, frame: Frame) -> bool:
        image = Image.fromarray(frame)
        image = opennsfw2.preprocess_image(image, opennsfw2.Preprocessing.YAHOO)
        views = np.expand_dims(image, axis=0)
        _, probability = self._predictor.predict(views)[0]
        return probability > MAX_PROBABILITY

    def predict_image(self, image_path: str) -> bool:
        probability = opennsfw2.predict_image(image_path)
        return probability > MAX_PROBABILITY

    def predict_video(self, video_path: str, frame_interval: int = 100) -> bool:
        _, probabilities = opennsfw2.predict_video_frames(video_path, frame_interval=frame_interval)
        return any(prob > MAX_PROBABILITY for prob in probabilities)

 auto_media_publisher/modifiers/face_swap/face_swaper_avpy.py

New file
@@ -0,0 +1,166 @@
import av
import torch
import threading, queue
from threading import current_thread
import numpy as np
import insightface
import cv2
import time
from insightface.app.common import Face
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 insightface.app import FaceAnalysis

SIMILAR_FACE_DISTANCE: float = 0.45     # < 0.4 同一人 0.4 -0.6 大概是同一人 >0.6 不是同一个

def _select_device():
    if torch.cuda.is_available():
        return 'cuda:0'
    elif torch.backends.mps.is_available():
        return 'mps'
    else:
        return 'cpu'
    
NUM_WORKERS = 1

#face_swapper = insightface.model_zoo.get_model(FACE_SWAPPER_MODEL_PATH, providers=['CUDAExecutionProvider'])
#face_enhancer = GFPGANer(model_path=FACE_ENHANCER_MODEL_PATH, upscale=1, device=_select_device())

frame_q = queue.Queue(maxsize=10)
processed_q = queue.Queue(maxsize=10)

THREAD_SEMAPHORE =threading.Semaphore(4)  # 不限制了。

def enhance_face(thread_face_enhancer,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]
    #print(temp_face.size)
    if temp_face.size:
        _, _, temp_face = thread_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(video_path):
    container = av.open(video_path)
    for idx, frame in enumerate(container.decode(video=0)):
        img = frame.to_ndarray(format='rgb24')
        frame_q.put((idx, img))
    for _ in range(NUM_WORKERS):
        frame_q.put((None, None))

def process_frame_loop(new_face):
    start_time = time.time()
    thread_face_analyser = FaceAnalysis(name='buffalo_l', root='./', providers=['CUDAExecutionProvider'])
    thread_face_analyser.prepare(ctx_id=3, det_size=(640, 640))
    #thread_face_enhancer = GFPGANer(model_path=FACE_ENHANCER_MODEL_PATH, upscale=1, device=_select_device())
    thread_face_enhancer = GFPGANer(model_path=FACE_ENHANCER_MODEL_PATH, upscale=1, device=_select_device())
    thread_face_swapper = insightface.model_zoo.get_model(FACE_SWAPPER_MODEL_PATH, providers=['CUDAExecutionProvider'])

    while True:
        idx, frame = frame_q.get()
        if idx is None:
            processed_q.put((None, None))
            break
        try:
            many_faces = thread_face_analyser.get(frame)
            if many_faces:
                old_face = many_faces[0]
                frame = thread_face_swapper.get(frame, old_face, new_face, paste_back=True)

                for sinlge_face in many_faces:
                    frame_ = enhance_face(thread_face_enhancer,sinlge_face, frame)                

            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 encode_video(output_path, fps, width, height):
    out = av.open(output_path, mode='w')

    #stream = out.add_stream('h264_nvenc', rate=fps)
    stream = out.add_stream('libx264', rate=fps)
    stream.width = width
    stream.height = height
    stream.pix_fmt = 'yuv420p'

    stream.options = {'crf': '18', 'preset': 'veryfast'}
    #stream.options = {'rc': 'constqp', 'qp': '18','preset': 'p7','profile': 'high'}

    next_index = 0
    buffer = {}
    end_count = 0

    while True:
        idx, frame = processed_q.get()
        if idx is None:
            end_count += 1
            if end_count >= NUM_WORKERS:
                break
            continue
        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 stream.encode(video_frame):
                out.mux(packet)
            next_index += 1

    # 清空编码缓存
    for packet in stream.encode():
        out.mux(packet)
    out.close()
# 启动任务
video_path = './temp/buhui.mp4'
output_path = './temp/buhui-wangrong-enhancer.mp4'
new_face_image_path = './temp/wangrong.jpg'
face_analyser = FaceAnalysis(name='buffalo_l', root='./', providers=['CUDAExecutionProvider'])
face_analyser.prepare(ctx_id=0, det_size=(640, 640))

faces = face_analyser.get(cv2.imread(new_face_image_path))
if faces:
    new_face = sorted(faces, key=lambda x: x.bbox[0])[0]
        
    #threading.Thread(target=decode_video, args=(video_path,)).start()
    #threading.Thread(target=process_frame_loop, args=(new_face,)).start()
    #threading.Thread(target=encode_video, args=(output_path, 25,576,1080)).start()
    threads = []
        
    # 解码线程
    t_decode = threading.Thread(target=decode_video,args=(video_path,))
    threads.append(t_decode)
    t_decode.start()
    
    # 多处理线程
    for i in range(NUM_WORKERS):
        t = threading.Thread(target=process_frame_loop, args=(new_face,),name=f"process_worker_{i}")
        threads.append(t)
        t.start()
    
    # 编码线程
    t_encode = threading.Thread(target=encode_video,args=(output_path, 25,576,1080))
    threads.append(t_encode)
    t_encode.start()
    
    for t in threads:
        t.join()
    
    print("Processing complete.")



 auto_media_publisher/modifiers/face_swap/face_swaper_ffmpeg.py

New file
@@ -0,0 +1,238 @@
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()

 auto_media_publisher/modifiers/face_swap/face_swaper_ffmpeg_stream.py

New file
@@ -0,0 +1,241 @@
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.")


 auto_media_publisher/modifiers/face_swap/typing.py

New file
@@ -0,0 +1,7 @@

from typing import Any
from insightface.app.common import Face
import numpy

Face = Face
Frame = numpy.ndarray[Any, Any]

 auto_media_publisher/utils/media_input_output.py

New file
@@ -0,0 +1,93 @@
import av
import os
import pims
import numpy as np
from torch.utils.data import Dataset
from torchvision.transforms.functional import to_pil_image
from PIL import Image


class VideoReader(Dataset):
    def __init__(self, path, transform=None):
        self.video = pims.PyAVVideoReader(path)
        self.rate = self.video.frame_rate
        self.transform = transform
        
    @property
    def frame_rate(self):
        return self.rate
        
    def __len__(self):
        return len(self.video)
        
    def __getitem__(self, idx):
        frame = self.video[idx]
        frame = Image.fromarray(np.asarray(frame))
        if self.transform is not None:
            frame = self.transform(frame)
        return frame


class VideoWriter:
    def __init__(self, path, frame_rate, bit_rate=1000000):
        self.container = av.open(path, mode='w')
        self.stream = self.container.add_stream('h264', rate=round(frame_rate))
        self.stream.pix_fmt = 'yuv420p'
        self.stream.bit_rate = bit_rate
        self.stream.codec_context.options = {
            'crf': '18',  # 控制画质，越小越清晰（18是高质量）
            'preset': 'veryfast',     # 影响编码速度，medium 是1x，fast是2x faster是3x veryfast是4x superfast是6x ultrafast是10x，文件体积大约会增加(以medium为基准) 10%，15%，20% 30% 50%
            'tune': 'film',
        }
    
    def write(self, frames):
        # frames: [T, C, H, W]
        self.stream.width = frames.size(3)
        self.stream.height = frames.size(2)
        if frames.size(1) == 1:
            frames = frames.repeat(1, 3, 1, 1) # convert grayscale to RGB
        frames = frames.mul(255).byte().cpu().permute(0, 2, 3, 1).numpy()
        for t in range(frames.shape[0]):
            frame = frames[t]
            frame = av.VideoFrame.from_ndarray(frame, format='rgb24')
            self.container.mux(self.stream.encode(frame))
                
    def close(self):
        self.container.mux(self.stream.encode())
        self.container.close()


class ImageSequenceReader(Dataset):
    def __init__(self, path, transform=None):
        self.path = path
        self.files = sorted(os.listdir(path))
        self.transform = transform
        
    def __len__(self):
        return len(self.files)
    
    def __getitem__(self, idx):
        with Image.open(os.path.join(self.path, self.files[idx])) as img:
            img.load()
        if self.transform is not None:
            return self.transform(img)
        return img


class ImageSequenceWriter:
    def __init__(self, path, extension='jpg'):
        self.path = path
        self.extension = extension
        self.counter = 0
        os.makedirs(path, exist_ok=True)
    
    def write(self, frames):
        # frames: [T, C, H, W]
        for t in range(frames.shape[0]):
            to_pil_image(frames[t]).save(os.path.join(
                self.path, str(self.counter).zfill(4) + '.' + self.extension))
            self.counter += 1
            
    def close(self):
        pass