import gradio as gr
import numpy as np
import PIL

import tensorflow as tf
from tensorflow import keras

def predict(img):
  img_cropped = np.array(img, dtype='float32')[:100, 15:-15, :] / 255
  img_bw = np.mean(img_cropped, axis=-1)

  # predict
  img_input = np.expand_dims(img_bw, axis=0)
  prediction = model.predict(img_input)[0]
  animals = ['common bottlenose dolphin', 'fin whale', 'risso dolphin', 'short finned pilot whale', 'sperm whale']

  #bw image for display
  im = PIL.Image.fromarray(np.uint8(img_bw*255))

  return [{animals[i]: float(prediction[i]) for i in range(len(animals))}, im]


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

iface = gr.Interface(predict,\
                     inputs = gr.Image(shape=(130, 120)),\
                     outputs = [gr.outputs.Label(num_top_classes=5),\
                                gr.Image(shape=(100, 100), image_mode='L')],\
                     examples = ["examples/DBUAC-BP-14005.jpg",\
                                 "examples/DBUAC-GG-08001.jpg",\
                                 "examples/DBUAC-GMA-10006.jpg",\
                                 "examples/DBUAC-PM-09046.jpg",\
                                 "examples/DBUAC-TT-15070.jpg"])

iface.launch()