app.py

import gradio as gr
from PIL import Image, ImageOps, ImageDraw
import numpy as np
import cv2 # Using OpenCV for fast resizing and filtering
import mediapipe as mp
import math

# --- Asset Setup ---
# Initialize MediaPipe Face Mesh for landmark detection, now supporting up to 2 faces.
mp_face_mesh = mp.solutions.face_mesh
face_mesh = mp_face_mesh.FaceMesh(static_image_mode=True, max_num_faces=2, min_detection_confidence=0.5)


# --- Standard Filter Functions ---
def apply_grayscale(img_np):
    if img_np is None: return None
    return cv2.cvtColor(cv2.cvtColor(img_np, cv2.COLOR_RGB2GRAY), cv2.COLOR_GRAY2RGB)

def apply_sepia(img_np):
    if img_np is None: return None
    sepia_matrix = np.array([[0.393, 0.769, 0.189], [0.349, 0.686, 0.168], [0.272, 0.534, 0.131]]).T
    sepia_img = np.dot(img_np[...,:3], sepia_matrix)
    return np.clip(sepia_img, 0, 255).astype(np.uint8)

def apply_invert(img_np):
    if img_np is None: return None
    return cv2.bitwise_not(img_np)

def apply_posterize(img_np):
    if img_np is None: return None
    bits = 4
    shift = 8 - bits
    return ((img_np >> shift) << shift).astype(np.uint8)

def apply_solarize(img_np):
    if img_np is None: return None
    threshold = 128
    return np.where(img_np > threshold, 255 - img_np, img_np).astype(np.uint8)

def apply_vignette(img_np):
    if img_np is None: return None
    rows, cols = img_np.shape[:2]
    kernel_x = cv2.getGaussianKernel(cols, int(cols * 0.5))
    kernel_y = cv2.getGaussianKernel(rows, int(rows * 0.5))
    kernel = kernel_y * kernel_x.T
    mask = 255 * kernel / np.max(kernel)
    return np.clip(img_np * (mask[:, :, np.newaxis] / 255.0), 0, 255).astype(np.uint8)

def apply_contour(img_np):
    if img_np is None: return None
    gray = cv2.cvtColor(img_np, cv2.COLOR_RGB2GRAY)
    edges = cv2.Canny(gray, 100, 200)
    return cv2.cvtColor(255 - edges, cv2.COLOR_GRAY2RGB)

def apply_sharpen(img_np):
    if img_np is None: return None
    kernel = np.array([[-1, -1, -1], [-1,  9, -1], [-1, -1, -1]])
    return cv2.filter2D(img_np, -1, kernel)

def apply_cartoon(img_np):
    if img_np is None: return None
    gray = cv2.cvtColor(img_np, cv2.COLOR_RGB2GRAY)
    gray = cv2.medianBlur(gray, 5)
    edges = cv2.adaptiveThreshold(gray, 255, cv2.ADAPTIVE_THRESH_MEAN_C, cv2.THRESH_BINARY, 9, 9)
    color = cv2.bilateralFilter(img_np, 9, 250, 250)
    return cv2.bitwise_and(color, color, mask=edges)

def apply_sketch(img_np):
    if img_np is None: return None
    gray_img = cv2.cvtColor(img_np, cv2.COLOR_RGB2GRAY)
    invert_img = 255 - gray_img
    blur_img = cv2.GaussianBlur(invert_img, (21, 21), 0)
    invert_blur_img = 255 - blur_img
    sketch_img = cv2.divide(gray_img, invert_blur_img, scale=256.0)
    return cv2.cvtColor(sketch_img, cv2.COLOR_GRAY2RGB)

