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.")
|
