#%%
import gradio as gr
import detect
import inpaint
import numpy as np
from PIL import Image

CSS = """
img, canvas { image-rendering: pixelated !important; }
#banner { max-width: 100%; }
#banner img { width: 100%; height: auto; display: block; }
"""

def saturate(rgb):
    """Push each pixel to S=1, V=1 while preserving hue. Greyscale -> white.
    4-bit quantised: all arithmetic fits in uint8."""
    q     = rgb >> 4
    maxc  = q.max(axis=-1, keepdims=True)
    minc  = q.min(axis=-1, keepdims=True)
    delta = maxc - minc
    diff  = q - minc
    out   = (diff * 15) // np.maximum(delta, 1)
    out   = out * 17
    return np.where(delta > 0, out, 255)

#%%
def process(img):
    mask = detect.predict(img.copy())
    out  = inpaint.fix(img, mask)
    overlay = np.array(img)
    overlay[mask] = saturate(overlay[mask])
    det_img = Image.fromarray(overlay)
    return out, (img, out), det_img

with gr.Blocks(title="Pixel denoiser", css=CSS) as demo:
    gr.HTML(
        '<img src="/gradio_api/file=flux_vae_peppered_banner_v2.svg" alt="banner">',
        elem_id="banner",
    )
    gr.Markdown('The Flux2 VAE sometimes produces purple/green/yellow pixels in its output. This app detects and removes those.')
    with gr.Row():
        inp = gr.Image(type="pil", label="Input",     height=420, format='png')
        det = gr.Image(type="pil", label="Detection", height=420, format='png')
        out = gr.Image(type="pil", label="Cleaned",   height=420, format='png')
    cmp = gr.ImageSlider(label="Before / After", format='png')

    gr.Examples(
        examples=[["example1.png"], ["example2.png"]],
        inputs=inp,
        outputs=[out, cmp, det],
        fn=process,
        cache_examples=True,
    )
    inp.change(process, inputs=inp, outputs=[out, cmp, det])

demo.launch(allowed_paths=["."])