app.py

"""
DimensioDepth - Add Dimension to Everything
Advanced AI Depth Estimation with 3D Visualization

Powered by Depth-Anything V2 | Runs on Hugging Face Spaces
"""

import gradio as gr
import numpy as np
import cv2
from PIL import Image
from pathlib import Path
import sys

# Add backend to path
sys.path.append(str(Path(__file__).parent / "backend"))

# Import backend utilities
from backend.utils.image_processing import (
    depth_to_colormap,
    create_side_by_side
)

# Try to import REAL AI model
try:
    from backend.utils.transformers_depth import TransformersDepthEstimator
    print("[*] Loading REAL AI Depth-Anything V2 model...")
    depth_estimator = TransformersDepthEstimator(model_size="small")
    print("[+] REAL AI MODE ACTIVE!")
    USE_REAL_AI = True
except Exception as e:
    print(f"[!] Could not load AI models: {e}")
    print("[*] Falling back to DEMO MODE")
    from backend.utils.demo_depth import generate_smart_depth
    USE_REAL_AI = False


def estimate_depth(image, quality_mode="Fast (Preview)", colormap_style="Inferno"):
    """
    Estimate depth from an input image using REAL AI or DEMO MODE
    """
    try:
        # Convert PIL to numpy if needed
        if isinstance(image, Image.Image):
            image = np.array(image)

        # Generate depth map
        if USE_REAL_AI:
            depth = depth_estimator.predict(image)
            mode_text = "REAL AI (Depth-Anything V2)"
        else:
            depth = generate_smart_depth(image)
            mode_text = "DEMO MODE (Synthetic)"

        # Convert colormap style to cv2 constant
        colormap_dict = {
            "Inferno": cv2.COLORMAP_INFERNO,
            "Viridis": cv2.COLORMAP_VIRIDIS,
            "Plasma": cv2.COLORMAP_PLASMA,
            "Turbo": cv2.COLORMAP_TURBO,
            "Magma": cv2.COLORMAP_MAGMA,
            "Hot": cv2.COLORMAP_HOT,
            "Ocean": cv2.COLORMAP_OCEAN,
            "Rainbow": cv2.COLORMAP_RAINBOW
        }

        # Create colored depth map
        depth_colored = depth_to_colormap(depth, colormap_dict[colormap_style])

        # Create grayscale depth map
        depth_gray = (depth * 255).astype(np.uint8)
        depth_gray = cv2.cvtColor(depth_gray, cv2.COLOR_GRAY2RGB)

        # Processing info
        info = f"""
### ✅ Depth Estimation Complete!

**Mode**: {mode_text}
**Input Size**: {image.shape[1]}x{image.shape[0]}
**Output Size**: {depth.shape[1]}x{depth.shape[0]}
**Colormap**: {colormap_style}

{f"**Powered by**: Depth-Anything V2 SMALL (97MB)" if USE_REAL_AI else "**Processing**: Ultra-fast (<50ms) synthetic depth"}
"""

        return depth_colored, depth_gray, info

    except Exception as e:
        error_msg = f"### ❌ Error\n\n{str(e)}"
        print(f"Error during depth estimation: {e}")
        return None, None, error_msg


def create_side_by_side_comparison(image, quality_mode="Fast (Preview)", colormap_style="Inferno"):
    """Create side-by-side comparison of original and depth map"""
    try:
        if isinstance(image, Image.Image):
            image = np.array(image)

        # Get depth estimation
        if USE_REAL_AI:
            depth = depth_estimator.predict(image)
        else:
            depth = generate_smart_depth(image)

        # Convert colormap
        colormap_dict = {
            "Inferno": cv2.COLORMAP_INFERNO,
            "Viridis": cv2.COLORMAP_VIRIDIS,
            "Plasma": cv2.COLORMAP_PLASMA,
            "Turbo": cv2.COLORMAP_TURBO,
            "Magma": cv2.COLORMAP_MAGMA,
            "Hot": cv2.COLORMAP_HOT,
            "Ocean": cv2.COLORMAP_OCEAN,
            "Rainbow": cv2.COLORMAP_RAINBOW
        }

        # Create side-by-side
        comparison = create_side_by_side(image, depth, colormap=colormap_dict[colormap_style])

        return comparison

    except Exception as e:
        print(f"Error creating comparison: {e}")
        return None


def create_3d_visualization(image, depth_map, parallax_strength=0.5):
    """Create a simple 3D displacement visualization"""
    try:
        if image is None:
            return None

        if isinstance(image, Image.Image):
            image = np.array(image)

        if depth_map is None:
            # Generate depth if not provided
            if USE_REAL_AI:
                depth_map = depth_estimator.predict(image)
            else:
                depth_map = generate_smart_depth(image)
            depth_map = (depth_map * 255).astype(np.uint8)
        elif isinstance(depth_map, Image.Image):
            depth_map = np.array(depth_map)

        # Convert depth to grayscale if colored
        if len(depth_map.shape) == 3:
            depth_map = cv2.cvtColor(depth_map, cv2.COLOR_RGB2GRAY)

        # Normalize depth
        depth_norm = depth_map.astype(float) / 255.0

        # Create parallax effect (simple x-shift based on depth)
        h, w = image.shape[:2]
        result = image.copy()

        # Apply horizontal shift based on depth
        shift_amount = int(w * parallax_strength * 0.05)

        for y in range(h):
            for x in range(w):
                depth_val = depth_norm[y, x]
                shift = int(shift_amount * depth_val)
                new_x = min(max(x + shift, 0), w - 1)
                result[y, new_x] = image[y, x]

        return result

    except Exception as e:
        print(f"Error creating 3D viz: {e}")
        return image if image is not None else None


