import gradio as gr import numpy as np import cv2 from PIL import Image import torch import torch.nn.functional as F from transformers import SegformerImageProcessor, AutoModelForSemanticSegmentation import insightface from insightface.app import FaceAnalysis import onnxruntime as ort import os import tempfile import warnings from gfpgan import GFPGANer import argparse warnings.filterwarnings("ignore") class FaceHairSwapperGradio: def __init__(self): """Khởi tạo models cho Gradio app""" self.setup_models() def setup_models(self): """Setup các models cần thiết""" try: print("🔄 Đang tải SegFormer model...") self.processor = SegformerImageProcessor.from_pretrained("VanNguyen1214/get_face_and_hair") self.model = AutoModelForSemanticSegmentation.from_pretrained("VanNguyen1214/get_face_and_hair") print("✅ SegFormer model đã tải thành công!") print("🔄 Đang tải InsightFace model...") # Thử GPU trước, fallback CPU providers = ['CPUExecutionProvider'] if torch.cuda.is_available(): providers = ['CUDAExecutionProvider', 'CPUExecutionProvider'] try: self.face_app = FaceAnalysis(providers=providers) self.face_app.prepare(ctx_id=0, det_size=(640, 640)) print("✅ InsightFace đã tải thành công!") except Exception as e: print(f"⚠️ Lỗi tải InsightFace: {e}") # Thử với providers đơn giản hơn self.face_app = FaceAnalysis(providers=['CPUExecutionProvider']) self.face_app.prepare(ctx_id=-1, det_size=(320, 320)) print("✅ InsightFace đã tải thành công (CPU mode)!") # Load face swapper từ file local swapper_path = './models/inswapper_128.onnx' if os.path.exists(swapper_path): try: self.face_swapper = insightface.model_zoo.get_model(swapper_path, providers=providers) print("✅ Face Swapper đã tải từ file local!") except Exception as e: print(f"⚠️ Lỗi tải face swapper từ file local: {e}") # Thử tải từ model zoo try: self.face_swapper = insightface.model_zoo.get_model('inswapper_128.onnx', providers=providers) print("✅ Face Swapper đã tải từ model zoo!") except Exception as e2: print(f"❌ Không thể tải face swapper: {e2}") self.face_swapper = None else: print("⚠️ Không tìm thấy inswapper_128.onnx trong thư mục models, thử tải từ model zoo...") try: self.face_swapper = insightface.model_zoo.get_model('inswapper_128.onnx', providers=providers) print("✅ Face Swapper đã tải từ model zoo!") except Exception as e: print(f"❌ Không thể tải face swapper: {e}") self.face_swapper = None # Load GFPGAN print("🔄 Đang tải GFPGAN model...") gfpgan_path = './models/GFPGANv1.4.pth' if os.path.exists(gfpgan_path): try: self.gfpgan_enhancer = GFPGANer( model_path=gfpgan_path, upscale=2, arch='clean', channel_multiplier=2, bg_upsampler=None ) print("✅ GFPGAN model đã tải thành công!") except Exception as e: print(f"⚠️ Lỗi tải GFPGAN: {e}") self.gfpgan_enhancer = None else: print("⚠️ Không tìm thấy GFPGANv1.4.pth, bỏ qua GFPGAN enhancement") self.gfpgan_enhancer = None print("✅ Setup models hoàn tất!") except Exception as e: print(f"❌ Lỗi khi tải model: {e}") import traceback traceback.print_exc() raise e def enhance_face_gfpgan(self, image): """Tăng chất lượng khuôn mặt bằng GFPGAN""" if self.gfpgan_enhancer is None: return image try: # GFPGAN expects BGR format if len(image.shape) == 3 and image.shape[2] == 3: input_img = cv2.cvtColor(image, cv2.COLOR_RGB2BGR) else: input_img = image # Enhance with GFPGAN _, _, enhanced_img = self.gfpgan_enhancer.enhance( input_img, has_aligned=False, only_center_face=False, paste_back=True ) # Convert back to RGB if enhanced_img is not None: enhanced_img = cv2.cvtColor(enhanced_img, cv2.COLOR_BGR2RGB) return enhanced_img else: return image except Exception as e: print(f"GFPGAN enhancement failed: {e}") return image def get_face_mask_insightface(self, image, expand_ratio=0.3): """Lấy mask khuôn mặt từ InsightFace với độ chính xác cao hơn""" try: image_np = np.array(image) height, width = image_np.shape[:2] face_mask = np.zeros((height, width), dtype=np.uint8) faces = self.face_app.get(image_np) if len(faces) > 0: # Chọn khuôn mặt lớn nhất face = max(faces, key=lambda x: (x.bbox[2] - x.bbox[0]) * (x.bbox[3] - x.bbox[1])) # Sử dụng landmarks để tạo mask chính xác hơn if hasattr(face, 'kps') and face.kps is not None: try: landmarks = face.kps.astype(int) # Tạo convex hull từ landmarks hull = cv2.convexHull(landmarks) cv2.fillPoly(face_mask, [hull], 1) # Mở rộng mask if expand_ratio > 0: kernel_size = max(3, int(min(width, height) * expand_ratio * 0.1)) if kernel_size % 2 == 0: kernel_size += 1 kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size, kernel_size)) face_mask = cv2.dilate(face_mask, kernel, iterations=1) except Exception as e: print(f"Landmark method failed: {e}, using bbox method") # Fallback về bbox method face_mask = self._create_face_mask_from_bbox(face.bbox, width, height, expand_ratio) else: # Fallback về bbox method face_mask = self._create_face_mask_from_bbox(face.bbox, width, height, expand_ratio) # Làm mượt mask face_mask = cv2.GaussianBlur(face_mask.astype(np.float32), (15, 15), 0) face_mask = (face_mask > 0.1).astype(np.uint8) return face_mask except Exception as e: print(f"Error in get_face_mask_insightface: {e}") # Return empty mask if all fails return np.zeros((image.size[1], image.size[0]), dtype=np.uint8) def _create_face_mask_from_bbox(self, bbox, width, height, expand_ratio): """Helper method to create face mask from bounding box""" face_mask = np.zeros((height, width), dtype=np.uint8) x1, y1, x2, y2 = bbox.astype(int) center_x, center_y = (x1 + x2) // 2, (y1 + y2) // 2 w, h = x2 - x1, y2 - y1 new_w = int(w * (1 + expand_ratio)) new_h = int(h * (1 + expand_ratio)) new_x1 = max(0, center_x - new_w // 2) new_y1 = max(0, center_y - new_h // 2) new_x2 = min(width, center_x + new_w // 2) new_y2 = min(height, center_y + new_h // 2) if new_x2 > new_x1 and new_y2 > new_y1: mask_region = np.zeros((new_y2 - new_y1, new_x2 - new_x1), dtype=np.uint8) center_region_x = (new_x2 - new_x1) // 2 center_region_y = (new_y2 - new_y1) // 2 if center_region_x > 0 and center_region_y > 0: cv2.ellipse(mask_region, (center_region_x, center_region_y), (center_region_x, center_region_y), 0, 0, 360, 1, -1) face_mask[new_y1:new_y2, new_x1:new_x2] = mask_region return face_mask def get_hair_mask(self, image): """Lấy mask tóc từ SegFormer""" inputs = self.processor(images=image, return_tensors="pt") with torch.no_grad(): outputs = self.model(**inputs) logits = outputs.logits.cpu() upsampled_logits = F.interpolate( logits, size=image.size[::-1], mode="bilinear", align_corners=False, ) pred_seg = upsampled_logits.argmax(dim=1)[0].numpy() hair_mask = (pred_seg == 2).astype(np.uint8) return hair_mask def swap_face_insightface(self, source_image, target_image, enhance_result=True): """Hoán đổi khuôn mặt với InsightFace và tùy chọn enhance""" if self.face_swapper is None: print("⚠️ Face swapper không khả dụng, sử dụng blend method") return target_image try: source_faces = self.face_app.get(source_image) target_faces = self.face_app.get(target_image) if len(source_faces) == 0 or len(target_faces) == 0: print("⚠️ Không tìm thấy khuôn mặt trong một hoặc cả hai ảnh") return target_image # Chọn khuôn mặt lớn nhất source_face = max(source_faces, key=lambda x: (x.bbox[2] - x.bbox[0]) * (x.bbox[3] - x.bbox[1])) target_face = max(target_faces, key=lambda x: (x.bbox[2] - x.bbox[0]) * (x.bbox[3] - x.bbox[1])) # Thực hiện face swap result = self.face_swapper.get(target_image, target_face, source_face, paste_back=True) # Enhance kết quả với GFPGAN nếu có if enhance_result and self.gfpgan_enhancer is not None: result = self.enhance_face_gfpgan(result) return result except Exception as e: print(f"❌ Lỗi trong face swap: {e}") return target_image def blend_images_with_mask(self, base_image, overlay_image, mask, blur_kernel=5): """ Blend hai ảnh với mask được cải thiện. Đảm bảo mọi shape đều khớp nhau. Args: base_image: numpy.ndarray (H, W, 3) overlay_image: numpy.ndarray (H, W, 3) mask: numpy.ndarray (H, W) blur_kernel: int (làm mượt mask) Returns: numpy.ndarray (H, W, 3) """ # Đảm bảo overlay_image cùng shape với base_image if overlay_image.shape != base_image.shape: overlay_image = cv2.resize(overlay_image, (base_image.shape[1], base_image.shape[0])) # Đảm bảo mask cùng shape (chỉ cần 2 chiều height, width) if mask.shape[:2] != base_image.shape[:2]: mask = cv2.resize(mask, (base_image.shape[1], base_image.shape[0])) mask_float = mask.astype(np.float32) if blur_kernel > 0: mask_float = cv2.GaussianBlur(mask_float, (blur_kernel, blur_kernel), 0) mask_float = mask_float / mask_float.max() if mask_float.max() > 0 else mask_float mask_3d = np.stack([mask_float] * 3, axis=-1) result = base_image.astype(np.float32) * (1 - mask_3d) + overlay_image.astype(np.float32) * mask_3d return result.astype(np.uint8) def post_process_result(self, image): """Post-processing để cải thiện kết quả cuối cùng""" # Color correction lab = cv2.cvtColor(image, cv2.COLOR_RGB2LAB) l, a, b = cv2.split(lab) # Áp dụng CLAHE cho kênh L clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8)) l = clahe.apply(l) # Merge lại lab = cv2.merge([l, a, b]) result = cv2.cvtColor(lab, cv2.COLOR_LAB2RGB) return result def process_swap(self, full_face_image, hair_head_image, enhance_faces=True, progress=gr.Progress()): """ Hàm chính để xử lý face/hair swap cho Gradio với GFPGAN enhancement Args: full_face_image: ảnh có khuôn mặt (từ Gradio) hair_head_image: ảnh có tóc (từ Gradio) enhance_faces: có sử dụng GFPGAN không progress: Gradio progress bar Returns: PIL Image: kết quả """ if full_face_image is None or hair_head_image is None: return None, "❌ Vui lòng tải lên cả hai ảnh!" try: progress(0.1, desc="🔄 Đang xử lý ảnh...") # Convert to PIL nếu cần if isinstance(full_face_image, np.ndarray): full_face_image = Image.fromarray(full_face_image) if isinstance(hair_head_image, np.ndarray): hair_head_image = Image.fromarray(hair_head_image) # Chuyển về RGB if full_face_image.mode != 'RGB': full_face_image = full_face_image.convert('RGB') if hair_head_image.mode != 'RGB': hair_head_image = hair_head_image.convert('RGB') progress(0.2, desc="📏 Resize ảnh...") # Resize về cùng kích thước target_size = (512, 512) full_face_resized = full_face_image.resize(target_size, Image.Resampling.LANCZOS) hair_head_resized = hair_head_image.resize(target_size, Image.Resampling.LANCZOS) # Convert to numpy full_face_array = np.array(full_face_resized) hair_head_array = np.array(hair_head_resized) progress(0.3, desc="💇 Phân tích tóc...") hair_mask = self.get_hair_mask(hair_head_resized) progress(0.5, desc="👤 Phân tích khuôn mặt...") face_mask = self.get_face_mask_insightface(full_face_resized, expand_ratio=0.2) progress(0.7, desc="🔄 Thực hiện face swap...") try: face_swapped = self.swap_face_insightface( full_face_array, hair_head_array, enhance_result=enhance_faces ) except Exception as e: print(f"Face swap thất bại: {e}, dùng blend method...") face_swapped = self.blend_images_with_mask( hair_head_array, full_face_array, face_mask, blur_kernel=7 ) progress(0.85, desc="🎨 Kết hợp tóc...") final_result = self.blend_images_with_mask( face_swapped, hair_head_array, hair_mask, blur_kernel=3 ) progress(0.95, desc="✨ Post-processing...") final_result = self.post_process_result(final_result) progress(1.0, desc="✅ Hoàn thành!") result_image = Image.fromarray(final_result) return result_image, "✅ Hoàn thành thành công!" except Exception as e: error_msg = f"❌ Lỗi xử lý: {str(e)}" print(error_msg) import traceback traceback.print_exc() return None, error_msg # Khởi tạo model (chỉ một lần) print("🚀 Đang khởi tạo Face Hair Swapper với GFPGAN...") try: swapper = FaceHairSwapperGradio() print("✅ Khởi tạo thành công!") except Exception as e: print(f"❌ Lỗi khởi tạo: {e}") swapper = None def gradio_interface(full_face_img, hair_head_img, enhance_faces): """Interface function cho Gradio""" if swapper is None: return None, "❌ Model chưa được khởi tạo!" return swapper.process_swap(full_face_img, hair_head_img, enhance_faces) # Tạo Gradio Interface def create_gradio_app(): """Tạo Gradio app với enhanced features""" with gr.Blocks( title="Face & Hair Swap with GFPGAN", theme=gr.themes.Soft(), css=""" .gradio-container { max-width: 1200px !important; } .image-container { height: 400px !important; } """ ) as demo: gr.HTML("""

