Update app.py

app.py

@@ -5,70 +5,96 @@ from torchvision import models
 from PIL import Image
 import gradio as gr
 import numpy as np
+import cv2
 from sklearn.metrics.pairwise import cosine_similarity
 
-# โหลดโมเดล pretrained ResNet50 (ตัด fully connected layer ออก)
+# โหลด ResNet50 และตัด FC Layer ออก
 resnet = models.resnet50(pretrained=True)
 resnet.eval()
 feature_extractor = torch.nn.Sequential(*list(resnet.children())[:-1])
 
-# กำหนด transform รูปภาพให้ตรงกับ input โมเดล
+# Transform ให้เข้ากับ input โมเดล
 transform = transforms.Compose([
-    transforms.Resize((224,224)),
+    transforms.Resize((224, 224)),
     transforms.ToTensor(),
-    transforms.Normalize(mean=[0.485, 0.456, 0.406],
-                         std=[0.229, 0.224, 0.225]),
+    transforms.Normalize(
+        mean=[0.485, 0.456, 0.406],
+        std=[0.229, 0.224, 0.225]
+    ),
 ])
 
-# กำหนดชื่อไฟล์โลโก้ธนาคารในโฟลเดอร์เดียวกับ app.py
+# รายชื่อไฟล์โลโก้ที่ใช้เป็นฐานข้อมูล
 bank_logo_files = [
     "scb.jpg",
     "kbank.jpg",
     "bbl.jpeg",
-    # เพิ่มชื่อไฟล์โลโก้ธนาคารของคุณที่นี่
+    # เพิ่มรูปอื่นๆ ได้
 ]
 
-# โหลดโลโก้ธนาคารและสร้าง embedding
+# ===== 📌 ฟังก์ชันช่วย: ตัดโลโก้จากภาพด้วย OpenCV =====
+def crop_logo(image_pil):
+    image_np = np.array(image_pil)
+    gray = cv2.cvtColor(image_np, cv2.COLOR_RGB2GRAY)
+    blurred = cv2.GaussianBlur(gray, (5, 5), 0)
+    _, thresh = cv2.threshold(blurred, 200, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)
+
+    # หา contour ที่ใหญ่สุด (คิดว่าเป็นโลโก้)
+    contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+    if contours:
+        largest = max(contours, key=cv2.contourArea)
+        x, y, w, h = cv2.boundingRect(largest)
+        cropped = image_np[y:y+h, x:x+w]
+        return Image.fromarray(cropped)
+    else:
+        return image_pil  # ถ้าไม่เจอ contour ก็ส่งภาพเดิมกลับ
+
+# ===== 📌 โหลดฐานข้อมูลโลโก้และสร้าง embedding =====
 def load_bank_logo_embeddings():
     bank_embeddings = {}
     for filename in bank_logo_files:
         if os.path.isfile(filename):
             image = Image.open(filename).convert("RGB")
+            image = crop_logo(image)
             img_t = transform(image).unsqueeze(0)
             with torch.no_grad():
                 emb = feature_extractor(img_t).squeeze().numpy()
-                emb = emb / np.linalg.norm(emb)  # Normalize embedding
+                emb = emb / np.linalg.norm(emb)
             bank_name = os.path.splitext(filename)[0]
             bank_embeddings[bank_name] = emb
         else:
-            print(f"Warning: ไฟล์ {filename} ไม่พบในโฟลเดอร์")
+            print(f"⚠️ ไม่พบไฟล์: {filename}")
     return bank_embeddings
 
 bank_embeddings = load_bank_logo_embeddings()
 
-# ฟังก์ชันทำนายธนาคารโดยเทียบ similarity
+# ===== 📌 ฟังก์ชันเทียบภาพที่อัปโหลดกับฐานข้อมูล =====
 def predict_bank_logo(input_img):
     if input_img.mode != "RGB":
         input_img = input_img.convert("RGB")
-    img_t = transform(input_img).unsqueeze(0)
+
+    cropped_img = crop_logo(input_img)
+    img_t = transform(cropped_img).unsqueeze(0)
+
     with torch.no_grad():
         emb = feature_extractor(img_t).squeeze().numpy()
         emb = emb / np.linalg.norm(emb)
+
     sims = {}
     for bank, bank_emb in bank_embeddings.items():
         sim = cosine_similarity([emb], [bank_emb])[0][0]
         sims[bank] = sim
+
     sorted_sims = dict(sorted(sims.items(), key=lambda item: item[1], reverse=True))
-    top3 = {k: round(v, 4) for k,v in list(sorted_sims.items())[:3]}
+    top3 = {k: round(v, 4) for k, v in list(sorted_sims.items())[:3]}
     return top3
 
-# สร้าง Gradio interface
+# ===== 📌 Gradio UI =====
 interface = gr.Interface(
     fn=predict_bank_logo,
-    inputs=gr.Image(type="pil", label="Upload bank logo"),
-    outputs=gr.Label(num_top_classes=3, label="Top 3 Matching Bank Logos"),
-    title="Bank Logo Similarity Checker",
-    description="Upload a bank logo image and find the most similar bank logos from the database."
+    inputs=gr.Image(type="pil", label="📤 Upload Bank Logo Image"),
+    outputs=gr.Label(num_top_classes=3, label="🏦 Top 3 Matching Bank Logos"),
+    title="🔍 Bank Logo Similarity Checker",
+    description="ระบบเปรียบเที��บโลโก้ธนาคารจากภาพที่อัปโหลด โดยใช้โมเดล ResNet50 + Cosine Similarity พร้อมการตัดโลโก้อัตโนมัติ"
 )
 
 if __name__ == "__main__":