app.py

import gradio as gr
import numpy as np
from PIL import Image, ImageFilter
import time

# ---------- small image utils ----------

def pil_to_np(img_pil):
    arr = np.asarray(img_pil.convert("RGB")).astype(np.float32) / 255.0
    return arr

def np_to_pil(arr):
    arr = np.clip(arr * 255.0, 0, 255).astype(np.uint8)
    return Image.fromarray(arr)

def resize_max_side(img_pil, max_side=1600):
    w, h = img_pil.size
    scale = min(1.0, max_side / max(w, h))
    if scale < 1.0:
        return img_pil.resize((int(w*scale), int(h*scale)), Image.LANCZOS)
    return img_pil

def rgb_to_hsv_np(rgb):
    # rgb: HxWx3 in [0,1]
    r, g, b = rgb[..., 0], rgb[..., 1], rgb[..., 2]
    mx = np.max(rgb, axis=-1)
    mn = np.min(rgb, axis=-1)
    diff = mx - mn + 1e-8

    # Hue
    h = np.zeros_like(mx)
    mask = diff > 1e-8
    r_is_max = (mx == r) & mask
    g_is_max = (mx == g) & mask
    b_is_max = (mx == b) & mask
    h[r_is_max] = (g[r_is_max] - b[r_is_max]) / diff[r_is_max]
    h[g_is_max] = 2.0 + (b[g_is_max] - r[g_is_max]) / diff[g_is_max]
    h[b_is_max] = 4.0 + (r[b_is_max] - g[b_is_max]) / diff[b_is_max]
    h = (h / 6.0) % 1.0

    # Saturation
    s = np.where(mx <= 1e-8, 0, diff / (mx + 1e-8))
    v = mx
    return np.stack([h, s, v], axis=-1)

def hsv_to_rgb_np(hsv):
    h, s, v = hsv[..., 0], hsv[..., 1], hsv[..., 2]
    i = np.floor(h * 6).astype(int)
    f = h * 6 - i
    p = v * (1 - s)
    q = v * (1 - f * s)
    t = v * (1 - (1 - f) * s)
    i_mod = i % 6

    r = np.select(
        [i_mod == 0, i_mod == 1, i_mod == 2, i_mod == 3, i_mod == 4, i_mod == 5],
        [v, q, p, p, t, v])
    g = np.select(
        [i_mod == 0, i_mod == 1, i_mod == 2, i_mod == 3, i_mod == 4, i_mod == 5],
        [t, v, v, q, p, p])
    b = np.select(
        [i_mod == 0, i_mod == 1, i_mod == 2, i_mod == 3, i_mod == 4, i_mod == 5],
        [p, p, t, v, v, q])
    rgb = np.stack([r, g, b], axis=-1)
    return np.clip(rgb, 0, 1)

def unsharp_mask(img_pil, radius=1.2, amount=0.7):
    # classic local contrast boost
    blurred = img_pil.filter(ImageFilter.GaussianBlur(radius=radius))
    arr = pil_to_np(img_pil)
    arr_blur = pil_to_np(blurred)
    out = np.clip(arr + amount * (arr - arr_blur), 0, 1)
    return np_to_pil(out)

# ---------- core adjustments ----------

