app.py

from __future__ import annotations
import os
import math
import uuid
import json
from dataclasses import dataclass
from typing import List, Optional, Tuple

import gradio as gr
from PIL import Image, ImageOps
import pandas as pd
import numpy as np
import mediapipe as mp

# ----------------------------
# Globals & configuration
# ----------------------------
DATASET_PATH = os.getenv("HAIRSTYLE_DATASET", "data/enhanced_full_hairstyle_dataset.csv")
HAIRSTYLE_FOLDER = os.getenv("HAIRSTYLE_FOLDER", "hairstyles")
RESULTS_DIR = os.getenv("RESULTS_DIR", "generated_results")

os.makedirs(RESULTS_DIR, exist_ok=True)

# Tune these if your images tend to sit too high/low by default
DEFAULT_VERT_OFFSET_PCT = -0.25  # relative to style_forehead_height
DEFAULT_HORIZ_OFFSET_PX = 0

# MediaPipe indices used
LM_LEFT_EYE_OUTER = 33
LM_RIGHT_EYE_OUTER = 263
LM_FOREHEAD_TOP = 10
LM_FOREHEAD_LEFT = 103
LM_FOREHEAD_RIGHT = 332

# Initialize MediaPipe FaceMesh once (safer with concurrency=1 in Gradio queue)
mp_face_mesh = mp.solutions.face_mesh
FACE_MESH = mp_face_mesh.FaceMesh(
    static_image_mode=True,
    max_num_faces=1,
    refine_landmarks=True,
    min_detection_confidence=0.5
)

@dataclass
class Style:
    name: str
    gender: str
    img_path: str
    img_rgba: Optional[Image.Image]
    style_forehead_w: int
    style_forehead_h: int


def _safe_read_dataset(path: str) -> pd.DataFrame:
    if not os.path.exists(path):
        # Create an empty frame with expected columns to avoid crashes
        cols = ["name", "gender", "forehead_width_px", "forehead_height_px", "image_file"]
        return pd.DataFrame(columns=cols)
    df = pd.read_csv(path)
    # Normalize columns and fill NaNs
    for col in ["name", "gender", "image_file"]:
        if col not in df.columns:
            df[col] = ""
        df[col] = df[col].fillna("")
    for col in ["forehead_width_px", "forehead_height_px"]:
        if col not in df.columns:
            df[col] = 0
        df[col] = pd.to_numeric(df[col], errors="coerce").fillna(0).astype(int)
    return df


def _load_styles(df: pd.DataFrame) -> List[Style]:
    styles: List[Style] = []
    if not os.path.exists(HAIRSTYLE_FOLDER):
        return styles
    for _, row in df.iterrows():
        img_file = row.get("image_file", "").strip()
        if not img_file:
            continue
        path = os.path.join(HAIRSTYLE_FOLDER, img_file)
        if not os.path.exists(path):
            continue
        try:
            img = Image.open(path).convert("RGBA")
        except Exception:
            img = None
        styles.append(
            Style(
                name=str(row.get("name", "Style")).strip() or "Style",
                gender=str(row.get("gender", "All")).strip() or "All",
                img_path=path,
                img_rgba=img,
                style_forehead_w=int(row.get("forehead_width_px", 0) or 0),
                style_forehead_h=int(row.get("forehead_height_px", 0) or 0),
            )
        )
    return styles


def _to_rgb(image: Image.Image) -> Image.Image:
    return image.convert("RGB") if image.mode != "RGB" else image


def get_face_landmarks(img_rgb: Image.Image):
    """Return MediaPipe face landmarks for a PIL RGB image or None."""
    np_img = np.array(img_rgb)
    results = FACE_MESH.process(np_img)
    if results.multi_face_landmarks:
        return results.multi_face_landmarks[0]
    return None


def _rotation_angle_rad(landmarks, w: int, h: int) -> float:
    """Estimate roll angle using outer eye corners."""
    left = landmarks.landmark[LM_LEFT_EYE_OUTER]
    right = landmarks.landmark[LM_RIGHT_EYE_OUTER]
    x1, y1 = left.x * w, left.y * h
    x2, y2 = right.x * w, right.y * h
    # angle of the line from left to right; positive means head tilted CCW
    angle = math.atan2(y2 - y1, x2 - x1)
    return angle


def _compute_forehead_metrics(landmarks, w: int, h: int) -> Tuple[int, Tuple[int, int]]:
    left = landmarks.landmark[LM_FOREHEAD_LEFT]
    right = landmarks.landmark[LM_FOREHEAD_RIGHT]
    top = landmarks.landmark[LM_FOREHEAD_TOP]
    forehead_width_px = int(abs((right.x - left.x) * w))
    top_x = int(top.x * w)
    top_y = int(top.y * h)
    return forehead_width_px, (top_x, top_y)


