admin/LC_NeoRefacer
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

first commit

Browse Source
This commit is contained in:
zhh
2025-10-17 16:04:42 +08:00
parent ad4dccfe49
commit d8d558e185
11 changed files with 510 additions and 88 deletions
Download Patch File Download Diff File Expand all files Collapse all files

178
refacer.py
View File

@@ -1,6 +1,9 @@
import cv2
import onnxruntime as rt
import sys
from utils.minio_client import oss_get_image, minio_client
sys.path.insert(1, './recognition')
from scrfd import SCRFD
from arcface_onnx import ArcFaceONNX
@@ -40,9 +43,11 @@ if sys.platform in ("win32", "win64"):
if hasattr(rt, "preload_dlls"):
rt.preload_dlls()
class RefacerMode(Enum):
CPU, CUDA, COREML, TENSORRT = range(1, 5)
class Refacer:
def __init__(self, force_cpu=False, colab_performance=False):
self.disable_similarity = False
@@ -55,46 +60,45 @@ class Refacer:
self.__check_providers()
self.total_mem = psutil.virtual_memory().total
self.__init_apps()
def _partial_face_blend(self, original_frame, swapped_frame, face):
h_frame, w_frame = original_frame.shape[:2]
x1, y1, x2, y2 = map(int, face.bbox)
x1 = max(0, min(x1, w_frame-1))
y1 = max(0, min(y1, h_frame-1))
x1 = max(0, min(x1, w_frame - 1))
y1 = max(0, min(y1, h_frame - 1))
x2 = max(0, min(x2, w_frame))
y2 = max(0, min(y2, h_frame))
if x2 <= x1 or y2 <= y1:
print(f"Invalid bbox: {x1},{y1},{x2},{y2}")
return swapped_frame
w = x2 - x1
h = y2 - y1
cutoff = int(h * (1.0 - self.blend_height_ratio))
swap_crop = swapped_frame[y1:y2, x1:x2].copy()
orig_crop = original_frame[y1:y2, x1:x2].copy()
mask = np.ones((h, w, 3), dtype=np.float32)
transition = 40
if cutoff < h:
blend_start = max(cutoff - transition // 2, 0)
blend_end = min(cutoff + transition // 2, h)
if blend_end > blend_start:
alpha = np.linspace(1.0, 0.0, blend_end - blend_start)[:, np.newaxis, np.newaxis]
mask[blend_start:blend_end, :, :] = alpha
mask[blend_end:, :, :] = 0.0
blended_crop = (swap_crop.astype(np.float32) * mask + orig_crop.astype(np.float32) * (1.0 - mask)).astype(np.uint8)
blended_frame = swapped_frame.copy()
blended_frame[y1:y2, x1:x2] = blended_crop
return blended_frame
def __download_with_progress(self, url, output_path):
response = requests.get(url, stream=True)
@@ -238,7 +242,7 @@ class Refacer:
faces = self.__get_faces(frame, max_num=0)
if not faces:
return frame
if self.disable_similarity:
for face in faces:
swapped = self.face_swapper.get(frame, face, self.replacement_faces[0][1], paste_back=True)
@@ -253,9 +257,9 @@ class Refacer:
faces = self.__get_faces(frame, max_num=0)
if not faces:
return frame
faces = sorted(faces, key=lambda face: face.bbox[0])
if self.multiple_faces_mode:
for idx, face in enumerate(faces):
if idx >= len(self.replacement_faces):
@@ -309,33 +313,33 @@ class Refacer:
original_name = osp.splitext(osp.basename(video_path))[0]
timestamp = str(int(time.time()))
filename = f"{original_name}_preview.mp4" if preview else f"{original_name}_{timestamp}.mp4"
self.__check_video_has_audio(video_path)
if preview:
os.makedirs("output/preview", exist_ok=True)
output_video_path = os.path.join('output', 'preview', filename)
else:
os.makedirs("output", exist_ok=True)
output_video_path = os.path.join('output', filename)
self.prepare_faces(faces, disable_similarity=disable_similarity, multiple_faces_mode=multiple_faces_mode)
self.first_face = False if multiple_faces_mode else (faces[0].get("origin") is None or disable_similarity)
self.partial_reface_ratio = partial_reface_ratio
cap = cv2.VideoCapture(video_path, cv2.CAP_FFMPEG)
total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
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))
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
output = cv2.VideoWriter(output_video_path, fourcc, fps, (frame_width, frame_height))
frames = []
frame_index = 0
skip_rate = 10 if preview else 1
with tqdm(total=total_frames, desc="Extracting frames") as pbar:
while cap.isOpened():
flag, frame = cap.read()
@@ -349,28 +353,24 @@ class Refacer:
gc.collect()
frame_index += 1
pbar.update()
cap.release()
if frames:
self.reface_group(faces, frames, output)
output.release()
converted_path = self.__convert_video(video_path, output_video_path, preview=preview)
if video_path.lower().endswith(".gif"):
if preview:
gif_output_path = os.path.join("output", "preview", os.path.basename(converted_path).replace(".mp4", ".gif"))
else:
gif_output_path = os.path.join("output", "gifs", os.path.basename(converted_path).replace(".mp4", ".gif"))
self.__generate_gif(converted_path, gif_output_path)
return converted_path, gif_output_path
return converted_path, None
return converted_path, None
def __generate_gif(self, video_path, gif_output_path):
os.makedirs(os.path.dirname(gif_output_path), exist_ok=True)
@@ -396,60 +396,64 @@ class Refacer:
return new_path
def reface_image(self, image_path, faces, disable_similarity=False, multiple_faces_mode=False, partial_reface_ratio=0.0):
self.prepare_faces(faces, disable_similarity=disable_similarity, multiple_faces_mode=multiple_faces_mode)
self.first_face = False if multiple_faces_mode else (faces[0].get("origin") is None or disable_similarity)
self.partial_reface_ratio = partial_reface_ratio
ext = osp.splitext(image_path)[1].lower()
os.makedirs("output", exist_ok=True)
original_name = osp.splitext(osp.basename(image_path))[0]
timestamp = str(int(time.time()))
if ext in ['.tif', '.tiff']:
pil_img = Image.open(image_path)
frames = []
page_count = 0
try:
while True:
pil_img.seek(page_count)
page_count += 1
except EOFError:
pass
pil_img = Image.open(image_path)
with tqdm(total=page_count, desc="Processing TIFF pages") as pbar:
for page in range(page_count):
pil_img.seek(page)
bgr_image = cv2.cvtColor(np.array(pil_img.convert('RGB')), cv2.COLOR_RGB2BGR)
refaced_bgr = self.process_first_face(bgr_image.copy()) if self.first_face else self.process_faces(bgr_image.copy())
enhanced_bgr = enhance_image_memory(refaced_bgr)
enhanced_rgb = cv2.cvtColor(enhanced_bgr, cv2.COLOR_BGR2RGB)
enhanced_pil = Image.fromarray(enhanced_rgb)
frames.append(enhanced_pil)
pbar.update(1)
output_path = os.path.join("output", f"{original_name}_{timestamp}.tif")
frames[0].save(output_path, save_all=True, append_images=frames[1:], compression="tiff_deflate")
print(f"Saved multipage refaced TIFF to {output_path}")
return output_path
else:
bgr_image = cv2.imread(image_path)
if bgr_image is None:
raise ValueError("Failed to read input image")
refaced_bgr = self.process_first_face(bgr_image.copy()) if self.first_face else self.process_faces(bgr_image.copy())
refaced_rgb = cv2.cvtColor(refaced_bgr, cv2.COLOR_BGR2RGB)
pil_img = Image.fromarray(refaced_rgb)
filename = f"{original_name}_{timestamp}.jpg"
output_path = os.path.join("output", filename)
pil_img.save(output_path, format='JPEG', quality=100, subsampling=0)
output_path = enhance_image(output_path)
print(f"Saved refaced image to {output_path}")
return output_path
self.prepare_faces(faces, disable_similarity=disable_similarity, multiple_faces_mode=multiple_faces_mode)
self.first_face = False if multiple_faces_mode else (faces[0].get("origin") is None or disable_similarity)
self.partial_reface_ratio = partial_reface_ratio
ext = osp.splitext(image_path)[1].lower() #
# ext = image_path.rsplit('.',1)[1].lower()
os.makedirs("output", exist_ok=True) #
original_name = osp.splitext(osp.basename(image_path))[0]
timestamp = str(int(time.time()))
if ext in ['.tif', '.tiff']:
pil_img = Image.open(image_path) #
# pil_img = oss_get_image(oss_client=minio_client, path=image_path, data_type="PIL")
frames = []
page_count = 0
try:
while True:
pil_img.seek(page_count)
page_count += 1
except EOFError:
pass
pil_img = Image.open(image_path) #
# pil_img = oss_get_image(oss_client=minio_client, path=image_path, data_type="PIL")
with tqdm(total=page_count, desc="Processing TIFF pages") as pbar:
for page in range(page_count):
pil_img.seek(page)
bgr_image = cv2.cvtColor(np.array(pil_img.convert('RGB')), cv2.COLOR_RGB2BGR)
refaced_bgr = self.process_first_face(bgr_image.copy()) if self.first_face else self.process_faces(bgr_image.copy())
enhanced_bgr = enhance_image_memory(refaced_bgr)
enhanced_rgb = cv2.cvtColor(enhanced_bgr, cv2.COLOR_BGR2RGB)
enhanced_pil = Image.fromarray(enhanced_rgb)
frames.append(enhanced_pil)
pbar.update(1)
output_path = os.path.join("output", f"{original_name}_{timestamp}.tif")
frames[0].save(output_path, save_all=True, append_images=frames[1:], compression="tiff_deflate")
print(f"Saved multipage refaced TIFF to {output_path}")
return output_path
else:
bgr_image = cv2.imread(image_path) #
# bgr_image = oss_get_image(oss_client=minio_client, path=image_path, data_type="cv2")
if bgr_image is None:
raise ValueError("Failed to read input image")
refaced_bgr = self.process_first_face(bgr_image.copy()) if self.first_face else self.process_faces(bgr_image.copy())
refaced_rgb = cv2.cvtColor(refaced_bgr, cv2.COLOR_BGR2RGB)
pil_img = Image.fromarray(refaced_rgb)
filename = f"{original_name}_{timestamp}.jpg"
output_path = os.path.join("output", filename)
pil_img.save(output_path, format='JPEG', quality=100, subsampling=0)
output_path = enhance_image(output_path)
print(f"Saved refaced image to {output_path}")
return output_path
def extract_faces_from_image(self, image_path, max_faces=5):
frame = cv2.imread(image_path)
@@ -508,4 +512,4 @@ class Refacer:
'h264_videotoolbox': '0',
'h264_nvenc': '0',
'libx264': '0'
}
}