def apply_pixelate(img_np):
    if img_np is None: return None
    h, w = img_np.shape[:2]
    pixel_size = 16
    temp = cv2.resize(img_np, (w // pixel_size, h // pixel_size), interpolation=cv2.INTER_NEAREST)
    return cv2.resize(temp, (w, h), interpolation=cv2.INTER_NEAREST)

# --- Advanced Filters ---

def apply_hdr_effect(img_np):
    """Simulates an HDR effect by enhancing details."""
    if img_np is None: return None
    return cv2.detailEnhance(img_np, sigma_s=12, sigma_r=0.15)

def apply_color_splash(img_np, color_str):
    """Keeps a selected color and converts the rest of the image to grayscale."""
    if img_np is None or color_str is None: return img_np

    try:
        if color_str.startswith('#'):
            h = color_str.lstrip('#')
            rgb_color = tuple(int(h[i:i+2], 16) for i in (0, 2, 4))
        elif color_str.startswith('rgb'):
            parts = color_str.strip('rgb()').split(',')
            rgb_color = tuple(int(p.strip()) for p in parts)
        else:
            rgb_color = (0, 0, 0)
    except (ValueError, IndexError):
        print(f"Warning: Could not parse color '{color_str}'. Defaulting to black.")
        rgb_color = (0, 0, 0)

    hsv_img = cv2.cvtColor(img_np, cv2.COLOR_RGB2HSV)
    hsv_color = cv2.cvtColor(np.uint8([[rgb_color]]), cv2.COLOR_RGB2HSV)[0][0]
    hue_tolerance = 10
    lower_bound = np.array([max(0, hsv_color[0] - hue_tolerance), 50, 50])
    upper_bound = np.array([min(179, hsv_color[0] + hue_tolerance), 255, 255])
    mask = cv2.inRange(hsv_img, lower_bound, upper_bound)
    mask_inv = cv2.bitwise_not(mask)
    gray_img = cv2.cvtColor(img_np, cv2.COLOR_RGB2GRAY)
    gray_img_3_channel = cv2.cvtColor(gray_img, cv2.COLOR_GRAY2RGB)
    colored_part = cv2.bitwise_and(img_np, img_np, mask=mask)
    grayscale_part = cv2.bitwise_and(gray_img_3_channel, gray_img_3_channel, mask=mask_inv)
    return cv2.add(colored_part, grayscale_part)

def apply_sunburst_glow(img_np):
    """Adds a sunburst/lens flare effect from the brightest point."""
    if img_np is None: return None
    gray = cv2.cvtColor(img_np, cv2.COLOR_RGB2GRAY)
    minVal, maxVal, minLoc, maxLoc = cv2.minMaxLoc(gray)
    overlay = img_np.copy()
    for i in range(12):
        angle = i * 30 * np.pi / 180
        length = np.random.randint(int(maxVal/2), int(maxVal*1.5))
        pt2_x = int(maxLoc[0] + length * np.cos(angle))
        pt2_y = int(maxLoc[1] + length * np.sin(angle))
        cv2.line(overlay, maxLoc, (pt2_x, pt2_y), (255, 255, 220), 1)
    glow = cv2.GaussianBlur(overlay, (0,0), sigmaX=30, sigmaY=30)
    return cv2.addWeighted(img_np, 0.8, glow, 0.4, 0)

def apply_dreamy_glow(img_np):
    """Adds a soft, dreamy glow effect to the image."""
    if img_np is None: return None
    blurred = cv2.GaussianBlur(img_np, (0,0), sigmaX=15, sigmaY=15)
    return cv2.addWeighted(img_np, 1.0, blurred, 0.6, 0)


# --- Sunglasses Functions ---
def _create_sunglasses_mask(style="aviator"):
    """Creates a PIL image mask for a given sunglass style."""
    sunglasses = Image.new('RGBA', (300, 150), (0, 0, 0, 0))
    draw = ImageDraw.Draw(sunglasses)

    if style == "aviator":
        # Left Lens
        draw.polygon([(20, 50), (130, 40), (120, 110), (30, 110)], fill=(20, 20, 20, 200), outline='gold', width=3)
        # Right Lens
        draw.polygon([(170, 40), (280, 50), (270, 110), (180, 110)], fill=(20, 20, 20, 200), outline='gold', width=3)
        # Bridge
        draw.line((130, 45, 170, 45), fill='gold', width=5)
        draw.line((130, 55, 170, 55), fill='gold', width=5)
    
    elif style == "retro_square":
        # Left Lens
        draw.rectangle((20, 40, 130, 110), fill=(10, 10, 10, 210), outline='white', width=6)
        # Right Lens
        draw.rectangle((170, 40, 280, 110), fill=(10, 10, 10, 210), outline='white', width=6)
        # Bridge
        draw.rectangle((130, 60, 170, 75), fill='white')

    return sunglasses

def apply_sunglasses(img_np, style="aviator"):
    """Applies a specific style of sunglasses to all detected faces (up to 2)."""
    if img_np is None: return img_np
    results = face_mesh.process(img_np)
    pil_image = Image.fromarray(img_np)
    if results.multi_face_landmarks:
        for face_landmarks in results.multi_face_landmarks:
            landmarks = np.array([(lm.x * img_np.shape[1], lm.y * img_np.shape[0]) for lm in face_landmarks.landmark])
            left_eye, right_eye = landmarks[33], landmarks[263]
            eye_center = (left_eye + right_eye) / 2
            eye_width = np.linalg.norm(left_eye - right_eye)
            angle = math.degrees(math.atan2(right_eye[1] - left_eye[1], right_eye[0] - left_eye[0]))
            
            sunglasses_img = _create_sunglasses_mask(style)
            
            w, h = int(eye_width * 1.8), int(eye_width * 1.8 * sunglasses_img.height / sunglasses_img.width)
            resized_sunglasses = sunglasses_img.resize((w, h), Image.Resampling.LANCZOS)
            rotated_sunglasses = resized_sunglasses.rotate(angle, expand=True, resample=Image.Resampling.BICUBIC)
            pos_x, pos_y = int(eye_center[0] - rotated_sunglasses.width / 2), int(eye_center[1] - rotated_sunglasses.height / 2)
            
            pil_image.paste(rotated_sunglasses, (pos_x, pos_y), rotated_sunglasses)
            
    return np.array(pil_image)


# --- Main Processing Function ---
def process_image(image, filter_name, splash_color):
    if image is None: return None
    img_np = np.array(image.convert("RGB")) if isinstance(image, Image.Image) else image
    
    filter_map = {
        "Grayscale": apply_grayscale, "Sepia": apply_sepia, "Invert": apply_invert,
        "Posterize": apply_posterize, "Solarize": apply_solarize, "Vignette": apply_vignette,
        "Contour": apply_contour, "Sharpen": apply_sharpen, "Cartoon": apply_cartoon,
        "Sketch": apply_sketch, "Pixelate": apply_pixelate,
        "HDR Effect": apply_hdr_effect, "Sunburst Glow": apply_sunburst_glow,
        "Dreamy Glow": apply_dreamy_glow,
        "Color Splash": lambda img: apply_color_splash(img, splash_color),
        "Aviator Sunglasses": lambda img: apply_sunglasses(img, style="aviator"),
        "Retro Square Sunglasses": lambda img: apply_sunglasses(img, style="retro_square"),
        "None": lambda img: img
    }
    
    filter_function = filter_map.get(filter_name, lambda img: img)
    return filter_function(img_np.copy())

# --- Gradio UI ---
css = """
#title { text-align: center; color: #1d1e22; font-size: 2.8em; font-weight: 700; }
#subtitle { text-align: center; color: #57606a; font-size: 1.2em; }
.gradio-container { max-width: 1280px !important; margin: auto !important; }
"""

with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
    gr.Markdown("# Advanced Image & Face Filter Studio", elem_id="title")
    gr.Markdown("Apply classic, artistic, and face-aware effects to your images.", elem_id="subtitle")

    filters_standard = ["Grayscale", "Sepia", "Invert", "Posterize", "Solarize", "Vignette", "Contour", "Sharpen"]
    filters_artistic = ["Cartoon", "Sketch", "Pixelate"]
    filters_advanced = ["HDR Effect", "Color Splash", "Sunburst Glow", "Dreamy Glow"]
    filters_face = ["Aviator Sunglasses", "Retro Square Sunglasses"]
    
    all_filters = ["None"] + filters_standard + filters_artistic + filters_advanced + filters_face

    with gr.Row(equal_height=False):
        with gr.Column(scale=2):
            input_image = gr.Image(sources=["upload", "webcam"], type="pil", label="Input Image")
        
        with gr.Column(scale=1):
            gr.Markdown("### Filter Controls")
            filter_radio = gr.Radio(all_filters, label="Select a Filter", value="None")
            color_picker = gr.ColorPicker(label="Color to Keep (for Color Splash)", value="#FF0000", visible=False)
            apply_button = gr.Button("Apply Filter", variant="primary")

        with gr.Column(scale=2):
            output_image = gr.Image(label="Filtered Output")

    def master_update_function(img, selected_filter, splash_color):
        """This function is the single point of truth for applying filters."""
        processed_img = process_image(img, selected_filter, splash_color)
        color_picker_visibility = gr.update(visible=True) if selected_filter == "Color Splash" else gr.update(visible=False)
        return processed_img, color_picker_visibility

    trigger_inputs = [input_image, filter_radio, color_picker]
    trigger_outputs = [output_image, color_picker]

    filter_radio.change(master_update_function, inputs=trigger_inputs, outputs=trigger_outputs)
    apply_button.click(master_update_function, inputs=trigger_inputs, outputs=trigger_outputs)
    input_image.change(master_update_function, inputs=trigger_inputs, outputs=trigger_outputs)
    color_picker.change(master_update_function, inputs=trigger_inputs, outputs=trigger_outputs)


if __name__ == "__main__":
    demo.launch(debug=True)