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

# Ensure model folder exists
os.makedirs("model", exist_ok=True)

# Download the model from Hugging Face if not already present
model_path = "model/mobnet_model.keras"
if not os.path.exists(model_path):
    url = "https://huggingface.co/ahmzakif/TrashNet-Classification/resolve/main/model/mobnet_model.keras"
    r = requests.get(url)
    with open(model_path, "wb") as f:
        f.write(r.content)

# Load Keras model
model = tf.keras.models.load_model(model_path)

# TrashNet classes
classes = ["cardboard", "glass", "metal", "paper", "plastic", "trash"]

# Image preprocessing
def predict(image: Image.Image):
    image = image.convert("RGB").resize((224, 224))
    x = np.array(image, dtype=np.float32) / 255.0
    x = np.expand_dims(x, axis=0)
    
    preds = model.predict(x)[0]
    scores = {classes[i]: float(preds[i]) for i in range(len(classes))}
    top_class = max(scores, key=scores.get)
    
    return {"prediction": top_class, "scores": scores}

# Gradio interface
iface = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil"),
    outputs="json",
    title="TrashNet Classification API",
    description="Upload an image of trash to get its classification."
)

iface.launch()