Update app.py

app.py

@@ -6,19 +6,28 @@ import tensorflow as tf
 from tensorflow.keras.applications.efficientnet import preprocess_input
 import requests
 from io import BytesIO
+from ultralytics import YOLO
 
-# load model + label encoder
+# ==== Load models ====
 MODEL_SAVE_PATH = "guava_model.keras"
 LABEL_ENCODER_PATH = "label_encoder.pkl"
 
+# ENB0
 model = tf.keras.models.load_model(MODEL_SAVE_PATH)
 with open(LABEL_ENCODER_PATH, "rb") as f:
     label_encoder = pickle.load(f)
 
+# YOLOv8 (model bạn đã train)
+YOLO_MODEL_PATH = "yolov8_guava.pt"
+yolo_model = YOLO(YOLO_MODEL_PATH)
+
 IMG_SIZE = model.input_shape[1:3]
 
+
+# ======================================================
+#          Load image (Upload or URL)
+# ======================================================
 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")
@@ -26,21 +35,25 @@ def load_image_from_url(url):
     except:
         return None
 
-def predict_fn(img, url):
-    """img: numpy image (upload), url: string"""
-    
-    # Ưu tiên dùng URL nếu có
+
+# ======================================================
+#        Prediction function (ENB0 + YOLOv8)
+# ======================================================
+def compare_models(img, url):
+
+    # --- Ưu tiên URL ---
     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
+            return "❌ Không tải được ảnh từ URL!", None, None, None
     else:
-        # sử dụng ảnh upload
         if img is None:
-            return "❌ Chưa cung cấp ảnh!", None
+            return "❌ Chưa cung cấp ảnh!", None, None, None
         img_pil = Image.fromarray(img).convert("RGB")
 
-    # preprocess
+    # ======================================================
+    #                EfficientNetB0 Prediction
+    # ======================================================
     img_resized = img_pil.resize(IMG_SIZE)
     arr = np.array(img_resized).astype("float32")
     arr = preprocess_input(arr)
@@ -49,24 +62,51 @@ def predict_fn(img, url):
     preds = model.predict(arr)
     idx = int(np.argmax(preds, axis=1)[0])
     confidence = float(np.max(preds))
-    label = label_encoder.inverse_transform([idx])[0]
+    label_enb0 = label_encoder.inverse_transform([idx])[0]
+    enb0_text = f"{label_enb0} (Conf {confidence:.2f})"
+
+    # ======================================================
+    #                    YOLOv8 Prediction
+    # ======================================================
+    results = yolo_model(img_pil)
+    result = results[0]
+
+    # Lấy label + conf cao nhất
+    if len(result.boxes) > 0:
+        best = result.boxes[0]
+        yolo_label_id = int(best.cls[0])
+        yolo_conf = float(best.conf[0])
+        yolo_label = yolo_model.model.names[yolo_label_id]
+        yolo_text = f"{yolo_label} (Conf {yolo_conf:.2f})"
+
+        # Vẽ bounding box
+        img_yolo = result.plot()   # return np array BGR
+        img_yolo = Image.fromarray(img_yolo[..., ::-1])  # convert to RGB
+    else:
+        yolo_text = "Không phát hiện!"
+        img_yolo = img_pil
 
-    return f"✅ {label} ", img_pil
+    # return 4 outputs
+    return enb0_text, yolo_text, img_pil, img_yolo
 
 
-# Giao diện Gradio
+# ======================================================
+#                   Gradio UI
+# ======================================================
 demo = gr.Interface(
-    fn=predict_fn,
+    fn=compare_models,
     inputs=[
         gr.Image(type="numpy", label="Upload Image"),
-        gr.Textbox(label="Hoặc dán URL ảnh online")
+        gr.Textbox(label="Hoặc nhập URL ảnh")
     ],
     outputs=[
-        gr.Textbox(label="Prediction"),
-        gr.Image(label="Preview Image")
+        gr.Textbox(label="ENB0 Prediction"),
+        gr.Textbox(label="YOLOv8 Prediction"),
+        gr.Image(label="Original Image"),
+        gr.Image(label="YOLOv8 Detection Image")
     ],
-    title="Guava Classifier",
-    description="Upload ảnh Ổi hoặc nhập URL ảnh để phân loại."
+    title="Guava Classifier — YOLOv8 vs EfficientNetB0",
+    description="So sánh kết quả phân loại giữa YOLOv8 và EfficientNetB0."
 )
 
 demo.launch(inline=True)