import tensorflow as tf
import numpy as np
from PIL import Image
import numpy as np
import gradio as gr

loaded_model = tf.keras.models.load_model("final_model.h5")


def classify_organic_recycle(image):
    pil_image = Image.fromarray(image)
    pil_image = pil_image.resize((224, 224))
    resized_image = np.array(pil_image)
    resized_image = resized_image.reshape(-1, 224, 224, 3)
    prediction = loaded_model.predict(resized_image).tolist()[0]

    if prediction[0] > prediction[1]:
        return {
            "class": "Organic",
            "score": prediction[0],
        }
    else:
        return {
            "class": "Recycle",
            "score": prediction[1],
        }


demo = gr.Interface(
    fn=classify_organic_recycle,
    inputs="image",
    outputs="text",
    title="Organic vs Recycle",
    description="This model classifies images as organic or recycle",
)


if __name__ == "__main__":
    demo.launch()