import gradio as gr
import tensorflow as tf
import numpy as np
from PIL import Image
#from drought_detection.utilities import transform_user_img, make_fig, get_dataframe_data
STORAGE_LOCATION = f'gs://wagon-data-batch913-drought_detection/SavedModel/Model_3band_RGB' # GCP path
model = tf.saved_model.load(STORAGE_LOCATION)

def predict(image):
    image = transform_user_img(image)
    with tf.compat.v1.Session(graph=tf.Graph()) as sess:
        tf.compat.v1.saved_model.loader.load(sess, [tf.compat.v1.saved_model.tag_constants.SERVING], STORAGE_LOCATION)
        graph = tf.compat.v1.get_default_graph()
        prediction = sess.run('StatefulPartitionedCall:0',
                              feed_dict={'serving_default_keras_layer_input:0': image})
        return prediction
def drought_prediction(image):
    prediction = predict(image)
    result = max(prediction[0])
    if result == prediction[0][0]:
        return "Drought :desert: :sweat: :sunny: : Looks like your region is likely suffering from drought and there's not enough plant life in the area to feed any cows."
    elif result == prediction[0][1]:
        return "Drought risk :cow: :warning: Looks like your region is at risk of drought. The forage quality can only feed one cow per 20sqm area."
    elif result == prediction[0][2]:
        return "Possible drought risk :cow: :seedling: :cow: : Looks like your region is not suffering from a drought, however, it can likely only feed 2 cows per 20sqm area."
    elif result == prediction[0][3]:
        return "No Drought :herb: :cow: :cow: :cow: :deciduous_tree: : Looks like your region is healthy and can feed at least 3 cows per 20sqm area! Happy foraging!"

        
def image_prediction(image):
    df = get_dataframe_data(predict(image))
    x = df.iloc[:, 2]
    y = df.iloc[:, 1]
    fig = make_fig(df, x, y)
    return fig

def image_upload(file):
    img = Image.open(file.name)
    return img

inputs = gr.inputs.Image(upload=False)
outputs = gr.outputs.Image(type="plotly")
title = "Drought conditions in Northern Kenya"
description = "Applying deep learning and computer vision for drought resilience, using satellite images and human expert labels to detect drought conditions in Northern Kenya"

gr.Interface(fn=image_prediction, inputs=inputs, outputs=outputs, title=title, description=description, allow_flagging=False).launch()