def apply_adjustments(img_pil,
                      exposure_stops=0.0,
                      contrast=0.0,
                      saturation=0.0,
                      warmth=0.0,        # + warm, - cool
                      hue_shift_deg=0.0,
                      gamma=1.0,
                      clarity=0.0,
                      lift=0.0):         # lift blacks / fade
    """All params are gentle, designed to stay natural."""
    img = resize_max_side(img_pil)
    arr = pil_to_np(img)

    # exposure (stops)
    if abs(exposure_stops) > 1e-6:
        arr = np.clip(arr * (2.0 ** exposure_stops), 0, 1)

    # contrast (simple S-curve around mid 0.5)
    if abs(contrast) > 1e-6:
        arr = np.clip((arr - 0.5) * (1.0 + contrast) + 0.5, 0, 1)

    # lift blacks (fade)
    if abs(lift) > 1e-6:
        arr = np.clip(arr + lift * (1.0 - arr), 0, 1)

    # warmth (white-balance tilt)
    if abs(warmth) > 1e-6:
        wb = np.array([1.0 + warmth, 1.0, 1.0 - warmth], dtype=np.float32)
        arr = np.clip(arr * wb, 0, 1)

    # HSV tweaks (saturation + hue shift + gamma)
    hsv = rgb_to_hsv_np(arr)
    if abs(saturation) > 1e-6:
        hsv[..., 1] = np.clip(hsv[..., 1] * (1.0 + saturation), 0, 1)
    if abs(hue_shift_deg) > 1e-6:
        hsv[..., 0] = (hsv[..., 0] + hue_shift_deg / 360.0) % 1.0
    if abs(gamma - 1.0) > 1e-6:
        hsv[..., 2] = np.clip(hsv[..., 2] ** (1.0 / gamma), 0, 1)
    arr = hsv_to_rgb_np(hsv)

    out = np_to_pil(arr)
    # clarity via unsharp mask
    if abs(clarity) > 1e-6:
        out = unsharp_mask(out, radius=1.2, amount=clarity)
    return out

# ---------- aesthetic scoring (fast heuristic) ----------

def aesthetic_score_fast(img_pil):
    arr = pil_to_np(img_pil)
    # luminance
    Y = 0.2126 * arr[..., 0] + 0.7152 * arr[..., 1] + 0.0722 * arr[..., 2]
    brightness = float(np.mean(Y))
    contrast = float(np.std(Y))
    # saturation
    s = rgb_to_hsv_np(arr)[..., 1]
    sat = float(np.mean(s))

    # targets tuned for mass-appeal feed aesthetics (roughly)
    target_b = 0.62
    target_sat = 0.35

    score_b = 1.0 - min(abs(brightness - target_b) / 0.62, 1.0)
    score_c = min(max((contrast - 0.04) / 0.26, 0.0), 1.0)
    score_s = 1.0 - min(abs(sat - target_sat) / 0.35, 1.0)

    # clipping penalties
    clip_hi = float((Y > 0.98).mean())
    clip_lo = float((Y < 0.02).mean())
    penalty_clip = min(clip_hi * 4.0 + clip_lo * 2.5, 1.5)

    # white balance cast penalty (channel means too far apart)
    means = arr.reshape(-1, 3).mean(axis=0)
    cast = float(np.max(means) - np.min(means))
    penalty_cast = min(cast * 2.0, 1.0)

    # simple skin guard: if skin-ish pixels oversaturated, penalize
    hsv = rgb_to_hsv_np(arr)
    h, s_, v = hsv[..., 0], hsv[..., 1], hsv[..., 2]
    skin_mask = (h < (50/360)) | (h > (345/360))
    skin_mask &= (s_ > 0.23) & (v > 0.35)
    skin_sat = float(s_[skin_mask].mean()) if np.any(skin_mask) else 0.0
    penalty_skin = max(0.0, (skin_sat - 0.65) * 2.0)

    raw = 0.4 * score_b + 0.35 * score_c + 0.25 * score_s
    penalties = penalty_clip + penalty_cast + penalty_skin
    final = max(0.0, min(1.0, raw - 0.4 * penalties))
    return final, {
        "brightness": round(brightness, 3),
        "contrast": round(contrast, 3),
        "saturation": round(sat, 3),
        "clip_hi%": round(clip_hi * 100, 2),
        "clip_lo%": round(clip_lo * 100, 2)
    }

# ---------- vibe presets ----------

