import gradio as gr
import numpy as np  
import tensorflow as tf


def classify_image(inp):
  new_image = np.zeros((120, 120, 4))

# Copy the original image to the first three channels of the new array
  new_image[:,:,0:3] = inp[:,:,:]

  model = tf.keras.models.load_model('poke.h5')
  prediction=np.argmax(model.predict(np.expand_dims(new_image,axis=0)))
  prediction=int(prediction)
  if prediction==1:
    return "Water"
  else:
    return "Fire"



gr.Interface(fn=classify_image, 
             inputs=gr.Image(shape=(120,120)),
             outputs=gr.Label(),
            
             ).launch()