app.py

import numpy as np
import tensorflow as tf
from skimage import color
from PIL import Image
import gradio as gr

# 1. Carga tu modelo
model = tf.keras.models.load_model('generator_final.h5', compile=False)
IMG_SIZE = 256

def lab_to_rgb(lab_norm):
    # a) Desnormaliza L*, a*, b*
    L  = (lab_norm[..., 0] + 1.0) * 50.0
    ab = lab_norm[..., 1:] * 128.0
    lab = np.concatenate([L[..., None], ab], axis=-1)
    # b) Convierte de Lab a RGB en [0,1] y escala a [0,255]
    rgb = color.lab2rgb(lab.clip([0, -128, -128], [100, 128, 128]))
    return (rgb * 255).astype(np.uint8)

def colorize(input_image):
    # 1. Asegura un objeto PIL RGB (descarta alfa, escala de grises, etc.)
    pil_img = Image.fromarray(input_image).convert('RGB')
    # 2. Redimensiona a 256×256
    pil_img = pil_img.resize((IMG_SIZE, IMG_SIZE))
    # 3. Pasa a array float32 en [0,1]
    img = np.array(pil_img).astype('float32') / 255.0
    # 4. Convierte a Lab y extrae L* (0–100)
    lab = color.rgb2lab(img)
    L   = lab[..., 0]
    # 5. Normaliza L* a [-1,1]
    L_norm = (L / 50.0) - 1.0
    # 6. Prepara input shape (1,256,256,1)
    input_L = np.expand_dims(L_norm[..., None], axis=0)
    # 7. Predice ab normalizados
    ab_pred = model.predict(input_L)[0]
    # 8. Reconstruye Lab normalizado y devuelve RGB uint8
    lab_norm = np.concatenate([L_norm[..., None], ab_pred], axis=-1)
    return lab_to_rgb(lab_norm)

# Interfaz Gradio (igual que antes)
iface = gr.Interface(
    fn=colorize,
    inputs=gr.Image(type='numpy', label="Blanco y negro o color"),
    outputs=gr.Image(type='numpy', label="Coloreada"),
    title="Colorización GAN + U-Net",
    description="Sube cualquier imagen; se redimensiona y colorea."
)

if __name__ == "__main__":
    iface.launch()