import cv2
|
import numpy as np
|
import insightface
|
from insightface.app import FaceAnalysis
|
from insightface.app.common import Face
|
import threading
|
from gfpgan.utils import GFPGANer
|
from typing import Any
|
|
Frame = np.ndarray
|
|
|
class FaceProcessor:
|
"""
|
调用cv的模块分解mp4，组装mp4，视频质量差
|
"""
|
def __init__(self, face_enhancer_model: str = 'GFPGANv1.4.pth', face_swapper_model: str = 'inswapper_128.onnx', device: str = 'cuda'):
|
self.FACE_ENHANCER = None
|
self.FACE_SWAPPER = None
|
self.THREAD_SEMAPHORE = threading.Semaphore()
|
self.THREAD_LOCK = threading.Lock()
|
|
# Set device
|
self.device = device # 'cuda' for GPU, 'cpu' for CPU
|
|
# Load models
|
self.face_enhancer_model = face_enhancer_model
|
self.face_swapper_model = face_swapper_model
|
|
# Initialize face analysis app
|
self.app = FaceAnalysis(providers=['CUDAExecutionProvider' if device == 'cuda' else 'CPUExecutionProvider'])
|
self.app.prepare(ctx_id=0, det_size=(640, 640))
|
|
def get_face_enhancer(self) -> GFPGANer:
|
with self.THREAD_LOCK:
|
if self.FACE_ENHANCER is None:
|
self.FACE_ENHANCER = GFPGANer(
|
model_path=self.face_enhancer_model, upscale=1, device=self.device)
|
return self.FACE_ENHANCER
|
|
def get_face_swapper(self) -> Any:
|
with self.THREAD_LOCK:
|
if self.FACE_SWAPPER is None:
|
self.FACE_SWAPPER = insightface.model_zoo.get_model(
|
self.face_swapper_model, providers=['CUDAExecutionProvider' if self.device == 'cuda' else 'CPUExecutionProvider'])
|
return self.FACE_SWAPPER
|
|
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:
|
with self.THREAD_SEMAPHORE:
|
_, _, temp_face = self.get_face_enhancer().enhance(
|
temp_face,
|
paste_back=True
|
)
|
temp_frame[start_y:end_y, start_x:end_x] = temp_face
|
return temp_frame
|
|
def swap_face(self, source_face: Face, target_face: Face, temp_frame: Frame) -> Frame:
|
return self.get_face_swapper().get(temp_frame, target_face, source_face, paste_back=True)
|
|
def process_video(self, src_mp4: str, image_path: str, dst_mp4: str) -> None:
|
# Read the target image (source face)
|
target_image = cv2.imread(image_path)
|
target_faces = self.app.get(target_image)
|
if len(target_faces) == 0:
|
print("No face detected in the target image.")
|
return
|
target_face = target_faces[0]
|
|
# Open the source video
|
cap = cv2.VideoCapture(src_mp4)
|
if not cap.isOpened():
|
print("Error opening video stream or file")
|
return
|
|
# Get video properties
|
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
|
fps = cap.get(cv2.CAP_PROP_FPS)
|
frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
|
frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
|
|
# Create VideoWriter object to save the output video
|
out = cv2.VideoWriter(dst_mp4, fourcc, fps, (frame_width, frame_height))
|
|
# Get the total frame count
|
total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
|
|
# Process each frame of the video
|
frame_count = 0
|
while cap.isOpened():
|
ret, frame = cap.read()
|
if not ret:
|
break
|
|
# Detect faces in the current frame
|
faces = self.app.get(frame)
|
if len(faces) > 0:
|
# Swap face from the target image to the current frame
|
swapped_face_frame = self.swap_face(target_face, faces[0], frame)
|
self.enhance_face(faces[0], swapped_face_frame)
|
out.write(swapped_face_frame)
|
else:
|
# If no face detected, just write the original frame
|
out.write(frame)
|
|
# Increment frame count
|
frame_count += 1
|
|
# Calculate and print progress every 5%
|
progress = (frame_count / total_frames) * 100
|
if progress >= 5 * (frame_count // (total_frames // 20)):
|
print(f"Processing: {int(progress)}% completed")
|
|
cap.release()
|
out.release()
|
print(f"Video processing complete. Output saved to {dst_mp4}")
|
|
|
# Usage example:
|
|
# Initialize the FaceProcessor (use 'cuda' for GPU)
|
face_processor = FaceProcessor(device='cpu')
|
|
# Process the video by swapping faces from the target image
|
# face_processor.process_video('g1.mp4', 'from-face/girl-1.png', './g1-demo_output.mp4')
|
# face_processor.process_video('g1.mp4', 'from-face/girl-2.png', './g2-demo_output.mp4')
|
face_processor.process_video('g3.mp4', 'lilizhen.jpg', './lilizhen_output.mp4')
|
face_processor.process_video('g3.mp4', 'from-face/girl-1.png', './g1-demo_output.mp4')
|
face_processor.process_video('g3.mp4', 'from-face/girl-2.png', './g2-demo_output.mp4')
|
