import gradio as gr
import tensorflow as tf
import numpy as np
from huggingface_hub import from_pretrained_keras

def predict_diagnosis(loops):
  if loops.shape[2] == 4:
      loops = np.delete(loops, 3, axis=2)

  n_img = np.ndarray((1, 255, 906, 3))
  for i in range(n_img.shape[0]):
      n_img[i] = tf.keras.utils.img_to_array(tf.image.resize(loops, [255, 906]))

  n_img = n_img.astype(np.float32)/255.0

  model = from_pretrained_keras("layai/priya-loop-model")
  single_data = model.predict(n_img)
  single_data = single_data[0]
  return np.argmax(single_data)

input = gr.Image()
output = gr.Textbox(label='Predicted diagnosis')
iloops = gr.Interface( fn=predict_diagnosis,
              inputs=input,
              outputs=output,
              capture_session=True)

iloops.launch(debug=True)