def _paste_rgba(base: Image.Image, overlay: Image.Image, pos: Tuple[int, int]) -> Image.Image:
    canvas = base.copy().convert("RGBA")
    tmp = Image.new("RGBA", canvas.size, (0, 0, 0, 0))
    x, y = pos
    tmp.paste(overlay, (x, y), overlay)
    return Image.alpha_composite(canvas, tmp)


def apply_hairstyle_impl(
    upload_img: Optional[Image.Image],
    webcam_img: Optional[Image.Image],
    input_source: str,
    style_index: Optional[int],
    scale_tweak: float,
    vert_offset: int,
    horiz_offset: int,
    opacity: float,
) -> Tuple[Optional[Image.Image], str]:
    user_img = upload_img if input_source == "Upload" else webcam_img
    if user_img is None:
        return None, "❌ No image from selected source."
    if style_index is None or style_index < 0 or style_index >= len(STYLES):
        return _to_rgb(user_img), "ℹ️ Select a hairstyle from the gallery."

    style = STYLES[style_index]
    if style.img_rgba is None:
        return _to_rgb(user_img), f"⚠️ Could not load image for: {style.name}"

    try:
        img_rgb = _to_rgb(user_img)
        w, h = img_rgb.size
        lms = get_face_landmarks(img_rgb)
        if not lms:
            return img_rgb, "⚠️ No face detected. Showing original image. Try a clearer, front‑facing photo."

        # Compute rotation and size
        angle_rad = _rotation_angle_rad(lms, w, h)
        forehead_w_px, (top_x, top_y) = _compute_forehead_metrics(lms, w, h)

        style_fw = max(style.style_forehead_w, 1)
        style_fh = max(style.style_forehead_h, 1)

        scale_ratio = (forehead_w_px / style_fw) * float(scale_tweak)

        new_w = max(int(style.img_rgba.width * scale_ratio), 1)
        new_h = max(int(style.img_rgba.height * scale_ratio), 1)

        # Rotate hair to match head roll
        hair = style.img_rgba.resize((new_w, new_h), resample=Image.LANCZOS)
        angle_deg = math.degrees(angle_rad)
        hair = hair.rotate(angle=-angle_deg, expand=True, resample=Image.BICUBIC)

        # Compute placement
        attach_y = top_y - int(style_fh * scale_ratio)
        attach_y += int(DEFAULT_VERT_OFFSET_PCT * style_fh * scale_ratio)
        attach_y += int(vert_offset)
        attach_x = top_x - hair.width // 2 + int(horiz_offset) + int(DEFAULT_HORIZ_OFFSET_PX)

        # Clamp within canvas (x can be <0 to allow partial paste, but we clamp y >= 0)
        attach_y = max(0, attach_y)

        # Optional opacity tweak
        if 0 <= opacity < 1:
            a = hair.split()[-1]
            a = ImageOps.autocontrast(a)
            a = a.point(lambda px: int(px * opacity))
            hair = Image.merge("RGBA", (*hair.split()[:3], a))

        composed = _paste_rgba(img_rgb, hair, (attach_x, attach_y)).convert("RGB")
        return composed, "✅ Success! Tip: fine‑tune scale/offsets if needed."

    except Exception as e:
        return _to_rgb(user_img), f"❌ Error: {str(e)}"


# ----------------------------
# Load data once
# ----------------------------
DATASET_DF = _safe_read_dataset(DATASET_PATH)
STYLES: List[Style] = _load_styles(DATASET_DF)

# Precompute gallery data (image + caption)
GALLERY_ITEMS: List[Tuple[Image.Image, str]] = []
for s in STYLES:
    if s.img_rgba is not None:
        thumb = s.img_rgba.copy()
        GALLERY_ITEMS.append((thumb, s.name))


# ----------------------------
# Gradio helpers
# ----------------------------

def update_gallery(gender: str):
    if gender == "All":
        indices = list(range(len(STYLES)))
    else:
        indices = [i for i, s in enumerate(STYLES) if s.gender.lower() == gender.lower()]
    filtered = []
    for i in indices:
        s = STYLES[i]
        if s.img_rgba is not None:
            filtered.append((s.img_rgba, s.name))
    return filtered, indices


def select_hairstyle(evt: gr.SelectData, filtered_inds: List[int]):
    if filtered_inds and 0 <= evt.index < len(filtered_inds):
        return int(filtered_inds[evt.index])
    return None


def update_source(source: str):
    return gr.update(visible=source == "Upload"), gr.update(visible=source == "Webcam")


def on_apply(upload_img, webcam_img, input_source, selected_index, scale_tweak, vert_offset, horiz_offset, opacity):
    img, msg = apply_hairstyle_impl(
        upload_img, webcam_img, input_source, selected_index, scale_tweak, vert_offset, horiz_offset, opacity
    )
    return img, msg


