app.py

import gradio as gr
import numpy as np
from PIL import Image, ImageDraw, ImageFont
import cv2
import os
import io
import base64
import time
import random

# Global variables
FEATURE_TYPES = ["Eyes", "Nose", "Lips", "Face Shape", "Hair", "Body"]
MODIFICATION_PRESETS = {
    "Eyes": ["Larger", "Smaller", "Change Color", "Change Shape"],
    "Nose": ["Refine", "Reshape", "Resize"],
    "Lips": ["Fuller", "Thinner", "Change Color"],
    "Face Shape": ["Slim", "Round", "Define Jawline", "Soften Features"],
    "Hair": ["Change Color", "Change Style", "Add Volume"],
    "Body": ["Slim", "Athletic", "Curvy", "Muscular"]
}

# Feature detection function
def detect_features(image):
    """Detect facial features in the image using OpenCV."""
    if image is None:
        return None, "Please upload an image first."
    
    # Convert to numpy array if it's a PIL Image
    if isinstance(image, Image.Image):
        img_array = np.array(image)
    else:
        img_array = image.copy()
    
    # Convert to grayscale for face detection
    gray = cv2.cvtColor(img_array, cv2.COLOR_RGB2GRAY)
    
    # Load pre-trained face detector
    face_cascade_path = cv2.data.haarcascades + 'haarcascade_frontalface_default.xml'
    eye_cascade_path = cv2.data.haarcascades + 'haarcascade_eye.xml'
    
    face_cascade = cv2.CascadeClassifier(face_cascade_path)
    eye_cascade = cv2.CascadeClassifier(eye_cascade_path)
    
    # Detect faces
    faces = face_cascade.detectMultiScale(gray, 1.3, 5)
    
    # Create a copy for visualization
    visualization = img_array.copy()
    
    # Dictionary to store detected features
    detected_features = {
        "faces": [],
        "eyes": [],
        "nose": [],
        "lips": []
    }
    
    # Draw rectangles around detected faces
    for (x, y, w, h) in faces:
        # Store face coordinates
        detected_features["faces"].append((x, y, w, h))
        
        # Draw face rectangle
        cv2.rectangle(visualization, (x, y), (x+w, y+h), (0, 255, 0), 2)
        
        # Region of interest for the face
        roi_gray = gray[y:y+h, x:x+w]
        roi_color = visualization[y:y+h, x:x+w]
        
        # Detect eyes
        eyes = eye_cascade.detectMultiScale(roi_gray)
        for (ex, ey, ew, eh) in eyes:
            # Store eye coordinates (relative to the face)
            detected_features["eyes"].append((x+ex, y+ey, ew, eh))
            
            # Draw eye rectangle
            cv2.rectangle(roi_color, (ex, ey), (ex+ew, ey+eh), (255, 0, 0), 2)
        
        # Approximate nose position (center of face)
        nose_x = x + w//2 - 15
        nose_y = y + h//2 - 10
        nose_w = 30
        nose_h = 30
        detected_features["nose"].append((nose_x, nose_y, nose_w, nose_h))
        
        # Draw nose rectangle
        cv2.rectangle(visualization, (nose_x, nose_y), (nose_x+nose_w, nose_y+nose_h), (0, 0, 255), 2)
        
        # Approximate lips position (lower third of face)
        lips_x = x + w//4
        lips_y = y + int(h * 0.7)
        lips_w = w//2
        lips_h = h//6
        detected_features["lips"].append((lips_x, lips_y, lips_w, lips_h))
        
        # Draw lips rectangle
        cv2.rectangle(visualization, (lips_x, lips_y), (lips_x+lips_w, lips_y+lips_h), (255, 0, 255), 2)
    
    # Add labels
    font = cv2.FONT_HERSHEY_SIMPLEX
    if len(detected_features["faces"]) > 0:
        cv2.putText(visualization, 'Face', (faces[0][0], faces[0][1]-10), font, 0.8, (0, 255, 0), 2)
    
    if len(detected_features["eyes"]) > 0:
        cv2.putText(visualization, 'Eye', (detected_features["eyes"][0][0], detected_features["eyes"][0][1]-5), font, 0.5, (255, 0, 0), 2)
    
    if len(detected_features["nose"]) > 0:
        cv2.putText(visualization, 'Nose', (detected_features["nose"][0][0], detected_features["nose"][0][1]-5), font, 0.5, (0, 0, 255), 2)
    
    if len(detected_features["lips"]) > 0:
        cv2.putText(visualization, 'Lips', (detected_features["lips"][0][0], detected_features["lips"][0][1]-5), font, 0.5, (255, 0, 255), 2)
    
    return Image.fromarray(visualization), detected_features

