import joblib
import json
import numpy as np
import gradio as gr
from tensorflow.keras.applications import MobileNetV2
from tensorflow.keras.applications.mobilenet_v2 import preprocess_input
from tensorflow.keras.preprocessing import image

# --- Paths to your model and labels ---
MODEL_PATH = "my_model_k7.pkl"
LABELS_PATH = "labels.json"

# --- Load model and labels ---
model_data = joblib.load(MODEL_PATH)      # this is a dict with model + labels
clf = model_data["model"]                 # your KNN classifier
labels = model_data.get("labels", None)   # try to load labels from inside dict

# If labels are stored separately, override
with open(LABELS_PATH, "r") as f:
    labels = json.load(f)

# --- Feature extractor (MobileNetV2 embeddings) ---
feature_extractor = MobileNetV2(weights="imagenet",
                                include_top=False,
                                pooling="avg")

def predict(img):
    # preprocess image
    img = image.img_to_array(img)
    img = np.expand_dims(img, axis=0)
    img = preprocess_input(img)

    # extract features
    features = feature_extractor.predict(img)

    # run through classifier
    pred_idx = int(clf.predict(features)[0])
    pred_label = labels[str(pred_idx)] if isinstance(labels, dict) else labels[pred_idx]

    return pred_label

# --- Gradio UI ---
iface = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil", label="Upload plant image"),
    outputs=gr.Label(num_top_classes=1, label="Predicted class"),
    title="🌿 UAE Flora Classifier",
    description="Upload a plant photo and I’ll predict the species using a KNN classifier over MobileNetV2 embeddings."
)

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