app.py

import gradio as gr
import cv2, numpy as np
from PIL import Image

# -------------------------------------------------
# 🔧 CORE FUNCTION
# -------------------------------------------------
def remove_folds(image, intensity=0.5):
    if image is None:
        return None

    img = np.array(image.convert("RGB"))
    img_lab = cv2.cvtColor(img, cv2.COLOR_RGB2LAB)
    l, a, b = cv2.split(img_lab)

    # --- (1) Smooth illumination correction ---
    ksize = int(49 + intensity * 100)
    ksize = ksize + 1 if ksize % 2 == 0 else ksize
    l_blur = cv2.GaussianBlur(l, (ksize, ksize), 0)
    l_equal = cv2.divide(l, l_blur, scale=128)

    # --- (2) Frequency-domain smoothing ---
    freq = np.fft.fft2(l_equal)
    freqshift = np.fft.fftshift(freq)
    rows, cols = l_equal.shape
    crow, ccol = rows // 2, cols // 2
    mask = np.ones((rows, cols), np.uint8)
    r = int(10 + intensity * 30)
    mask[crow - r:crow + r, ccol - r:ccol + r] = 0
    freqshift = freqshift * mask
    ishift = np.fft.ifftshift(freqshift)
    l_flat = np.abs(np.fft.ifft2(ishift))
    l_flat = cv2.normalize(l_flat, None, 0, 255, cv2.NORM_MINMAX).astype(np.uint8)

    # --- (3) Merge back and convert ---
    img_lab = cv2.merge([l_flat, a, b])
    result = cv2.cvtColor(img_lab, cv2.COLOR_LAB2RGB)
    return Image.fromarray(result)

# -------------------------------------------------
# 🎛️ GRADIO BLOCKS UI (with slider throttle)
# -------------------------------------------------
with gr.Blocks(title="🪄 Saree Fold Remover (Lighting Only, No Distortion)") as demo:
    gr.Markdown(
        """
        ## 🪄 Saree Fold Remover  
        Upload a flat saree or fabric image.<br>
        Adjust the **Fold Intensity** slider — higher values remove deeper folds.<br>
        Lighting only correction, no geometric distortion.
        """
    )

    with gr.Row():
        inp = gr.Image(label="Upload Saree Image", type="pil")
        intensity = gr.Slider(
            0.0, 1.0, value=0.5, step=0.05,
            label="Fold Intensity",
            info="Higher = stronger flattening (lighting only)"
        )

    out = gr.Image(label="Flat, Fold-Free Output")

    # Button to trigger with throttle behavior
    run_btn = gr.Button("✨ Remove Folds")

    # Run when button is pressed (instead of live refresh)
    run_btn.click(fn=remove_folds, inputs=[inp, intensity], outputs=out)

demo.launch()