import gradio as gr
import tensorflow as tf
from tensorflow.keras.models import load_model
from PIL import Image
import numpy as np
import requests
from io import BytesIO
from huggingface_hub import hf_hub_download


# Download the model file from the repo
local_path = hf_hub_download(
    repo_id="Ben041/soil-type-classifier",
    filename="soil_type_model.h5"
)

model = load_model(local_path)

# --- Define class labels (matching the 11 categories) ---
class_labels = [
    "Alluvial soil", "Black Soil", "Cinder Soil",
    "Clayey soils", "Laterite soil", "Loamy soil",
    "Peat Soil", "Sandy loam", "Sandy soil",
    "Yellow Soil"
]

# --- Prediction helper ---
def predict_soil_type(model, image_path, class_labels):
    img = Image.open(image_path).convert("RGB")
    img = img.resize((224, 224))
    arr = np.array(img) / 255.0
    arr = np.expand_dims(arr, axis=0)
    preds = model.predict(arr)
    idx = np.argmax(preds[0])
    confidence = float(preds[0][idx])
    soil_type = class_labels[idx]
    return soil_type, confidence

# --- Gradio wrapper ---
def classify(image):
    # image comes as a PIL image via Gradio
    image.save("temp.jpg")
    soil_type, confidence = predict_soil_type(model, "temp.jpg", class_labels)
    return { "Soil Type": soil_type, "Confidence": f"{confidence*100:.2f}%" }

# --- Build Gradio interface ---
app = gr.Interface(
    fn=classify,
    inputs=gr.Image(type="pil", label="Upload Soil Image"),
    outputs=gr.JSON(label="Prediction Result"),
    title="🌱 Soil Type Classifier",
    description="CNN model trained to classify 11 soil types from images."
)


if __name__ == "__main__":
    app.launch(show_error=True)