Update app.py

app.py

@@ -8,24 +8,29 @@ import requests
 from io import BytesIO
 from ultralytics import YOLO
 
-# ==== Load models ====
+# ======================================================
+#                 Load Models
+# ======================================================
 MODEL_SAVE_PATH = "guava_model.keras"
 LABEL_ENCODER_PATH = "label_encoder.pkl"
+YOLO_MODEL_PATH = "yolov8_guava.pt"
 
-# ENB0
+# Load ENB0 model
 model = tf.keras.models.load_model(MODEL_SAVE_PATH)
+
+# Load label encoder
 with open(LABEL_ENCODER_PATH, "rb") as f:
     label_encoder = pickle.load(f)
 
-# YOLOv8 (model bạn đã train)
-YOLO_MODEL_PATH = "yolov8_guava.pt"
+# Load YOLOv8
 yolo_model = YOLO(YOLO_MODEL_PATH)
 
-IMG_SIZE = model.input_shape[1:3]
+# Lấy đúng input size của model
+IMG_SIZE = (model.input_shape[1], model.input_shape[2])
 
 
 # ======================================================
-#          Load image (Upload or URL)
+#            Load image from URL
 # ======================================================
 def load_image_from_url(url):
     try:
@@ -37,61 +42,56 @@ def load_image_from_url(url):
 
 
 # ======================================================
-#        Prediction function (ENB0 + YOLOv8)
+#                Predict Function
 # ======================================================
 def compare_models(img, url):
 
     # --- Ưu tiên URL ---
-    if url and url.strip() != "":
+    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, None, None
+            return "❌ URL không hợp lệ!", None, None, None
     else:
         if img is None:
             return "❌ Chưa cung cấp ảnh!", None, None, None
+
         img_pil = Image.fromarray(img).convert("RGB")
 
-    # ======================================================
-    #                EfficientNetB0 Prediction
-    # ======================================================
+    # ===================== ENB0 =======================
     img_resized = img_pil.resize(IMG_SIZE)
-    arr = np.array(img_resized).astype("float32")
+    arr = np.array(img_resized, dtype=np.float32)
     arr = preprocess_input(arr)
-    arr = np.expand_dims(arr, 0)
+    arr = np.expand_dims(arr, axis=0)
 
     preds = model.predict(arr)
     idx = int(np.argmax(preds, axis=1)[0])
     confidence = float(np.max(preds))
-    label_enb0 = label_encoder.inverse_transform([idx])[0]
-    enb0_text = f"{label_enb0} (Conf {confidence:.2f})"
+    label = label_encoder.inverse_transform([idx])[0]
 
-    # ======================================================
-    #                    YOLOv8 Prediction
-    # ======================================================
+    enb0_text = f"{label} (Conf {confidence:.2f})"
+
+    # ===================== YOLO =======================
     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
+        cls_id = int(best.cls[0])
+        conf = float(best.conf[0])
+        yolo_label = yolo_model.model.names[cls_id]
+        yolo_text = f"{yolo_label} (Conf {conf:.2f})"
+
+        img_yolo = result.plot()
+        img_yolo = Image.fromarray(img_yolo[..., ::-1])
     else:
         yolo_text = "Không phát hiện!"
         img_yolo = img_pil
 
-    # return 4 outputs
     return enb0_text, yolo_text, img_pil, img_yolo
 
 
 # ======================================================
-#                   Gradio UI
+#                    Gradio UI
 # ======================================================
 demo = gr.Interface(
     fn=compare_models,
@@ -109,4 +109,4 @@ demo = gr.Interface(
     description="So sánh kết quả phân loại giữa YOLOv8 và EfficientNetB0."
 )
 
-demo.launch(inline=True)
+demo.launch()