import keras.backend as K
import gradio as gr
import numpy as np

def psnr(y_true, y_pred):
    return -10*K.log(K.mean(K.flatten((y_true - y_pred))**2)) / np.log(10)

from keras.models import load_model
model = load_model("./MyNet.h5", custom_objects={'psnr': psnr, 'val_psnr': psnr})

# ラップする関数
#def sepia(input_img):
#    sepia_img = np.asarray(input_img)
#    sepia_img = sepia_img.astype('float32')
#    sepia_img = sepia_img / 255.0
#    sepia_img = model.predict(sepia_img)
#    return sepia_img

def sepia(inp):
    inp = inp.reshape((-1, 256, 256, 3))
    inp = np.asarray(inp)
    inp = inp.astype('float32')
    inp = inp / 255.0
    sepia_img = model.predict(inp)
    return sepia_img



# シンプルなUIを作成
demo = gr.Interface(
    fn=sepia, 
    inputs=gr.Image(shape=(256, 256)), 
    outputs="image"
).launch()


#image = gr.inputs.Image(shape=(256,256))
#image = np.asarray(image)
#image = image.astype('float32')
#image = image / 255.0
#decoded_imgs = model.predict(image)

#decoded_imgs.reshape(256,256,3)
#prediction=model.predict(img_4d)[0]


#gr.Interface(inputs=image, outputs=decoded_imgs,interpretation='default').launch(debug='True')