Spaces:
Running
Running
| import os | |
| import torch | |
| import torch.nn as nn | |
| import numpy as np | |
| import cv2 | |
| import traceback | |
| import gc | |
| from PIL import Image, ImageFilter, ImageEnhance, ImageDraw, ImageFont | |
| from torchvision.transforms import functional as TF | |
| from scipy.ndimage import label | |
| import antialiased_cnns | |
| import mediapipe as mp | |
| from skimage.exposure import match_histograms | |
| from transformers import SegformerImageProcessor, SegformerForSemanticSegmentation | |
| from ultralytics import YOLO | |
| from gfpgan import GFPGANer | |
| import urllib.request | |
| import gradio as gr | |
| # ========================= CONFIG ========================= | |
| AGING_MODEL_PATH = "face_aging_model/best_unet_model.pth" | |
| BEARD_MODEL_PATH = "models/best_hair_117_epoch_v4.pt" | |
| GFPGAN_MODEL_PATH = "GFPGANv1.4.pth" | |
| SAFE_IMG_SIZE = 1024 | |
| SOURCE_AGE = 20 | |
| TARGET_AGE = 90 | |
| WRINKLE_STRENGTH = 0.42 | |
| CONTRAST_BOOST = 1.10 | |
| SHARPNESS_BOOST = 1.20 | |
| ALPHA_HAIR = 0.95 | |
| BLUR_RADIUS = 7 | |
| EDGE_SMOOTHING = True | |
| USE_GFPGAN = True | |
| GFPGAN_UPSCALE = 2 | |
| GFPGAN_WEIGHT = 0.65 | |
| DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| print(f"🚀 Device: {DEVICE}") | |
| os.environ["HF_HOME"] = "/tmp/hf_cache" | |
| os.makedirs("/tmp/hf_cache", exist_ok=True) | |
| # Global models (Lazy Loading) | |
| age_model = None | |
| face_processor = None | |
| face_parser = None | |
| beard_model = None | |
| gfpgan_restorer = None | |
| mp_face_mesh = mp.solutions.face_mesh | |
| # ================== DOWNLOAD HELPER ================== | |
| def download_file(url, filename): | |
| if os.path.exists(filename): | |
| print(f"✅ {filename} already exists.") | |
| return True | |
| print(f"🔄 Downloading {filename}... (~350 MB)") | |
| try: | |
| urllib.request.urlretrieve(url, filename) | |
| print(f"✅ Download completed: {filename}") | |
| return True | |
| except Exception as e: | |
| print(f"❌ Download failed: {e}") | |
| return False | |
| # ================== LOAD GFPGAN ================== | |
| def load_gfpgan(): | |
| global gfpgan_restorer | |
| if gfpgan_restorer is not None: | |
| return gfpgan_restorer | |
| if not os.path.exists(GFPGAN_MODEL_PATH): | |
| model_url = 'https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth' | |
| success = download_file(model_url, GFPGAN_MODEL_PATH) | |
| if not success: | |
| print("❌ GFPGAN model download failed.") | |
| return None | |
| print("🔄 Loading GFPGAN v1.4...") | |
| try: | |
| gfpgan_restorer = GFPGANer( | |
| model_path=GFPGAN_MODEL_PATH, | |
| upscale=GFPGAN_UPSCALE, | |
| arch='clean', | |
| channel_multiplier=2, | |
| bg_upsampler=None, | |
| device=DEVICE | |
| ) | |
| print("✅ GFPGAN loaded successfully!") | |
| return gfpgan_restorer | |
| except Exception as e: | |
| print(f"❌ GFPGAN load failed: {e}") | |
| return None | |
| # ================== LOAD AGING MODEL (UNet) ================== | |
| def load_aging_model(): | |
| global age_model | |
| if age_model is not None: | |
| return age_model | |
| print("Loading UNet aging model...") | |
| class DownLayer(nn.Module): | |
| def __init__(self, in_ch, out_ch): | |
| super().__init__() | |
| self.layer = nn.Sequential( | |
| nn.MaxPool2d(2, stride=1), | |
| antialiased_cnns.BlurPool(in_ch, stride=2), | |
| nn.Conv2d(in_ch, out_ch, 3, padding=1), | |
| nn.LeakyReLU(inplace=True), | |
| nn.Conv2d(out_ch, out_ch, 3, padding=1), | |
| nn.LeakyReLU(inplace=True) | |
| ) | |
| def forward(self, x): | |
| return self.layer(x) | |
| class UpLayer(nn.Module): | |
| def __init__(self, in_ch, out_ch): | |
| super().__init__() | |
| self.blur_upsample = nn.Sequential( | |
| nn.ConvTranspose2d(in_ch, out_ch, 2, stride=2), | |
| antialiased_cnns.BlurPool(out_ch, stride=1) | |
| ) | |
| self.layer = nn.Sequential( | |
| nn.Conv2d(out_ch * 2, out_ch, 3, padding=1), | |
| nn.LeakyReLU(inplace=True), | |
| nn.Conv2d(out_ch, out_ch, 3, padding=1), | |
| nn.LeakyReLU(inplace=True) | |
| ) | |
| def forward(self, x, skip): | |
| x = self.blur_upsample(x) | |
| x = torch.cat([x, skip], dim=1) | |
| return self.layer(x) | |
| class UNet(nn.Module): | |
| def __init__(self): | |
| super().__init__() | |
| self.init_conv = nn.Sequential( | |
| nn.Conv2d(5, 64, 3, padding=1), nn.LeakyReLU(inplace=True), | |
| nn.Conv2d(64, 64, 3, padding=1), nn.LeakyReLU(inplace=True) | |
| ) | |
| self.down1 = DownLayer(64, 128) | |
| self.down2 = DownLayer(128, 256) | |
| self.down3 = DownLayer(256, 512) | |
| self.down4 = DownLayer(512, 1024) | |
| self.up1 = UpLayer(1024, 512) | |
| self.up2 = UpLayer(512, 256) | |
| self.up3 = UpLayer(256, 128) | |
| self.up4 = UpLayer(128, 64) | |
| self.final_conv = nn.Conv2d(64, 3, 1) | |
| def forward(self, x): | |
| x0 = self.init_conv(x) | |
| x1 = self.down1(x0) | |
| x2 = self.down2(x1) | |
| x3 = self.down3(x2) | |
| x4 = self.down4(x3) | |
| x = self.up1(x4, x3) | |
| x = self.up2(x, x2) | |
| x = self.up3(x, x1) | |
| x = self.up4(x, x0) | |
| return self.final_conv(x) | |
| age_model = UNet().to(DEVICE) | |
| state = torch.load(AGING_MODEL_PATH, map_location=DEVICE, weights_only=True) | |
| age_model.load_state_dict(state) | |
| age_model.eval() | |
| if DEVICE.type == "cuda" and int(torch.__version__.split('.')[0]) >= 2: | |
| print("Compiling UNet with torch.compile...") | |
| age_model = torch.compile(age_model, mode="reduce-overhead") | |
| print("✅ Aging model loaded!") | |
| return age_model | |
| # ================== LOAD FACE PARSER ================== | |
| def load_face_parser(): | |
| global face_processor, face_parser | |
| if face_parser is not None: | |
| return face_processor, face_parser | |
| print("Loading Segformer face-parsing...") | |
| face_processor = SegformerImageProcessor.from_pretrained("jonathandinu/face-parsing") | |
| face_parser = SegformerForSemanticSegmentation.from_pretrained("jonathandinu/face-parsing") | |
| face_parser.to(DEVICE) | |
| face_parser.eval() | |
| if DEVICE.type == "cuda" and int(torch.__version__.split('.')[0]) >= 2: | |
| print("Compiling Segformer with torch.compile...") | |
| face_parser = torch.compile(face_parser, mode="reduce-overhead") | |
| print("✅ Face parser loaded!") | |
| return face_processor, face_parser | |
| # ================== LOAD BEARD MODEL ================== | |
| def load_beard_model(): | |
| global beard_model | |
| if beard_model is None: | |
| print("Loading Beard Detection Model (YOLO)...") | |
| beard_model = YOLO(BEARD_MODEL_PATH) | |
| return beard_model | |
| # ================== MASK FUNCTIONS ================== | |
| def get_lips_mask(pil_image: Image.Image) -> np.ndarray: | |
| img_np = np.array(pil_image) | |
| h, w = img_np.shape[:2] | |
| lips_mask = np.zeros((h, w), dtype=np.uint8) | |
| with mp_face_mesh.FaceMesh(static_image_mode=True, max_num_faces=1, refine_landmarks=True, | |
| min_detection_confidence=0.5) as face_mesh: | |
| rgb_image = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR) | |
| results = face_mesh.process(rgb_image) | |
| if results.multi_face_landmarks: | |
| for face_landmarks in results.multi_face_landmarks: | |
| lip_landmarks = [61, 146, 91, 181, 84, 17, 314, 405, 321, 375, 291, 308, 324, 318, 402, 317, 14, 87, 178, 88, 95] | |
| points = [] | |
| for idx in lip_landmarks: | |
| landmark = face_landmarks.landmark[idx] | |
| x = int(landmark.x * w) | |
| y = int(landmark.y * h) | |
| points.append([x, y]) | |
| if points: | |
| points_np = np.array(points, np.int32) | |
| hull = cv2.convexHull(points_np) | |
| cv2.fillConvexPoly(lips_mask, hull, 255) | |
| kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (7, 7)) | |
| lips_mask = cv2.dilate(lips_mask, kernel, iterations=2) | |
| lips_mask = cv2.GaussianBlur(lips_mask.astype(np.float32), (15, 15), 4) | |
| lips_mask = np.clip(lips_mask / 255.0, 0, 1) | |
| return lips_mask | |
| return np.zeros((h, w), dtype=np.float32) | |
| def exclude_lips_from_mask(beard_mask: np.ndarray, pil_image: Image.Image) -> np.ndarray: | |
| if np.sum(beard_mask) == 0: | |
| return beard_mask | |
| lips_mask = get_lips_mask(pil_image) | |
| lips_region = (lips_mask > 0.3).astype(np.float32) | |
| kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5)) | |
| lips_region = cv2.dilate(lips_region, kernel, iterations=1) | |
| beard_mask = beard_mask * (1.0 - lips_region) | |
| beard_mask = cv2.GaussianBlur(beard_mask, (5, 5), 1) | |
| return beard_mask | |
| def get_beard_mask(pil_image: Image.Image) -> np.ndarray: | |
| temp_path = "temp_input.jpg" | |
| try: | |
| pil_image.save(temp_path) | |
| model = load_beard_model() | |
| results = model(temp_path, device=DEVICE.type, conf=0.25, iou=0.5, verbose=False, | |
| half=True if DEVICE.type == "cuda" else False) | |
| img_np = np.array(pil_image) | |
| h, w = img_np.shape[:2] | |
| beard_mask = np.zeros((h, w), dtype=np.uint8) | |
| if results[0].masks is not None: | |
| for i, cls in enumerate(results[0].boxes.cls): | |
| if int(cls) == 0: # beard class | |
| mask = results[0].masks.data[i].cpu().numpy() | |
| mask = cv2.resize(mask, (w, h)) | |
| mask = (mask > 0.4).astype(np.uint8) * 255 | |
| beard_mask = cv2.bitwise_or(beard_mask, mask) | |
| if np.sum(beard_mask) > 0: | |
| beard_mask_float = beard_mask.astype(np.float32) / 255.0 | |
| beard_mask_float = cv2.dilate(beard_mask_float, cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (7, 7)), iterations=2) | |
| beard_mask_float = cv2.morphologyEx(beard_mask_float, cv2.MORPH_OPEN, cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3)), iterations=1) | |
| beard_mask_float = cv2.morphologyEx(beard_mask_float, cv2.MORPH_CLOSE, cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5)), iterations=2) | |
| beard_mask_float = exclude_lips_from_mask(beard_mask_float, pil_image) | |
| beard_mask_float = cv2.GaussianBlur(beard_mask_float, (7, 7), 2) | |
| beard_mask_float = np.clip(beard_mask_float, 0, 1) | |
| return beard_mask_float | |
| return np.zeros((h, w), dtype=np.float32) | |
| finally: | |
| if os.path.exists(temp_path): | |
| os.remove(temp_path) | |
| def clean_mask(mask, min_area=150): | |
| mask = mask.astype(np.uint8) | |
| labeled, num = label(mask) | |
| new_mask = np.zeros_like(mask) | |
| for i in range(1, num + 1): | |
| if np.sum(labeled == i) >= min_area: | |
| new_mask[labeled == i] = 1 | |
| return new_mask | |
| def get_hair_mask_segformer(pil_image: Image.Image) -> np.ndarray: | |
| processor, parser = load_face_parser() | |
| inputs = processor(images=pil_image, return_tensors="pt").to(DEVICE) | |
| with torch.no_grad(): | |
| outputs = parser(**inputs) | |
| logits = outputs.logits | |
| upsampled = torch.nn.functional.interpolate(logits, size=pil_image.size[::-1], mode="bilinear", align_corners=False) | |
| probs = torch.softmax(upsampled, dim=1)[0] | |
| hair_prob = probs[13].cpu().numpy() | |
| hair_mask = (hair_prob > 0.12).astype(np.uint8) | |
| face_classes = list(range(1, 6)) + list(range(8, 13)) + [17, 18] | |
| parsing = upsampled.argmax(dim=1).squeeze(0).cpu().numpy() | |
| face_mask = np.isin(parsing, face_classes).astype(np.uint8) | |
| kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (7, 7)) | |
| face_mask = cv2.dilate(face_mask, kernel, iterations=1) | |
| hair_mask = hair_mask * (1 - face_mask) | |
| hair_mask = cv2.morphologyEx(hair_mask, cv2.MORPH_OPEN, cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3)), iterations=1) | |
| hair_mask = cv2.morphologyEx(hair_mask, cv2.MORPH_CLOSE, cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (11, 11)), iterations=2) | |
| hair_mask = clean_mask(hair_mask, min_area=100) | |
| hair_mask = cv2.GaussianBlur(hair_mask.astype(np.float32), (5, 5), 1.5) | |
| hair_mask = np.clip(hair_mask, 0, 1) | |
| return hair_mask | |
| def apply_hair_and_beard_color(image: Image.Image, hair_mask: np.ndarray, beard_mask: np.ndarray): | |
| combined_mask = np.maximum(hair_mask, beard_mask) | |
| if np.sum(combined_mask) == 0: | |
| return image | |
| combined_mask = cv2.GaussianBlur(combined_mask, (BLUR_RADIUS*2+1, BLUR_RADIUS*2+1), BLUR_RADIUS) | |
| combined_mask = np.clip(combined_mask, 0, 1) | |
| if EDGE_SMOOTHING: | |
| combined_mask = cv2.bilateralFilter(combined_mask.astype(np.float32), 9, 75, 75) | |
| combined_mask = np.clip(combined_mask, 0, 1) | |
| combined_mask = np.clip(combined_mask * 1.2, 0, 1) | |
| img_np = np.array(image).astype(np.float32) | |
| target_color = np.array([255, 255, 255], dtype=np.float32) | |
| gray = cv2.cvtColor(img_np.astype(np.uint8), cv2.COLOR_RGB2GRAY).astype(np.float32) / 255.0 | |
| lum_factor = 0.6 + 0.4 * gray | |
| white_layer = target_color * lum_factor[..., np.newaxis] | |
| alpha = ALPHA_HAIR | |
| result = (1 - alpha * combined_mask[..., np.newaxis]) * img_np + (alpha * combined_mask[..., np.newaxis]) * white_layer | |
| result = np.clip(result, 0, 255).astype(np.uint8) | |
| result_pil = Image.fromarray(result) | |
| result_pil = result_pil.filter(ImageFilter.UnsharpMask(1.2, 140, 2)) | |
| return result_pil | |
| def post_correct_aged(original: Image.Image, aged: Image.Image) -> Image.Image: | |
| orig_np = np.array(original) | |
| aged_np = np.array(aged) | |
| matched = match_histograms(aged_np, orig_np, channel_axis=-1) | |
| matched_img = Image.fromarray(np.clip(matched, 0, 255).astype(np.uint8)) | |
| matched_img = ImageEnhance.Brightness(matched_img).enhance(1.10) | |
| matched_img = ImageEnhance.Contrast(matched_img).enhance(1.06) | |
| return matched_img | |
| def enhance_texture(img: Image.Image) -> Image.Image: | |
| img = img.filter(ImageFilter.UnsharpMask(2, 160, 3)) | |
| img = ImageEnhance.Contrast(img).enhance(CONTRAST_BOOST) | |
| img = ImageEnhance.Sharpness(img).enhance(SHARPNESS_BOOST) | |
| return img | |
| def create_comparison(orig, raw_aged, final): | |
| W = 640 | |
| def rsz(img): | |
| ratio = img.height / img.width if img.width else 1 | |
| return img.resize((W, int(W * ratio)), Image.LANCZOS) | |
| o, r, f = rsz(orig), rsz(raw_aged), rsz(final) | |
| H = max(o.height, r.height, f.height) | |
| canvas = Image.new("RGB", (W*3, H), (255, 255, 255)) | |
| canvas.paste(o, (0, (H - o.height)//2)) | |
| canvas.paste(r, (W, (H - r.height)//2)) | |
| canvas.paste(f, (W*2, (H - f.height)//2)) | |
| draw = ImageDraw.Draw(canvas) | |
| try: | |
| font = ImageFont.truetype("arial.ttf", 28) | |
| except: | |
| font = ImageFont.load_default() | |
| draw.text((W//4, 8), "Original", (0, 0, 0), font=font) | |
| draw.text((W + W//5, 8), "Aged Raw", (0, 0, 0), font=font) | |
| draw.text((W*2 + W//6, 8), "Final Result", (0, 0, 0), font=font) | |
| return canvas | |
| # ================== MAIN PROCESSING FUNCTION ================== | |
| def process_face_aging(input_image: Image.Image): | |
| if input_image is None: | |
| raise gr.Error("Please upload a clear photo of a young person!") | |
| try: | |
| print(f"→ Processing image: {input_image.size}") | |
| orig = input_image.convert("RGB") | |
| ow, oh = orig.size | |
| img_resized = orig.resize((SAFE_IMG_SIZE, SAFE_IMG_SIZE), Image.LANCZOS) | |
| rgb_tensor = TF.to_tensor(img_resized) | |
| src_age = torch.full((1, SAFE_IMG_SIZE, SAFE_IMG_SIZE), SOURCE_AGE / 100.0) | |
| tgt_age = torch.full((1, SAFE_IMG_SIZE, SAFE_IMG_SIZE), TARGET_AGE / 100.0) | |
| cond_input = torch.cat([rgb_tensor, src_age, tgt_age], dim=0).unsqueeze(0).to(DEVICE) | |
| with torch.no_grad(): | |
| aging_net = load_aging_model() | |
| raw_output = aging_net(cond_input).squeeze(0) | |
| raw_aged = TF.to_pil_image(raw_output.clamp(0, 1)).resize((ow, oh), Image.LANCZOS) | |
| alpha = WRINKLE_STRENGTH | |
| blended = (1 - alpha) * rgb_tensor.unsqueeze(0) + alpha * raw_output | |
| blended = blended.clamp(0, 1).squeeze(0) | |
| final_aged = TF.to_pil_image(blended).resize((ow, oh), Image.LANCZOS) | |
| final_aged = enhance_texture(final_aged) | |
| final_aged = post_correct_aged(orig, final_aged) | |
| print(" Generating hair mask...") | |
| hair_mask = get_hair_mask_segformer(final_aged) | |
| print(" Generating beard mask...") | |
| beard_mask = get_beard_mask(final_aged) | |
| print(" Applying white hair & beard...") | |
| final_img = apply_hair_and_beard_color(final_aged, hair_mask, beard_mask) | |
| if USE_GFPGAN: | |
| print(" Applying GFPGAN face restoration...") | |
| gfpgan = load_gfpgan() | |
| if gfpgan is not None: | |
| try: | |
| img_cv = cv2.cvtColor(np.array(final_img), cv2.COLOR_RGB2BGR) | |
| _, _, restored_cv = gfpgan.enhance( | |
| img_cv, has_aligned=False, only_center_face=False, paste_back=True, weight=GFPGAN_WEIGHT | |
| ) | |
| final_img = Image.fromarray(cv2.cvtColor(restored_cv, cv2.COLOR_BGR2RGB)) | |
| except Exception as e: | |
| print(f" GFPGAN error: {e}") | |
| comparison = create_comparison(orig, raw_aged, final_img) | |
| print("✓ Processing completed!") | |
| gc.collect() | |
| return final_img, comparison | |
| except Exception as e: | |
| print(f"❌ Error: {str(e)}") | |
| traceback.print_exc() | |
| raise gr.Error(f"Processing failed: {str(e)}") | |
| # ================== GRADIO INTERFACE ================== | |
| with gr.Blocks(theme=gr.themes.Soft(), title="👴 Face Aging + White Hair & Beard Generator") as demo: | |
| gr.Markdown("# 👴 Face Aging + White Hair & Beard Generator") | |
| gr.Markdown("Upload a clear photo of a young person.<br>This tool will age them to ~90 years with realistic wrinkles and add natural white hair & beard.") | |
| with gr.Row(): | |
| input_img = gr.Image(type="pil", label="Upload Young Face Photo", height=450) | |
| with gr.Row(): | |
| output_img = gr.Image(type="pil", label="Final Aged Result (with White Hair & Beard)", height=450) | |
| comparison_img = gr.Image(type="pil", label="Comparison: Original | Raw Aged | Final", height=450) | |
| btn = gr.Button("🚀 Generate Aged Face", variant="primary") | |
| btn.click( | |
| fn=process_face_aging, | |
| inputs=input_img, | |
| outputs=[output_img, comparison_img], | |
| queue=True, | |
| concurrency_limit=2 # ← Yeh line concurrency_count ki jagah use hui | |
| ) | |
| if __name__ == "__main__": | |
| print("Starting Face Aging App...") | |
| demo.queue(max_size=8).launch( | |
| server_name="0.0.0.0", | |
| server_port=7860, | |
| share=False, # HF Spaces mein false rakho | |
| debug=False # Production mein False better | |
| ) |