# Basic image editing function
def edit_image(image, feature_type, modification_type, intensity, detected_features):
    """Apply basic image editing based on the selected feature and modification."""
    if image is None or detected_features is None:
        return image
    
    # Convert to numpy array if it's a PIL Image
    if isinstance(image, Image.Image):
        img_array = np.array(image)
    else:
        img_array = image.copy()
    
    # Create a copy for editing
    edited_img = img_array.copy()
    
    # Apply different edits based on feature type
    if feature_type == "Eyes" and len(detected_features["eyes"]) > 0:
        for (x, y, w, h) in detected_features["eyes"]:
            # Get the eye region
            eye_region = edited_img[y:y+h, x:x+w]
            
            if modification_type == "Larger":
                # Scale the eye region
                scale_factor = 1.0 + (intensity * 0.5)  # Scale up to 1.5x based on intensity
                new_h, new_w = int(h * scale_factor), int(w * scale_factor)
                
                # Resize the eye region
                resized_eye = cv2.resize(eye_region, (new_w, new_h))
                
                # Calculate offsets to center the resized eye
                offset_y = (new_h - h) // 2
                offset_x = (new_w - w) // 2
                
                # Create a larger region to paste the resized eye
                y1 = max(0, y - offset_y)
                y2 = min(edited_img.shape[0], y + h + offset_y)
                x1 = max(0, x - offset_x)
                x2 = min(edited_img.shape[1], x + w + offset_x)
                
                # Blend the resized eye with the original image
                alpha = 0.7  # Blend factor
                try:
                    # Crop the resized eye to fit the target region
                    crop_y1 = max(0, offset_y - (y - y1))
                    crop_y2 = crop_y1 + (y2 - y1)
                    crop_x1 = max(0, offset_x - (x - x1))
                    crop_x2 = crop_x1 + (x2 - x1)
                    
                    cropped_eye = resized_eye[crop_y1:crop_y2, crop_x1:crop_x2]
                    
                    # Ensure dimensions match before blending
                    if cropped_eye.shape[0] == (y2 - y1) and cropped_eye.shape[1] == (x2 - x1):
                        edited_img[y1:y2, x1:x2] = cv2.addWeighted(
                            edited_img[y1:y2, x1:x2], 1-alpha, cropped_eye, alpha, 0
                        )
                except Exception as e:
                    print(f"Error resizing eye: {e}")
            
            elif modification_type == "Smaller":
                # Scale the eye region
                scale_factor = 1.0 - (intensity * 0.3)  # Scale down to 0.7x based on intensity
                new_h, new_w = int(h * scale_factor), int(w * scale_factor)
                
                # Resize the eye region
                resized_eye = cv2.resize(eye_region, (new_w, new_h))
                
                # Calculate offsets to center the resized eye
                offset_y = (h - new_h) // 2
                offset_x = (w - new_w) // 2
                
                # Create a background (use the surrounding area)
                background = edited_img[y:y+h, x:x+w].copy()
                
                # Paste the resized eye onto the background
                background[offset_y:offset_y+new_h, offset_x:offset_x+new_w] = resized_eye
                
                # Blend the result with the original image
                edited_img[y:y+h, x:x+w] = background
            
            elif modification_type == "Change Color":
                # Apply a color tint to the eye region
                # Generate a random color based on intensity
                blue = random.randint(0, 255)
                green = random.randint(0, 255)
                red = random.randint(0, 255)
                
                # Create a color overlay
                overlay = np.ones(eye_region.shape, dtype=np.uint8) * np.array([blue, green, red], dtype=np.uint8)
                
                # Blend the overlay with the eye region
                alpha = intensity * 0.7  # Adjust alpha based on intensity
                edited_img[y:y+h, x:x+w] = cv2.addWeighted(eye_region, 1-alpha, overlay, alpha, 0)
    
    elif feature_type == "Nose" and len(detected_features["nose"]) > 0:
        for (x, y, w, h) in detected_features["nose"]:
            # Get the nose region
            nose_region = edited_img[y:y+h, x:x+w]
            
            if modification_type == "Refine":
                # Apply a subtle blur to refine the nose
                blurred_nose = cv2.GaussianBlur(nose_region, (5, 5), 0)
                
                # Blend the blurred nose with the original
                alpha = intensity * 0.8
                edited_img[y:y+h, x:x+w] = cv2.addWeighted(nose_region, 1-alpha, blurred_nose, alpha, 0)
            
            elif modification_type == "Reshape" or modification_type == "Resize":
                # Apply a subtle transformation
                scale_x = 1.0 + (intensity * 0.4 - 0.2)  # Scale between 0.8x and 1.2x
                scale_y = 1.0 + (intensity * 0.4 - 0.2)
                
                # Create transformation matrix
                center = (w // 2, h // 2)
                M = cv2.getRotationMatrix2D(center, 0, scale_x)
                
                # Apply transformation
                transformed_nose = cv2.warpAffine(nose_region, M, (w, h))
                
                # Blend the transformed nose with the original
                alpha = 0.7
                edited_img[y:y+h, x:x+w] = cv2.addWeighted(nose_region, 1-alpha, transformed_nose, alpha, 0)
    
    elif feature_type == "Lips" and len(detected_features["lips"]) > 0:
        for (x, y, w, h) in detected_features["lips"]:
            # Get the lips region
            lips_region = edited_img[y:y+h, x:x+w]
            
            if modification_type == "Fuller":
                # Scale the lips region
                scale_factor = 1.0 + (intensity * 0.3)  # Scale up to 1.3x based on intensity
                new_h, new_w = int(h * scale_factor), int(w * scale_factor)
                
                # Resize the lips region
                resized_lips = cv2.resize(lips_region, (new_w, new_h))
                
                # Calculate offsets to center the resized lips
                offset_y = (new_h - h) // 2
                offset_x = (new_w - w) // 2
                
                # Create a larger region to paste the resized lips
                y1 = max(0, y - offset_y)
                y2 = min(edited_img.shape[0], y + h + offset_y)
                x1 = max(0, x - offset_x)
                x2 = min(edited_img.shape[1], x + w + offset_x)
                
                # Blend the resized lips with the original image
                alpha = 0.7  # Blend factor
                try:
                    # Crop the resized lips to fit the target region
                    crop_y1 = max(0, offset_y - (y - y1))
                    crop_y2 = crop_y1 + (y2 - y1)
                    crop_x1 = max(0, offset_x - (x - x1))
                    crop_x2 = crop_x1 + (x2 - x1)
                    
                    cropped_lips = resized_lips[crop_y1:crop_y2, crop_x1:crop_x2]
                    
                    # Ensure dimensions match before blending
                    if cropped_lips.shape[0] == (y2 - y1) and cropped_lips.shape[1] == (x2 - x1):
                        edited_img[y1:y2, x1:x2] = cv2.addWeighted(
                            edited_img[y1:y2, x1:x2], 1-alpha, cropped_lips, alpha, 0
                        )
                except Exception as e:
                    print(f"Error resizing lips: {e}")
            
            elif modification_type == "Thinner":
                # Scale the lips region
                scale_factor = 1.0 - (intensity * 0.3)  # Scale down to 0.7x based on intensity
                new_h, new_w = int(h * scale_factor), int(w)  # Only reduce height
                
                # Resize the lips region
                resized_lips = cv2.resize(lips_region, (new_w, new_h))
                
                # Calculate offsets to center the resized lips
                offset_y = (h - new_h) // 2
                offset_x = 0
                
                # Create a background (use the surrounding area)
                background = edited_img[y:y+h, x:x+w].copy()
                
                # Paste the resized lips onto the background
                background[offset_y:offset_y+new_h, offset_x:offset_x+new_w] = resized_lips
                
                # Blend the result with the original image
                edited_img[y:y+h, x:x+w] = background
            
            elif modification_type == "Change Color":
                # Apply a color tint to the lips
                # Use a reddish color for lips
                red_tint = np.ones(lips_region.shape, dtype=np.uint8) * np.array([50, 50, 200], dtype=np.uint8)
                
                # Blend the tint with the lips region
                alpha = intensity * 0.6  # Adjust alpha based on intensity
                edited_img[y:y+h, x:x+w] = cv2.addWeighted(lips_region, 1-alpha, red_tint, alpha, 0)
    
    elif feature_type == "Face Shape" and len(detected_features["faces"]) > 0:
        for (x, y, w, h) in detected_features["faces"]:
            # Get the face region
            face_region = edited_img[y:y+h, x:x+w]
            
            if modification_type == "Slim":
                # Apply a slimming effect by squeezing horizontally
                scale_x = 1.0 - (intensity * 0.2)  # Scale between 0.8x and 1.0x horizontally
                scale_y = 1.0  # Keep vertical scale the same
                
                # Create transformation matrix
                center = (w // 2, h // 2)
                M = cv2.getRotationMatrix2D(center, 0, 1.0)
                M[0, 0] = scale_x  # Modify the horizontal scale
                
                # Apply transformation
                transformed_face = cv2.warpAffine(face_region, M, (w, h))
                
                # Blend the transformed face with the original
                alpha = 0.7
                edited_img[y:y+h, x:x+w] = cv2.addWeighted(face_region, 1-alpha, transformed_face, alpha, 0)
            
            elif modification_type == "Round":
                # Apply a rounding effect
                # Create a circular mask
                mask = np.zeros((h, w), dtype=np.uint8)
                center = (w // 2, h // 2)
                radius = min(w, h) // 2
                cv2.circle(mask, center, radius, 255, -1)
                
                # Blur the edges of the face
                blurred_face = cv2.GaussianBlur(face_region, (21, 21), 0)
                
                # Blend based on the mask
                alpha = intensity * 0.5
                for i in range(h):
                    for j in range(w):
                        if mask[i, j] == 0:
                            # Outside the circle, blend more of the blurred face
                            edited_img[y+i, x+j] = cv2.addWeighted(
                                face_region[i, j].reshape(1, 3), 1-alpha, 
                                blurred_face[i, j].reshape(1, 3), alpha, 0
                            ).reshape(3)
    
    # For other features, apply a simpler effect
    else:
        # Add a visual indicator to show something happened
        # Draw a small colored rectangle in the corner to indicate processing
        color = (int(255 * intensity), 100, 200)
        cv2.rectangle(edited_img, (10, 10), (30, 30), color, -1)
        
        # Add text to indicate the modification
        font = cv2.FONT_HERSHEY_SIMPLEX
        cv2.putText(
            edited_img, 
            f"{feature_type}: {modification_type}", 
            (40, 25), 
            font, 
            0.7, 
            (255, 255, 255), 
            2
        )
    
    return Image.fromarray(edited_img)

# Main processing function
def process_image(image, feature_type, modification_type, intensity, custom_prompt="", use_custom_prompt=False):
    if image is None:
        return None, None, "Please upload an image first."
    
    # Step 1: Detect features and create visualization
    visualization, detected_features = detect_features(image)
    
    # Step 2: Apply edits based on detected features
    if isinstance(image, np.ndarray):
        processed_image = edit_image(image, feature_type, modification_type, intensity, detected_features)
    else:
        processed_image = edit_image(np.array(image), feature_type, modification_type, intensity, detected_features)
    
    # Get the instruction based on feature and modification
    if use_custom_prompt and custom_prompt:
        instruction = custom_prompt
    else:
        instruction = f"Applied {feature_type} modification: {modification_type} with intensity {intensity:.1f}"
    
    return processed_image, visualization, f"Edit applied: {instruction}\n\nNote: This is using CPU-based processing. For more advanced AI-powered edits, download the Pinokio local version which supports GPU acceleration."

# UI Components
def create_ui():
    with gr.Blocks(title="PortraitPerfectAI - Facial & Body Feature Editor") as app:
        gr.Markdown("# PortraitPerfectAI - Facial & Body Feature Editor")
        gr.Markdown("Upload an image and use the controls to edit specific facial and body features.")
        
        with gr.Row():
            with gr.Column(scale=1):
                # Input controls
                input_image = gr.Image(label="Upload Image", type="numpy")
                
                with gr.Group():
                    gr.Markdown("### Feature Selection")
                    feature_type = gr.Dropdown(
                        choices=FEATURE_TYPES, 
                        label="Select Feature", 
                        value="Eyes"
                    )
                    
                    # Initialize with choices for the default feature (Eyes)
                    modification_type = gr.Dropdown(
                        choices=MODIFICATION_PRESETS["Eyes"], 
                        label="Modification Type", 
                        value="Larger"
                    )
                    
                    intensity = gr.Slider(
                        minimum=0.1, 
                        maximum=1.0, 
                        value=0.5, 
                        step=0.1, 
                        label="Intensity"
                    )
                
                with gr.Group():
                    gr.Markdown("### Custom Prompt (Advanced)")
                    use_custom_prompt = gr.Checkbox(
                        label="Use Custom Prompt", 
                        value=False
                    )
                    custom_prompt = gr.Textbox(
                        label="Custom Prompt", 
                        placeholder="e.g., make the eyes blue and add long eyelashes"
                    )
                
                edit_button = gr.Button("Apply Edit", variant="primary")
                reset_button = gr.Button("Reset")
                status_text = gr.Textbox(label="Status", interactive=False)
            
            with gr.Column(scale=1):
                # Output display
                with gr.Tab("Edited Image"):
                    output_image = gr.Image(label="Edited Image", type="pil")
                
                with gr.Tab("Feature Detection"):
                    feature_visualization = gr.Image(label="Detected Features", type="pil")
                
                # Download Pinokio package section
                with gr.Accordion("Download Full Version for Local Use", open=True):
                    gr.Markdown("""
                    ### Get the Full AI-Powered Version
                    
                    For more advanced AI-powered editing with GPU acceleration:
                    
                    1. Download the Pinokio package below
                    2. Install [Pinokio](https://pinokio.computer/) on your computer
                    3. Follow the instructions in the PINOKIO_GUIDE.md file
                    
                    [Download Pinokio Package](pinokio-package.zip)
                    """)
                
                # Information about the application
                with gr.Accordion("About This Application", open=False):
                    gr.Markdown("""
                    ### PortraitPerfectAI
                    
                    This application allows you to make precise edits to facial and body features in uploaded images.
                    
                    **Features:**
                    - Edit facial features like eyes, nose, lips, and more
                    - Modify body proportions and characteristics
                    - Intuitive sliders and controls
                    - Non-destructive editing workflow
                    
                    **Note:** The web version uses CPU-based processing. For more advanced AI-powered editing with GPU acceleration, download the Pinokio package.
                    """)
        
        # Event handlers
        def update_modification_choices(feature):
            return gr.Dropdown(choices=MODIFICATION_PRESETS[feature])
        
        feature_type.change(
            fn=update_modification_choices,
            inputs=feature_type,
            outputs=modification_type
        )
        
        edit_button.click(
            fn=process_image,
            inputs=[
                input_image, 
                feature_type, 
                modification_type, 
                intensity, 
                custom_prompt, 
                use_custom_prompt
            ],
            outputs=[output_image, feature_visualization, status_text]
        )
        
        def reset_image():
            return None, None, "Image reset."
        
        reset_button.click(
            fn=reset_image,
            inputs=[],
            outputs=[output_image, feature_visualization, status_text]
        )
        
        # Add ethical usage notice
        gr.Markdown("""
        ## Ethical Usage Notice
        
        This tool is designed for creative and personal use. Please ensure:
        
        - You have appropriate rights to edit the images you upload
        - You use this tool responsibly and respect the dignity of individuals
        - You understand that AI-generated modifications are artificial and may not represent reality
        
        By using this application, you agree to these terms.
        """)
        
    return app

# Launch the app
if __name__ == "__main__":
    app = create_ui()
    app.launch(server_name="0.0.0.0", share=False)