Update app.py

app.py

@@ -1,40 +1,147 @@
+import os, json, traceback
+from typing import List, Tuple, Optional
+
 import gradio as gr
-import joblib
-import json
-import os
-from PIL import Image
 import numpy as np
-from sklearn.neighbors import KNeighborsClassifier
+from PIL import Image
+import joblib
+
+# -------------------------------
+# Config (change if you used other names)
+# -------------------------------
+MODEL_FILE = os.getenv("MODEL_FILE", "my_model_k7.pkl")   # e.g. "my_model.pkl" or "my_model_k5.pkl"
+LABELS_FILE = os.getenv("LABELS_FILE", "labels.json")     # the labels you saved
+
+IMG_SIZE = 224  # MobileNetV2 default input size
+
+# -------------------------------
+# Feature extractor (MobileNetV2)
+# -------------------------------
+# We use TF/Keras MobileNetV2 (without the top) + global average pooling ("avg") for embeddings.
+# Ref: Keras Applications docs and MobileNetV2 API.
+try:
+    import tensorflow as tf
+    from tensorflow.keras.applications.mobilenet_v2 import (
+        MobileNetV2,
+        preprocess_input,
+    )
+
+    _feature_model = MobileNetV2(
+        weights="imagenet", include_top=False, pooling="avg",
+        input_shape=(IMG_SIZE, IMG_SIZE, 3)
+    )
+
+    def embed_image(img: Image.Image) -> np.ndarray:
+        img = img.convert("RGB").resize((IMG_SIZE, IMG_SIZE))
+        x = np.array(img, dtype=np.float32)
+        x = np.expand_dims(x, axis=0)
+        x = preprocess_input(x)
+        feat = _feature_model.predict(x, verbose=0)[0]
+        return feat
 
-# Load trained model
-MODEL_PATH = "my_model_k7.pkl"
-LABELS_PATH = "labels.json"
+    FE_MSG = "TF MobileNetV2 embeddings (pooling=avg)"
+except Exception as e:
+    # Very robust fallback (no TF): simple resize + flatten
+    _feature_model = None
+    FE_MSG = "Fallback embeddings (resize+flatten) — consider enabling TensorFlow for better accuracy."
+    def embed_image(img: Image.Image) -> np.ndarray:
+        img = img.convert("RGB").resize((IMG_SIZE, IMG_SIZE))
+        return (np.array(img, dtype=np.float32).ravel() / 255.0)
 
-clf = joblib.load(MODEL_PATH)
+# -------------------------------
+# Load model & labels (if present)
+# -------------------------------
+def _load_labels(path: str) -> Optional[List[str]]:
+    if not os.path.exists(path):
+        return None
+    with open(path, "r", encoding="utf-8") as f:
+        data = json.load(f)
+    # Accept list OR dict ({"0": "classA", ...})
+    if isinstance(data, list):
+        return [str(x) for x in data]
+    if isinstance(data, dict):
+        try:
+            # keys are indices as strings
+            return [data[str(i)] for i in sorted(map(int, data.keys()))]
+        except Exception:
+            # fallback: values in insertion order
+            return list(map(str, data.values()))
+    return None
 
-with open(LABELS_PATH, "r") as f:
-    labels = json.load(f)
+def _load_model_and_labels() -> Tuple[Optional[object], Optional[List[str]], str]:
+    model, labels, note = None, None, ""
+    if os.path.exists(MODEL_FILE):
+        try:
+            model = joblib.load(MODEL_FILE)
+        except Exception as e:
+            note += f"⚠️ Could not load model '{MODEL_FILE}': {e}\n"
+    else:
+        note += f"❗ Model file not found: `{MODEL_FILE}`. Upload it in the **Files** tab.\n"
 
-# Feature extractor (resize + flatten)
-def preprocess(img):
-    img = img.convert("RGB").resize((224, 224))
-    return np.array(img).flatten().reshape(1, -1)
+    labels = _load_labels(LABELS_FILE)
+    if labels is None:
+        note += f"❗ Labels file not found or unreadable: `{LABELS_FILE}`. Upload a JSON labels file (list or index→name dict).\n"
 
-# Prediction function
-def predict(image):
-    features = preprocess(image)
-    pred = clf.predict(features)[0]
-    return f"🌿 Predicted Plant: {labels[int(pred)]}"
+    if model and labels:
+        note = f"✅ Model & labels loaded. Feature extractor: **{FE_MSG}**"
+    elif not note:
+        note = "❗ Model and/or labels not loaded yet."
+    return model, labels, note
 
-# Gradio UI
-demo = gr.Interface(
-    fn=predict,
-    inputs=gr.Image(type="pil", label="Upload a Plant Image"),
-    outputs="text",
-    title="🌍 UAE Flora Explorer",
-    description="Upload an image of a UAE plant and the model will identify it."
-)
+clf, LABELS, STATUS = _load_model_and_labels()
+
+# -------------------------------
+# Prediction
+# -------------------------------
+def predict(img: Image.Image):
+    if img is None:
+        return "Please upload an image.", None
+
+    if clf is None or LABELS is None:
+        return (
+            "Model/labels missing. Upload your `*.pkl` and `labels.json` in the Space **Files** tab, "
+            "then click **Restart**.", None
+        )
+
+    try:
+        feat = embed_image(img).reshape(1, -1)
+        # Predicted class
+        pred_idx = int(clf.predict(feat)[0])
+        pred_label = LABELS[pred_idx] if 0 <= pred_idx < len(LABELS) else f"Class #{pred_idx}"
+
+        # Top-3 probabilities (if available)
+        top3 = None
+        if hasattr(clf, "predict_proba"):
+            probs = clf.predict_proba(feat)[0]
+            order = np.argsort(probs)[::-1][:min(3, len(probs))]
+            top3 = [
+                {"Class": LABELS[i] if i < len(LABELS) else f"#{i}",
+                 "Probability": float(probs[i])}
+                for i in order
+            ]
+
+        return pred_label, top3
+    except Exception as e:
+        tb = traceback.format_exc(limit=2)
+        return f"Error during prediction: {e}\n{tb}", None
+
+# -------------------------------
+# UI
+# -------------------------------
+with gr.Blocks(title="UAE Flora Classifier") as demo:
+    gr.Markdown(
+        "# 🌿 UAE Flora Classifier\n"
+        "Upload a plant photo and I’ll predict the species using a KNN classifier over MobileNetV2 embeddings."
+    )
+    gr.Markdown(f"**Status:** {STATUS}\n\n**Model file:** `{MODEL_FILE}`  •  **Labels file:** `{LABELS_FILE}`")
+
+    with gr.Row():
+        in_img = gr.Image(type="pil", label="Upload plant image")
+        out_label = gr.Label(label="Predicted class")
+    out_table = gr.Dataframe(headers=["Class", "Probability"], label="Top-3 (if available)", interactive=False)
+
+    btn = gr.Button("Predict")
+    btn.click(fn=predict, inputs=[in_img], outputs=[out_label, out_table])
 
 if __name__ == "__main__":
     demo.launch()
-