# Create Gradio interface
with gr.Blocks(
    theme=gr.themes.Soft(primary_hue="blue", secondary_hue="purple"),
    title="DimensioDepth - Add Dimension to Everything"
) as demo:

    gr.Markdown("""
    # 🎨 DimensioDepth - Add Dimension to Everything

    ### Transform 2D images into stunning 3D depth visualizations

    **Running in DEMO MODE** - Ultra-fast synthetic depth estimation (no AI models needed!)

    ---
    """)

    with gr.Tabs():
        # Tab 1: Main Depth Estimation
        with gr.Tab("🎯 Depth Estimation"):
            with gr.Row():
                with gr.Column(scale=1):
                    input_image = gr.Image(
                        label="Upload Your Image",
                        type="pil",
                        height=400
                    )

                    colormap_style = gr.Dropdown(
                        choices=["Inferno", "Viridis", "Plasma", "Turbo", "Magma", "Hot", "Ocean", "Rainbow"],
                        value="Inferno",
                        label="Colormap Style",
                        info="Choose your depth visualization color scheme"
                    )

                    estimate_btn = gr.Button("🚀 Generate Depth Map", variant="primary", size="lg")

                with gr.Column(scale=1):
                    depth_colored = gr.Image(label="Depth Map (Colored)", height=400)
                    depth_gray = gr.Image(label="Depth Map (Grayscale)", height=400)

            processing_info = gr.Markdown()

            estimate_btn.click(
                fn=estimate_depth,
                inputs=[input_image, gr.State("Fast"), colormap_style],
                outputs=[depth_colored, depth_gray, processing_info]
            )

        # Tab 2: Side-by-Side Comparison
        with gr.Tab("🎭 Side-by-Side Comparison"):
            gr.Markdown("""
            ### Compare Original Image with Depth Map
            Perfect for analyzing depth estimation quality and understanding 3D structure.
            """)

            with gr.Row():
                with gr.Column(scale=1):
                    compare_input = gr.Image(label="Upload Image", type="pil", height=400)

                    compare_colormap = gr.Dropdown(
                        choices=["Inferno", "Viridis", "Plasma", "Turbo", "Magma", "Hot", "Ocean", "Rainbow"],
                        value="Turbo",
                        label="Colormap"
                    )

                    compare_btn = gr.Button("🎬 Create Comparison", variant="primary")

                with gr.Column(scale=1):
                    comparison_output = gr.Image(label="Side-by-Side Comparison", height=500)

            compare_btn.click(
                fn=create_side_by_side_comparison,
                inputs=[compare_input, gr.State("Fast"), compare_colormap],
                outputs=comparison_output
            )

        # Tab 3: 3D Parallax Effect
        with gr.Tab("🌊 3D Parallax Effect"):
            gr.Markdown("""
            ### Create 3D Depth Displacement Effect
            Generate a parallax effect to visualize the 3D structure of your image.
            """)

            with gr.Row():
                with gr.Column(scale=1):
                    parallax_input = gr.Image(label="Original Image", type="pil")
                    parallax_depth = gr.Image(label="Depth Map (optional)", type="pil")
                    parallax_strength = gr.Slider(
                        minimum=0, maximum=2, value=0.5, step=0.1,
                        label="Parallax Strength",
                        info="Control the 3D displacement effect intensity"
                    )
                    parallax_btn = gr.Button("✨ Generate 3D Effect", variant="primary")

                with gr.Column(scale=1):
                    parallax_output = gr.Image(label="3D Parallax Result", height=500)

            parallax_btn.click(
                fn=create_3d_visualization,
                inputs=[parallax_input, parallax_depth, parallax_strength],
                outputs=parallax_output
            )

    # Info section
    gr.Markdown("---")
    gr.Markdown("""
    ## 💡 About This Demo

    ### 🎨 Demo Mode Features:
    - ✅ **Ultra-fast processing** (<50ms per image)
    - ✅ **No model downloads** required
    - ✅ **Advanced edge detection** + intensity analysis
    - ✅ **Surprisingly good quality** for most use cases
    - ✅ **Perfect for testing** and prototyping

    ### 🚀 How It Works:
    Demo Mode uses sophisticated computer vision techniques:
    1. **Edge Detection** - Find object boundaries
    2. **Intensity Analysis** - Analyze brightness patterns
    3. **Gaussian Smoothing** - Create smooth depth transitions
    4. **Normalization** - Convert to depth values

    ### 💡 Tips for Best Results:
    - **Image Quality**: Higher resolution = better depth detail
    - **Lighting**: Well-lit images produce clearer depth maps
    - **Contrast**: Good contrast shows better depth separation
    - **Colormap**: Inferno for general use, Viridis for scientific viz

    ---

    ### 📊 Use Cases

    - 🎨 **Creative & Artistic**: Depth-enhanced photos, 3D effects
    - 🎬 **VFX & Film**: Depth map generation for compositing
    - 🔬 **Research**: Computer vision, depth perception studies
    - 📱 **Content Creation**: Engaging 3D effects for social media

    ---

    **Tech Stack**: Advanced CV Algorithms, OpenCV, NumPy, Gradio

    Made with ❤️ for the AI community
    """)


# Launch the app
if __name__ == "__main__":
    demo.launch()