app.py

"""
AI Photo Edit - Hugging Face Spaces App
A web-based AI photo editing tool with Restyle, Enhance, and Fix capabilities
"""

import gradio as gr
import numpy as np
from PIL import Image, ImageEnhance, ImageFilter
from io import BytesIO
import base64

# Style configurations
RESTYLE_OPTIONS = {
    "Anime": "Transform into anime style with vibrant colors and expressive features",
    "Chibi Sticker": "Transform into cute chibi style with big head and small body",
    "Superhero": "Transform into superhero comic book style with dramatic lighting",
    "Graffiti": "Transform into urban graffiti street art style",
    "Watercolour": "Transform into watercolor painting with soft edges",
    "Comic": "Transform into comic book style with bold lines and halftone dots",
}

SEASON_OPTIONS = {
    "Summer": "Bright summer with warm golden sunlight and vibrant colors",
    "Winter": "Snowy winter with frost, cold blue tones",
    "Monsoon": "Rainy monsoon with wet surfaces and moody lighting",
}

ENHANCE_OPTIONS = {
    "Auto Enhance": "Improve colors, contrast, and sharpness",
    "HDR Effect": "Expand dynamic range for vivid details",
    "Sharpen": "Increase clarity and edge definition",
    "Reduce Noise": "Remove grain while preserving details",
}

FIX_OPTIONS = {
    "Blur Background": "Add portrait-style depth effect",
    "Relight": "Adjust lighting and shadows",
    "Fix Colors": "Correct white balance and color cast",
}


def apply_style(image, style_name, intensity=1.0):
    """Apply visual style filters to image"""
    if image is None:
        return None

    img = Image.fromarray(image)

    # Convert to RGB if needed
    if img.mode != 'RGB':
        img = img.convert('RGB')

    style_filters = {
        "Anime": lambda i: apply_anime(i, intensity),
        "Chibi Sticker": lambda i: apply_anime(i, intensity * 0.8),
        "Superhero": lambda i: apply_dramatic(i, intensity),
        "Graffiti": lambda i: apply_vibrant(i, intensity * 1.2),
        "Watercolour": lambda i: apply_watercolor(i, intensity),
        "Comic": lambda i: apply_comic(i, intensity),
        "Summer": lambda i: apply_summer(i, intensity),
        "Winter": lambda i: apply_winter(i, intensity),
        "Monsoon": lambda i: apply_monsoon(i, intensity),
        "Auto Enhance": lambda i: apply_enhance(i, intensity),
        "HDR Effect": lambda i: apply_hdr(i, intensity),
        "Sharpen": lambda i: apply_sharpen(i, intensity),
        "Reduce Noise": lambda i: apply_denoise(i, intensity),
        "Blur Background": lambda i: apply_blur_bg(i, intensity),
        "Relight": lambda i: apply_relight(i, intensity),
        "Fix Colors": lambda i: apply_fix_colors(i, intensity),
    }

    if style_name in style_filters:
        img = style_filters[style_name](img)

    return np.array(img)


def apply_anime(img, intensity):
    """Anime style: high saturation, contrast"""
    enhancer = ImageEnhance.Color(img)
    img = enhancer.enhance(1.0 + 0.5 * intensity)
    enhancer = ImageEnhance.Contrast(img)
    img = enhancer.enhance(1.0 + 0.25 * intensity)
    enhancer = ImageEnhance.Brightness(img)
    img = enhancer.enhance(1.0 + 0.05 * intensity)
    return img


def apply_dramatic(img, intensity):
    """Dramatic superhero style"""
    enhancer = ImageEnhance.Contrast(img)
    img = enhancer.enhance(1.0 + 0.5 * intensity)
    enhancer = ImageEnhance.Color(img)
    img = enhancer.enhance(1.0 + 0.3 * intensity)
    enhancer = ImageEnhance.Brightness(img)
    img = enhancer.enhance(1.0 - 0.05 * intensity)
    return img


def apply_vibrant(img, intensity):
    """Vibrant graffiti style"""
    enhancer = ImageEnhance.Color(img)
    img = enhancer.enhance(1.0 + 1.0 * intensity)
    enhancer = ImageEnhance.Contrast(img)
    img = enhancer.enhance(1.0 + 0.4 * intensity)
    return img


def apply_watercolor(img, intensity):
    """Watercolor painting effect"""
    img = img.filter(ImageFilter.GaussianBlur(radius=0.5 * intensity))
    enhancer = ImageEnhance.Color(img)
    img = enhancer.enhance(1.0 + 0.2 * intensity)
    return img


