from huggingface_hub import from_pretrained_keras
import gradio as gr
from skimage.io import imsave
from skimage.color import lab2rgb
from tensorflow.keras.utils import img_to_array, load_img
import numpy as np

repo_id="saruizn/trabajo"

model = from_pretrained_keras(repo_id)

def predict(img):
    partes=[]
    imsave("input.jpg",(img * 255).astype(np.uint8))
    imag = img_to_array(load_img("input.jpg"))
    partes=divideImage(imag)
    i=0
    imagenes=[]
    for imgp in partes:
        nombre="result"+str(i)+".png"
        imsave(nombre,imgp)
        guardarImagen(nombre,i,imagenes)
        i=i+1
    imagenFinal=juntaImage(imagenes)
    return imagenFinal

def guardarImagen(foto,i,imagenes):
    color_me = []
    img = img_to_array(load_img(foto))
    img = resize(img ,(256,256))
    color_me.append(img)
    color_me = np.array(color_me, dtype=float)
    color_me = rgb2lab(1.0/255*color_me)[:,:,:,0]
    color_me = color_me.reshape(color_me.shape+(1,))
    
    output = model.predict(color_me)
    output = output * 128 * 8
    
      
    result = np.zeros((256, 256, 3))
    result[:,:,0] = color_me[0][:,:,0]
    result[:,:,1:] = output[0]
    imagenes.append(lab2rgb(result))
    imsave("result"+str(i)+".png", lab2rgb(result))

def divideImage(img):
    salida=[]
    for i in range(0,3):
        for j in range(0,5):
            salida.append(img[i*256:i*256+256,j*256:j*256+256])
        salida.append(img[i*256:i*256+256,1372-256:1372])
    for j in range(0,5):
        salida.append(img[929-256:929,j*256:j*256+256])
    salida.append(img[929-256:929,1372-256:1372])
    return salida

def juntaImage(X):
    salida=np.zeros((929, 1372, 3))
    im=0
    for i in range(0,3):
        for j in range(0,5):
            salida[i*256:i*256+256,j*256:j*256+256]=X[im]
            im=im+1
        salida[i*256:i*256+256,1372-256:1372]=X[im]
        im=im+1
    for j in range(0,5):
        salida[929-256:929,j*256:j*256+256]=X[im]
        im=im+1
    salida[929-256:929,1372-256:1372]=X[im]
    return salida




gr.Interface(fn=predict, inputs=gr.inputs.Image(shape=(929, 1372)), outputs="image").launch(share=False)