Spaces:
Running
Running
| 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(""" | |
| <div style="text-align: center; padding: 20px;"> | |
| <h1>🎭 Face & Hair Swap với GFPGAN Enhancement</h1> | |
| <p>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</p> | |
| </div> | |
| """) | |
| with gr.Row(): | |
| with gr.Column(): | |
| gr.HTML("<h3>📸 Ảnh có khuôn mặt</h3>") | |
| full_face_input = gr.Image( | |
| label="Tải ảnh có khuôn mặt", | |
| type="pil", | |
| height=400 | |
| ) | |
| with gr.Column(): | |
| gr.HTML("<h3>💇 Ảnh có tóc</h3>") | |
| 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("<h3>📋 Hướng dẫn sử dụng:</h3>") | |
| gr.HTML(""" | |
| <ul> | |
| <li><strong>Ảnh khuôn mặt:</strong> Tải ảnh có khuôn mặt rõ ràng, nhìn thẳng</li> | |
| <li><strong>Ảnh tóc:</strong> Tải ảnh có kiểu tóc đẹp mong muốn</li> | |
| <li><strong>GFPGAN Enhancement:</strong> Bật để tăng chất lượng khuôn mặt (khuyến nghị)</li> | |
| <li><strong>Chất lượng tốt nhất:</strong> Ảnh 512x512px, ánh sáng đều, không bị mờ</li> | |
| <li><strong>Models sử dụng:</strong> InsightFace + SegFormer + GFPGAN</li> | |
| </ul> | |
| """) | |
| gr.HTML("<h3>🔧 Thông tin kỹ thuật:</h3>") | |
| gr.HTML(""" | |
| <ul> | |
| <li><strong>Face Swap:</strong> inswapper_128.onnx (InsightFace)</li> | |
| <li><strong>Face Enhancement:</strong> GFPGANv1.4.pth</li> | |
| <li><strong>Hair Segmentation:</strong> SegFormer</li> | |
| <li><strong>Post-processing:</strong> CLAHE, Color Correction</li> | |
| </ul> | |
| """) | |
| 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() |