VIBES = {
    "Natural": dict(exposure_stops=0.10, contrast=0.08, saturation=0.06, warmth=0.02, clarity=0.06, gamma=1.0, lift=0.00),
    "Film":    dict(exposure_stops=0.05, contrast=-0.03, saturation=-0.02, warmth=0.05, clarity=0.03, gamma=0.95, lift=0.06),
    "Pop":     dict(exposure_stops=0.00, contrast=0.15, saturation=0.12, warmth=0.00, clarity=0.15, gamma=1.0,  lift=0.00),
    "Moody":   dict(exposure_stops=-0.15,contrast=0.10, saturation=-0.08,warmth=-0.03, clarity=0.05, gamma=1.05, lift=0.02),
    "Pastel":  dict(exposure_stops=0.10, contrast=-0.10, saturation=-0.15,warmth=0.03, clarity=0.02, gamma=0.90, lift=0.08),
}

# Keep recent result in memory so feedback buttons can store something meaningful
LAST_RESULT = {"winner": None, "scores": None}

def process(image, intensity):
    if image is None:
        raise gr.Error("Please upload a photo first.")

    # generate candidates
    candidates = []
    scores = []
    metrics = []
    for name, params in VIBES.items():
        out = apply_adjustments(image, **params)
        score, met = aesthetic_score_fast(out)
        candidates.append((name, out, score))
        scores.append(score)
        metrics.append((name, met))

    # pick winner
    candidates.sort(key=lambda x: x[2], reverse=True)
    winner_name, winner_img, winner_score = candidates[0]

    # blend intensity with original (0..100)
    t = float(intensity) / 100.0
    base = resize_max_side(image)
    wnp = pil_to_np(winner_img)
    onp = pil_to_np(base)
    blended = np_to_pil(onp * (1 - t) + wnp * t)

    # gallery: show all looks with their scores
    gallery = []
    for name, img, score in candidates:
        caption = f"{name} — score {score:.2f}"
        gallery.append((img, caption))

    # remember
    LAST_RESULT["winner"] = {
        "name": winner_name,
        "score": float(winner_score),
        "when": time.strftime("%Y-%m-%d %H:%M:%S")
    }
    LAST_RESULT["scores"] = {name: float(s) for name, _, s in candidates}

    # metrics text
    metrics_top = next(m for n, m in metrics if n == winner_name)
    info = f"Picked: **{winner_name}** (score {winner_score:.2f})"
    info += f"\n\nBrightness: {metrics_top['brightness']} | Contrast: {metrics_top['contrast']} | Saturation: {metrics_top['saturation']}"
    info += f"\nClipped Highlights: {metrics_top['clip_hi%']}% | Deep Shadows: {metrics_top['clip_lo%']}%"

    return blended, gallery, info

def feedback(good):
    if LAST_RESULT["winner"] is None:
        return "Upload a photo and generate a result first."
    # Append a tiny log in Space storage (ephemeral on free tier, good enough for MVP)
    try:
        with open("feedback_log.csv", "a", encoding="utf-8") as f:
            f.write(
                f"{LAST_RESULT['winner']['when']},{LAST_RESULT['winner']['name']},{LAST_RESULT['winner']['score']},{'up' if good else 'down'}\n"
            )
    except Exception:
        pass
    return "Thanks for the feedback! ✨"

# ---------- UI ----------

with gr.Blocks(title="One-Click Aesthetic") as demo:
    gr.Markdown(
        """
        # One-Click Aesthetic ✨
        Upload a photo and hit **Make it Aesthetic**.  
        The app tries a few tasteful looks and picks the one with the best predicted mass-appeal score.  
        Use the **Intensity** slider to control how strong the look is.
        """
    )
    with gr.Row():
        inp = gr.Image(label="Upload photo", type="pil")
        out = gr.Image(label="Aesthetic result")

    intensity = gr.Slider(0, 100, value=80, step=1, label="Intensity (blend)")

    go = gr.Button("Make it Aesthetic", variant="primary")
    info = gr.Markdown()
    gallery = gr.Gallery(label="Tried looks (ranked high → low)", show_label=True, columns=5, height="auto")

    with gr.Row():
        up = gr.Button("👍 Looks great")
        down = gr.Button("👎 Needs work")

    go.click(process, inputs=[inp, intensity], outputs=[out, gallery, info])
    up.click(lambda: feedback(True), inputs=None, outputs=info)
    down.click(lambda: feedback(False), inputs=None, outputs=info)

demo.launch()