Create app.py

app.py

@@ -0,0 +1,230 @@
+import os
+import subprocess
+
+os.system("pip install gradio==3.50")
+os.system("pip install dlib==19.24.2")
+
+#############################################
+
+import torch
+print(f"Is CUDA available: {torch.cuda.is_available()}")
+# True
+print(f"CUDA device: {torch.cuda.get_device_name(torch.cuda.current_device())}")
+
+###################################################
+
+
+from argparse import Namespace
+import pprint
+import numpy as np
+from PIL import Image
+import torch
+import torchvision.transforms as transforms
+import cv2
+import dlib
+import matplotlib.pyplot as plt
+import gradio as gr  # Importing Gradio as gr
+from tensorflow.keras.preprocessing.image import img_to_array
+from huggingface_hub import hf_hub_download, login
+from datasets.augmentations import AgeTransformer
+from utils.common import tensor2im
+from models.psp import pSp
+
+# Huggingface login
+login(token=os.getenv("TOKENKEY"))
+
+# Download models from Huggingface
+#age_prototxt = hf_hub_download(repo_id="AshanGimhana/Age_Detection_caffe", filename="age.prototxt")
+#caffe_model = hf_hub_download(repo_id="AshanGimhana/Age_Detection_caffe", filename="dex_imdb_wiki.caffemodel")
+sam_ffhq_aging = hf_hub_download(repo_id="AshanGimhana/Face_Agin_model", filename="sam_ffhq_aging.pt")
+
+# If 'mse' is a custom function needed, 
+#custom_objects = {'mse': MeanSquaredError()}
+new_age_model = load_model("age_prediction_model.h5")
+
+# Age prediction model setup
+age_net = cv2.dnn.readNetFromCaffe(age_prototxt, caffe_model)
+
+# Face detection and landmarks predictor setup
+detector = dlib.get_frontal_face_detector()
+predictor = dlib.shape_predictor("shape_predictor_68_face_landmarks.dat")
+
+# Load the pretrained aging model
+EXPERIMENT_TYPE = 'ffhq_aging'
+EXPERIMENT_DATA_ARGS = {
+    "ffhq_aging": {
+        "model_path": sam_ffhq_aging,
+        "transform": transforms.Compose([
+            transforms.Resize((256, 256)),
+            transforms.ToTensor(),
+            transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
+        ])
+    }
+}
+EXPERIMENT_ARGS = EXPERIMENT_DATA_ARGS[EXPERIMENT_TYPE]
+model_path = EXPERIMENT_ARGS['model_path']
+ckpt = torch.load(model_path, map_location='cpu')
+opts = ckpt['opts']
+pprint.pprint(opts)
+opts['checkpoint_path'] = model_path
+opts = Namespace(**opts)
+net = pSp(opts)
+net.eval()
+net.cuda()
+
+print('Model successfully loaded!')
+
+def check_image_quality(image):
+    # Convert the image to grayscale
+    gray_image = np.array(image.convert("L"))
+    
+    # Check for under/over-exposure using histogram
+    hist = exposure.histogram(gray_image)
+    low_exposure = hist[0][:5].sum() > 0.5 * hist[0].sum()  # Significant pixels in dark range
+    high_exposure = hist[0][-5:].sum() > 0.5 * hist[0].sum()  # Significant pixels in bright range
+    
+    # Check sharpness using Laplacian variance
+    sharpness = cv2.Laplacian(np.array(image), cv2.CV_64F).var()
+    low_sharpness = sharpness < 70  # Threshold for sharpness
+    
+    # Check overall quality
+    if low_exposure or high_exposure or low_sharpness:
+        return False  # Image quality is insufficient
+    return True  # Image quality is sufficient
+
+# Functions for face and mouth region
+def get_face_region(image):
+    gray = cv2.cvtColor(np.array(image), cv2.COLOR_BGR2GRAY)
+    faces = detector(gray)
+    if len(faces) > 0:
+        return faces[0]
+    return None
+
+def get_mouth_region(image):
+    gray = cv2.cvtColor(np.array(image), cv2.COLOR_BGR2GRAY)
+    faces = detector(gray)
+    for face in faces:
+        landmarks = predictor(gray, face)
+        mouth_points = [(landmarks.part(i).x, landmarks.part(i).y) for i in range(48, 68)]
+        return np.array(mouth_points, np.int32)
+    return None
+
+# Function to predict age
+def get_age(distr):
+    # Convert distribution to approximate age by scaling
+    age = distr * 4
+    return age
+
+def predict_age(image):
+    image = np.array(image.resize((64, 64)))
+    image = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY) 
+    image = image / 255.0  
+    image = np.expand_dims(image, axis=0)
+
+    # Predict age
+    val = new_age_model.predict(np.array(image))
+    age = val[0][0]
+    return int(age)
+
+# Function for color correction
+def color_correct(source, target):
+    mean_src = np.mean(source, axis=(0, 1))
+    std_src = np.std(source, axis=(0, 1))
+    mean_tgt = np.mean(target, axis=(0, 1))
+    std_tgt = np.std(target, axis=(0, 1))
+    src_normalized = (source - mean_src) / std_src
+    src_corrected = (src_normalized * std_tgt) + mean_tgt
+    return np.clip(src_corrected, 0, 255).astype(np.uint8)
+
+# Function to replace teeth
+def replace_teeth(temp_image, aged_image):
+    temp_image = np.array(temp_image)
+    aged_image = np.array(aged_image)
+    temp_mouth = get_mouth_region(temp_image)
+    aged_mouth = get_mouth_region(aged_image)
+    if temp_mouth is None or aged_mouth is None:
+        return aged_image
+    temp_mask = np.zeros_like(temp_image)
+    cv2.fillConvexPoly(temp_mask, temp_mouth, (255, 255, 255))
+    temp_mouth_region = cv2.bitwise_and(temp_image, temp_mask)
+    temp_mouth_bbox = cv2.boundingRect(temp_mouth)
+    aged_mouth_bbox = cv2.boundingRect(aged_mouth)
+    temp_mouth_crop = temp_mouth_region[temp_mouth_bbox[1]:temp_mouth_bbox[1] + temp_mouth_bbox[3], temp_mouth_bbox[0]:temp_mouth_bbox[0] + temp_mouth_bbox[2]]
+    temp_mask_crop = temp_mask[temp_mouth_bbox[1]:temp_mouth_bbox[1] + temp_mouth_bbox[3], temp_mouth_bbox[0]:temp_mouth_bbox[0] + temp_mouth_bbox[2]]
+    temp_mouth_crop_resized = cv2.resize(temp_mouth_crop, (aged_mouth_bbox[2], aged_mouth_bbox[3]))
+    temp_mask_crop_resized = cv2.resize(temp_mask_crop, (aged_mouth_bbox[2], aged_mouth_bbox[3]))
+    aged_mouth_crop = aged_image[aged_mouth_bbox[1]:aged_mouth_bbox[1] + aged_mouth_bbox[3], aged_mouth_bbox[0]:aged_mouth_bbox[0] + aged_mouth_bbox[2]]
+    temp_mouth_crop_resized = color_correct(temp_mouth_crop_resized, aged_mouth_crop)
+    center = (aged_mouth_bbox[0] + aged_mouth_bbox[2] // 2, aged_mouth_bbox[1] + aged_mouth_bbox[3] // 2)
+    seamless_teeth = cv2.seamlessClone(temp_mouth_crop_resized, aged_image, temp_mask_crop_resized, center, cv2.NORMAL_CLONE)
+    return seamless_teeth
+
+# Function to run alignment
+def run_alignment(image):
+    from scripts.align_all_parallel import align_face
+    temp_image_path = "/tmp/temp_image.jpg"
+    image.save(temp_image_path)
+    aligned_image = align_face(filepath=temp_image_path, predictor=predictor)
+    return aligned_image
+
+# Function to apply aging
+def apply_aging(image, target_age):
+    img_transforms = EXPERIMENT_DATA_ARGS[EXPERIMENT_TYPE]['transform']
+    input_image = img_transforms(image)
+    age_transformers = [AgeTransformer(target_age=target_age)]
+    results = []
+    for age_transformer in age_transformers:
+        with torch.no_grad():
+            input_image_age = [age_transformer(input_image.cpu()).to('cuda')]
+            input_image_age = torch.stack(input_image_age)
+            result_tensor = net(input_image_age.float(), randomize_noise=False, resize=False)[0]
+            result_image = tensor2im(result_tensor)
+            results.append(np.array(result_image))
+    final_result = results[0]
+    return final_result
+
+# Function to process the image
+def process_image(uploaded_image):
+    # Loading images for good and bad teeth
+    temp_images_good = [Image.open(f"good_teeth/G{i}.JPG") for i in range(1, 4)]
+    temp_images_bad = [Image.open(f"bad_teeth/B{i}.jpeg") for i in range(1, 5)]
+
+    # Predicting the age
+    predicted_age = predict_age(uploaded_image)
+    target_age = predicted_age + 5
+
+    # Aligning the face in the uploaded image
+    aligned_image = run_alignment(uploaded_image)
+
+    # Applying aging effect
+    aged_image = apply_aging(aligned_image, target_age=target_age)
+
+    # Randomly selecting teeth images
+    good_teeth_image = temp_images_good[np.random.randint(0, len(temp_images_good))]
+    bad_teeth_image = temp_images_bad[np.random.randint(0, len(temp_images_bad))]
+
+    # Replacing teeth in aged image
+    aged_image_good_teeth = replace_teeth(good_teeth_image, aged_image)
+    aged_image_bad_teeth = replace_teeth(bad_teeth_image, aged_image)
+
+    return aged_image_good_teeth, aged_image_bad_teeth, predicted_age, target_age
+
+# Gradio Interface
+def show_results(uploaded_image):
+    # Perform quality check
+    if not check_image_quality(uploaded_image):
+        return None, None, "Not_Allowed"
+
+    # If quality is acceptable, continue with processing
+    aged_image_good_teeth, aged_image_bad_teeth, predicted_age, target_age = process_image(uploaded_image)
+    return aged_image_good_teeth, aged_image_bad_teeth, f"Predicted Age: {predicted_age}, Target Age: {target_age}"
+
+iface = gr.Interface(
+    fn=show_results,
+    inputs=gr.Image(type="pil"),
+    outputs=[gr.Image(type="pil"), gr.Image(type="pil"), gr.Textbox()],
+    title="Aging Effect with Teeth Replacement",
+    description="Upload an image to apply an aging effect. The application will generate two results: one with good teeth and one with bad teeth."
+)
+
+iface.launch()