🎭 Face & Hair Swap với GFPGAN Enhancement

Hoán đổi khuôn mặt và tóc giữa hai bức ảnh sử dụng AI với chất lượng cao

""") with gr.Row(): with gr.Column(): gr.HTML("

📸 Ảnh có khuôn mặt

") full_face_input = gr.Image( label="Tải ảnh có khuôn mặt", type="pil", height=400 ) with gr.Column(): gr.HTML("

💇 Ảnh có tóc

") hair_head_input = gr.Image( label="Tải ảnh có tóc", type="pil", height=400 ) with gr.Row(): enhance_checkbox = gr.Checkbox( label="✨ Sử dụng GFPGAN để tăng chất lượng khuôn mặt", value=True, info="Tăng chất lượng và độ chi tiết của khuôn mặt trong kết quả" ) with gr.Row(): process_btn = gr.Button( "🚀 Bắt đầu hoán đổi", variant="primary", size="lg" ) with gr.Row(): with gr.Column(): result_output = gr.Image( label="Kết quả", type="pil", height=400 ) with gr.Column(): status_output = gr.Textbox( label="Trạng thái", lines=3, interactive=False ) # Event handlers process_btn.click( fn=gradio_interface, inputs=[full_face_input, hair_head_input, enhance_checkbox], outputs=[result_output, status_output], api_name="swap_face_hair" ) # Examples và hướng dẫn gr.HTML("

📋 Hướng dẫn sử dụng:

") gr.HTML(""" """) gr.HTML("

🔧 Thông tin kỹ thuật:

") gr.HTML(""" """) return demo # Chạy app if __name__ == "__main__": demo = create_gradio_app() demo.launch( server_name="0.0.0.0", server_port=7860, share=True, debug=True ) # Cho Hugging Face Spaces demo = create_gradio_app()