def apply_comic(img, intensity):
    """Comic book style"""
    enhancer = ImageEnhance.Contrast(img)
    img = enhancer.enhance(1.0 + 0.5 * intensity)
    enhancer = ImageEnhance.Color(img)
    img = enhancer.enhance(1.0 + 0.3 * intensity)
    return img


def apply_summer(img, intensity):
    """Warm summer tones"""
    enhancer = ImageEnhance.Color(img)
    img = enhancer.enhance(1.0 + 0.35 * intensity)
    enhancer = ImageEnhance.Brightness(img)
    img = enhancer.enhance(1.0 + 0.08 * intensity)
    # Add warm tint
    r, g, b = img.split()
    r = r.point(lambda x: min(255, x + int(10 * intensity)))
    img = Image.merge('RGB', (r, g, b))
    return img


def apply_winter(img, intensity):
    """Cool winter tones"""
    enhancer = ImageEnhance.Color(img)
    img = enhancer.enhance(1.0 - 0.25 * intensity)
    enhancer = ImageEnhance.Brightness(img)
    img = enhancer.enhance(1.0 + 0.12 * intensity)
    # Add cool tint
    r, g, b = img.split()
    b = b.point(lambda x: min(255, x + int(15 * intensity)))
    img = Image.merge('RGB', (r, g, b))
    return img


def apply_monsoon(img, intensity):
    """Moody monsoon effect"""
    enhancer = ImageEnhance.Color(img)
    img = enhancer.enhance(1.0 - 0.15 * intensity)
    enhancer = ImageEnhance.Brightness(img)
    img = enhancer.enhance(1.0 - 0.12 * intensity)
    enhancer = ImageEnhance.Contrast(img)
    img = enhancer.enhance(1.0 + 0.15 * intensity)
    return img


def apply_enhance(img, intensity):
    """Auto enhance"""
    enhancer = ImageEnhance.Contrast(img)
    img = enhancer.enhance(1.0 + 0.2 * intensity)
    enhancer = ImageEnhance.Color(img)
    img = enhancer.enhance(1.0 + 0.2 * intensity)
    enhancer = ImageEnhance.Sharpness(img)
    img = enhancer.enhance(1.0 + 0.1 * intensity)
    return img


def apply_hdr(img, intensity):
    """HDR effect"""
    enhancer = ImageEnhance.Contrast(img)
    img = enhancer.enhance(1.0 + 0.35 * intensity)
    enhancer = ImageEnhance.Color(img)
    img = enhancer.enhance(1.0 + 0.25 * intensity)
    return img


def apply_sharpen(img, intensity):
    """Sharpen image"""
    enhancer = ImageEnhance.Sharpness(img)
    img = enhancer.enhance(1.0 + 1.0 * intensity)
    return img


def apply_denoise(img, intensity):
    """Reduce noise (slight blur)"""
    img = img.filter(ImageFilter.GaussianBlur(radius=0.3 * intensity))
    enhancer = ImageEnhance.Sharpness(img)
    img = enhancer.enhance(1.0 + 0.2 * intensity)
    return img


def apply_blur_bg(img, intensity):
    """Blur background effect (simplified)"""
    img = img.filter(ImageFilter.GaussianBlur(radius=2 * intensity))
    return img


def apply_relight(img, intensity):
    """Adjust lighting"""
    enhancer = ImageEnhance.Brightness(img)
    img = enhancer.enhance(1.0 + 0.1 * intensity)
    enhancer = ImageEnhance.Contrast(img)
    img = enhancer.enhance(1.0 + 0.1 * intensity)
    return img


def apply_fix_colors(img, intensity):
    """Fix color cast"""
    enhancer = ImageEnhance.Color(img)
    img = enhancer.enhance(1.0 + 0.1 * intensity)
    enhancer = ImageEnhance.Contrast(img)
    img = enhancer.enhance(1.0 + 0.05 * intensity)
    return img


