import gradio as gr
import numpy as np
import PIL

import tensorflow as tf
from tensorflow import keras

def predict(img):
  input = np.array(img, dtype='float32') / 255

  sepia_filter = np.array([[.393, .769, .189], [.349, .686, .168], [.272, .534, .131]])
  sepia_img = input.dot(sepia_filter.T)
  sepia_img /= sepia_img.max()

  mask = np.argmax(model.predict(np.expand_dims(input, axis=0)), axis=-1)[0]

  result = np.copy(input)
  for r in range(result.shape[0]):
    for c in range(result.shape[1]):
      if mask[r, c] != 1:
        result[r, c] = sepia_img[r, c]

  return PIL.Image.fromarray(np.uint8(result*255))


model = keras.models.load_model('model')

iface = gr.Interface(predict,\
                    inputs = gr.Image(shape=(256, 256)),\
                    outputs = gr.Image(shape=(256, 256), image_mode='rgb'),\
                    examples = ["examples/english_setter_78.jpg",\
                                "examples/Ragdoll_60.jpg",\
                                "examples/pomeranian_74.jpg",\
                                "examples/Persian_137.jpg",\
                                "examples/saint_bernard_136.jpg"])

iface.launch()