Update app.py

app.py

@@ -1,59 +1,72 @@
 import gradio as gr
-import tensorflow as tf
-import numpy as np
 from PIL import Image
+import numpy as np
+import pickle
+import tensorflow as tf
 from tensorflow.keras.applications.efficientnet import preprocess_input
+import requests
+from io import BytesIO
 
-# ============================
-# Load TensorFlow SavedModel
-# ============================
-MODEL_PATH = "exported_model"
-IMG_SIZE = (224, 224)
-CLASS_NAMES = ["bad", "good", "very_good"]
+# load model + label encoder
+MODEL_SAVE_PATH = "best_model.h5"
+LABEL_ENCODER_PATH = "label_encoder.pkl"
 
-print("🔄 Loading SavedModel…")
-model = tf.saved_model.load(MODEL_PATH)
-infer = model.signatures["serving_default"]
-print("✅ Model loaded!")
+model = tf.keras.models.load_model(MODEL_SAVE_PATH)
+with open(LABEL_ENCODER_PATH, "rb") as f:
+    label_encoder = pickle.load(f)
 
+IMG_SIZE = model.input_shape[1:3]
 
-# ============================
-# Prediction Function
-# ============================
-def predict_guava_quality(image):
-    if image is None:
-        return "❌ Vui lòng tải ảnh!", 0.0
+def load_image_from_url(url):
+    """Tải ảnh từ URL và return PIL."""
+    try:
+        resp = requests.get(url, timeout=5)
+        img = Image.open(BytesIO(resp.content)).convert("RGB")
+        return img
+    except:
+        return None
 
-    img = Image.fromarray(image).convert("RGB")
-    img = img.resize(IMG_SIZE)
+def predict_fn(img, url):
+    """img: numpy image (upload), url: string"""
+    
+    # Ưu tiên dùng URL nếu có
+    if url and url.strip() != "":
+        img_pil = load_image_from_url(url)
+        if img_pil is None:
+            return "❌ Không tải được ảnh từ URL!", None
+    else:
+        # sử dụng ảnh upload
+        if img is None:
+            return "❌ Chưa cung cấp ảnh!", None
+        img_pil = Image.fromarray(img).convert("RGB")
 
-    arr = np.array(img).astype("float32")
+    # preprocess
+    img_resized = img_pil.resize(IMG_SIZE)
+    arr = np.array(img_resized).astype("float32")
     arr = preprocess_input(arr)
-    arr = np.expand_dims(arr, axis=0)
-
-    outputs = infer(tf.constant(arr))
-    preds = list(outputs.values())[0].numpy()[0]
+    arr = np.expand_dims(arr, 0)
 
-    idx = np.argmax(preds)
-    confidence = float(preds[idx])
-    label = CLASS_NAMES[idx]
+    preds = model.predict(arr)
+    idx = int(np.argmax(preds, axis=1)[0])
+    confidence = float(np.max(preds))
+    label = label_encoder.inverse_transform([idx])[0]
 
-    return f"🍈 {label}", confidence
+    return f"✅ {label} ", img_pil
 
 
-# ============================
-# Gradio UI
-# ============================
+# Giao diện Gradio
 demo = gr.Interface(
-    fn=predict_guava_quality,
-    inputs=gr.Image(type="numpy", label="Tải ảnh Ổi"),
+    fn=predict_fn,
+    inputs=[
+        gr.Image(type="numpy", label="Upload Image"),
+        gr.Textbox(label="Hoặc dán URL ảnh online")
+    ],
     outputs=[
-        gr.Textbox(label="Kết quả dự đoán"),
-        gr.Number(label="Độ tin cậy (0–1)")
+        gr.Textbox(label="Prediction"),
+        gr.Image(label="Preview Image")
     ],
-    title="Guava Quality Classifier",
-    description="Phân loại chất lượng quả Ổi: very_good / good / bad"
+    title="Guava Classifier",
+    description="Upload ảnh Ổi hoặc nhập URL ảnh để phân loại."
 )
 
-if __name__ == "__main__":
-    demo.launch()
+demo.launch(inline=True)