def on_random(filtered_indices: List[int]):
    if not filtered_indices:
        return None, "ℹ️ No styles available for current filter."
    import random
    return int(random.choice(filtered_indices)), "🎲 Random style selected!"


def on_save(result_img: Optional[Image.Image]):
    if result_img is None:
        return None, "⚠️ Generate a preview first."
    file_path = os.path.join(RESULTS_DIR, f"hairstyle_{uuid.uuid4().hex}.png")
    result_img.save(file_path, format="PNG")
    return file_path, "💾 Saved! Use the button below to download."


# ----------------------------
# UI
# ----------------------------
with gr.Blocks(theme=gr.themes.Soft(), css=".small-hint{font-size:12px;opacity:.8}") as demo:
    gr.Markdown("## 💇 Virtual Hairstyle Try‑On")
    gr.Markdown(
        "Upload a front‑facing photo or use your webcam. Click a hairstyle to select it, then fine‑tune using the controls."
    )

    status = gr.Textbox(label="Status", interactive=False)
    filtered_indices = gr.State([])

    with gr.Row():
        with gr.Column(scale=1):
            input_source = gr.Radio(["Upload", "Webcam"], value="Upload", label="Input Source")
            upload_col = gr.Column(visible=True)
            with upload_col:
                upload_img = gr.Image(sources=["upload"], type="pil", label="📷 Upload Your Photo (front‑facing)")
            webcam_col = gr.Column(visible=False)
            with webcam_col:
                webcam_img = gr.Image(sources=["webcam"], type="pil", label="📹 Live Webcam", streaming=True)

            gender_filter = gr.Dropdown(choices=["All", "Male", "Female"], value="All", label="🎭 Filter by Gender")

            hairstyle_gallery = gr.Gallery(
                label="🎨 Available Hairstyles (click to select)", columns=4, height=380, object_fit="contain"
            )

            selected_index = gr.Number(value=None, visible=False)
            selected_label = gr.Markdown("*No style selected*", elem_classes=["small-hint"]) 
            random_btn = gr.Button("🎲 Random Style")

        with gr.Column(scale=2):
            result_output = gr.Image(label="🔍 Preview Result", height=520)
            with gr.Row():
                scale_tweak = gr.Slider(0.7, 1.4, value=1.0, step=0.01, label="Scale tweak")
                opacity = gr.Slider(0.6, 1.0, value=1.0, step=0.01, label="Opacity")
            with gr.Row():
                vert_offset = gr.Slider(-150, 150, value=0, step=1, label="Vertical offset (px)")
                horiz_offset = gr.Slider(-150, 150, value=0, step=1, label="Horizontal offset (px)")

            with gr.Row():
                apply_btn = gr.Button("✨ Apply Hairstyle", variant="primary")
                save_btn = gr.Button("💾 Save Preview")
                dl = gr.DownloadButton("⬇️ Download PNG", file_name="hairstyle_result.png")

    # Visibility switching
    input_source.change(update_source, inputs=input_source, outputs=[upload_col, webcam_col])

    # Gallery filtering / selection
    def _update_label(i):
        if i is None or not isinstance(i, (int, float)):
            return "*No style selected*"
        idx = int(i)
        if 0 <= idx < len(STYLES):
            return f"**Selected:** {STYLES[idx].name}"
        return "*No style selected*"

    gender_filter.change(update_gallery, inputs=gender_filter, outputs=[hairstyle_gallery, filtered_indices])
    hairstyle_gallery.select(select_hairstyle, inputs=filtered_indices, outputs=selected_index)
    selected_index.change(_update_label, inputs=selected_index, outputs=selected_label)
    random_btn.click(on_random, inputs=filtered_indices, outputs=[selected_index, status])

    # Apply + live preview
    apply_inputs = [upload_img, webcam_img, input_source, selected_index, scale_tweak, vert_offset, horiz_offset, opacity]
    apply_btn.click(on_apply, inputs=apply_inputs, outputs=[result_output, status])

    # Live webcam auto-apply (gives a smooth preview). Keep concurrency=1 for FaceMesh safety.
    webcam_img.change(on_apply, inputs=apply_inputs, outputs=[result_output, status], every=0.6)

    # Save & download
    def _save_and_link(img):
        path, msg = on_save(img)
        # Update download component with the new file
        return msg, gr.update(value=path)

    save_btn.click(_save_and_link, inputs=[result_output], outputs=[status, dl])

    # Initial gallery
    demo.load(update_gallery, inputs=gender_filter, outputs=[hairstyle_gallery, filtered_indices])

# Limit concurrency to avoid MediaPipe thread issues, enable queue for responsiveness
if __name__ == "__main__":
    demo.queue(concurrency_count=1)
    demo.launch()