def apply_custom_prompt(image, prompt, intensity=1.0):
    """Apply style based on custom prompt keywords"""
    if image is None:
        return None

    prompt_lower = prompt.lower()

    # Detect style from prompt
    if any(word in prompt_lower for word in ['anime', 'cartoon', 'animated']):
        return apply_style(image, "Anime", intensity)
    elif any(word in prompt_lower for word in ['winter', 'snow', 'cold', 'frost']):
        return apply_style(image, "Winter", intensity)
    elif any(word in prompt_lower for word in ['summer', 'warm', 'sunny', 'bright']):
        return apply_style(image, "Summer", intensity)
    elif any(word in prompt_lower for word in ['vintage', 'sepia', 'old', 'retro']):
        img = Image.fromarray(image)
        if img.mode != 'RGB':
            img = img.convert('RGB')
        # Sepia effect
        enhancer = ImageEnhance.Color(img)
        img = enhancer.enhance(0.3)
        r, g, b = img.split()
        r = r.point(lambda x: min(255, x + 30))
        g = g.point(lambda x: min(255, x + 15))
        img = Image.merge('RGB', (r, g, b))
        return np.array(img)
    elif any(word in prompt_lower for word in ['black and white', 'grayscale', 'monochrome']):
        img = Image.fromarray(image)
        img = img.convert('L').convert('RGB')
        enhancer = ImageEnhance.Contrast(img)
        img = enhancer.enhance(1.1)
        return np.array(img)
    elif any(word in prompt_lower for word in ['comic', 'pop art']):
        return apply_style(image, "Comic", intensity)
    elif any(word in prompt_lower for word in ['watercolor', 'painting', 'artistic']):
        return apply_style(image, "Watercolour", intensity)
    elif any(word in prompt_lower for word in ['enhance', 'improve', 'better']):
        return apply_style(image, "Auto Enhance", intensity)
    else:
        # Default: subtle enhancement
        return apply_style(image, "Auto Enhance", intensity * 0.5)


# Create Gradio Interface
with gr.Blocks(
    title="AI Photo Edit",
    theme=gr.themes.Soft(
        primary_hue="blue",
        secondary_hue="slate",
    ),
    css="""
    .container { max-width: 1200px; margin: auto; }
    .title { text-align: center; margin-bottom: 1rem; }
    """
) as demo:

    gr.Markdown(
        """
        # AI Photo Edit
        Transform your photos with AI-powered styles, enhancements, and fixes.

        Upload an image and choose a style to apply!
        """
    )

    with gr.Row():
        with gr.Column(scale=2):
            input_image = gr.Image(label="Upload Image", type="numpy")
            output_image = gr.Image(label="Result", type="numpy")

        with gr.Column(scale=1):
            with gr.Tabs():
                with gr.TabItem("Restyle"):
                    gr.Markdown("### Artistic Styles")
                    restyle_dropdown = gr.Dropdown(
                        choices=list(RESTYLE_OPTIONS.keys()),
                        label="Choose Style",
                        value="Anime"
                    )
                    restyle_btn = gr.Button("Apply Restyle", variant="primary")

                    gr.Markdown("### Seasons")
                    season_dropdown = gr.Dropdown(
                        choices=list(SEASON_OPTIONS.keys()),
                        label="Choose Season",
                        value="Summer"
                    )
                    season_btn = gr.Button("Apply Season", variant="primary")

                    gr.Markdown("### Custom Prompt")
                    custom_prompt = gr.Textbox(
                        label="Describe your style",
                        placeholder="e.g., vintage sepia photograph from the 1920s"
                    )
                    custom_btn = gr.Button("Apply Custom", variant="primary")

                with gr.TabItem("Enhance"):
                    enhance_dropdown = gr.Dropdown(
                        choices=list(ENHANCE_OPTIONS.keys()),
                        label="Enhancement Type",
                        value="Auto Enhance"
                    )
                    enhance_btn = gr.Button("Apply Enhancement", variant="primary")

                with gr.TabItem("Fix"):
                    fix_dropdown = gr.Dropdown(
                        choices=list(FIX_OPTIONS.keys()),
                        label="Fix Type",
                        value="Relight"
                    )
                    fix_btn = gr.Button("Apply Fix", variant="primary")

            gr.Markdown("### Intensity")
            intensity_slider = gr.Slider(
                minimum=0.1,
                maximum=2.0,
                value=1.0,
                step=0.1,
                label="Effect Intensity"
            )

    # Event handlers
    restyle_btn.click(
        fn=apply_style,
        inputs=[input_image, restyle_dropdown, intensity_slider],
        outputs=output_image
    )

    season_btn.click(
        fn=apply_style,
        inputs=[input_image, season_dropdown, intensity_slider],
        outputs=output_image
    )

    custom_btn.click(
        fn=apply_custom_prompt,
        inputs=[input_image, custom_prompt, intensity_slider],
        outputs=output_image
    )

    enhance_btn.click(
        fn=apply_style,
        inputs=[input_image, enhance_dropdown, intensity_slider],
        outputs=output_image
    )

    fix_btn.click(
        fn=apply_style,
        inputs=[input_image, fix_dropdown, intensity_slider],
        outputs=output_image
    )

    gr.Markdown(
        """
        ---
        **Features:**
        - **Restyle**: Anime, Chibi, Superhero, Graffiti, Watercolour, Comic
        - **Seasons**: Summer, Winter, Monsoon
        - **Enhance**: Auto Enhance, HDR, Sharpen, Reduce Noise
        - **Fix**: Blur Background, Relight, Fix Colors

        Built with Gradio
        """
    )

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