Update app.py

app.py

@@ -9,7 +9,7 @@ from qdrant_client import QdrantClient
 from qdrant_client.http.models import Filter, FieldCondition, MatchValue
 import os
 
-# โหลด EfficientNet-B2
+# โหลด EfficientNet-B7
 model = EfficientNet.from_pretrained('efficientnet-b7')
 model.eval()
 
@@ -54,32 +54,43 @@ def get_image_embedding(image_path_or_url):
         embedding = model(image_tensor)
     return embedding
 
-# ค้นหาภาพที่คล้ายกัน
+# NEW: ดึง embedding จาก PIL.Image โดยตรง
+def get_image_embedding_from_pil(image: Image.Image):
+    image_tensor = preprocess(image).unsqueeze(0)
+    with torch.no_grad():
+        embedding = model(image_tensor)
+    return embedding
+
+# ปรับฟังก์ชันหลักให้เลือกเส้นทางที่ถูกต้อง
 def search_similar_images(image_upload, image_url):
     # กำหนดแหล่งภาพ
     if image_upload is not None:
-        image = image_upload.convert("RGB")
+        query_vector = get_image_embedding_from_pil(image_upload)
     elif image_url:
-        response = requests.get(image_url)
-        image = Image.open(BytesIO(response.content)).convert("RGB")
+        query_vector = get_image_embedding(image_url)
     else:
         return [], "❌ Please upload an image or provide an image URL."
 
-    query_vector = get_image_embedding(image)
+    # ป้องกันกรณีที่ embedding ผิดพลาด
+    if query_vector is None:
+        return [], "❌ Could not process the image."
 
     # Query จาก Qdrant
-    result = qdrant_client.search(
-        collection_name="image_collection_b7",
-        query_vector=query_vector[0].tolist(),
-        with_payload=True,
-    )
+    try:
+        result = qdrant_client.search(
+            collection_name="image_collection_b7",
+            query_vector=query_vector[0].tolist(),
+            with_payload=True,
+        )
+    except Exception as e:
+        return [], f"❌ Qdrant search failed: {str(e)}"
 
     if not result:
         return [], "❌ No matching images found."
 
     gallery_data = []
     for r in result:
-        image_url = r.payload['image_url']
+        image_url = r.payload.get('image_url', '')
         name = r.payload.get('name', 'Unknown')
         score = round(r.score, 3)
         caption = f"{name}\n(score: {score})"
@@ -87,8 +98,9 @@ def search_similar_images(image_upload, image_url):
 
     return gallery_data, f"✅ Found {len(result)} similar images."
 
+
 # Gradio Interface
-with gr.Blocks(title="Image Similarity Search with EfficientNet-B2") as demo:
+with gr.Blocks(title="Image Similarity Search with EfficientNet-B7") as demo:
     gr.Markdown("# 🔍 EfficientNet-B7 Image Search")
     gr.Markdown("Upload an image or paste a URL to find similar images using Qdrant.")
 
@@ -108,4 +120,4 @@ with gr.Blocks(title="Image Similarity Search with EfficientNet-B2") as demo:
     )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()