Update app.py

app.py

@@ -2,12 +2,19 @@ import gradio as gr
 import numpy as np
 from PIL import Image
 import torch
-from facenet_pytorch import InceptionResnetV1, MTCNN
-
-device = 'cpu'
+from facenet_pytorch import MTCNN, InceptionResnetV1
+import os
 
+# إعدادات
+device = "cpu"
 mtcnn = MTCNN(image_size=160, margin=20, device=device)
-model = InceptionResnetV1(pretrained='vggface2').eval()
+model = InceptionResnetV1(pretrained="vggface2").eval()
+
+DB_DIR = "db_images"
+os.makedirs(DB_DIR, exist_ok=True)
+
+embeddings = []
+images = []
 
 def get_embedding(img):
     face = mtcnn(img)
@@ -18,23 +25,58 @@ def get_embedding(img):
         emb = model(face)
     return emb[0].numpy()
 
-def compare(img1, img2):
-    e1 = get_embedding(img1)
-    e2 = get_embedding(img2)
+def rebuild_db():
+    global embeddings, images
+    embeddings, images = [], []
+    for fname in os.listdir(DB_DIR):
+        path = os.path.join(DB_DIR, fname)
+        try:
+            img = Image.open(path).convert("RGB")
+            emb = get_embedding(img)
+            if emb is not None:
+                embeddings.append(emb)
+                images.append(path)
+        except:
+            pass
+
+def add_to_db(img):
+    idx = len(os.listdir(DB_DIR))
+    path = os.path.join(DB_DIR, f"img_{idx}.jpg")
+    img.save(path)
+    rebuild_db()
+    return "تمت الإضافة"
+
+def search_similar(img):
+    if not embeddings:
+        return "قاعدة الصور فارغة", []
+
+    q = get_embedding(img)
+    if q is None:
+        return "لم يتم اكتشاف وجه", []
+
+    sims = []
+    for i, e in enumerate(embeddings):
+        s = np.dot(q, e) / (np.linalg.norm(q) * np.linalg.norm(e))
+        sims.append((s, images[i]))
 
-    if e1 is None or e2 is None:
-        return "لم يتم اكتشاف وجه واضح في إحدى الصورتين"
+    sims = sorted(sims, reverse=True)[:5]
+    results = [(Image.open(p), f"{round(s*100,2)}%") for s, p in sims]
+    return "نتائج البحث", results
 
-    sim = np.dot(e1, e2) / (np.linalg.norm(e1) * np.linalg.norm(e2))
-    return f"نسبة التشابه: {round(sim * 100, 2)} %"
+with gr.Blocks() as app:
+    gr.Markdown("## Face Search MVP")
 
-with gr.Blocks() as demo:
-    gr.Markdown("## Face Similarity MVP (Children Images)")
-    img1 = gr.Image(type="pil", label="صورة الطفل الأولى")
-    img2 = gr.Image(type="pil", label="صورة الطفل الثانية")
-    btn = gr.Button("قارن")
-    out = gr.Textbox()
+    with gr.Tab("إضافة صورة لقاعدة البيانات"):
+        img_add = gr.Image(type="pil")
+        btn_add = gr.Button("إضافة")
+        out_add = gr.Textbox()
+        btn_add.click(add_to_db, img_add, out_add)
 
-    btn.click(compare, inputs=[img1, img2], outputs=out)
+    with gr.Tab("بحث عن شبيه"):
+        img_q = gr.Image(type="pil")
+        btn_s = gr.Button("بحث")
+        out_txt = gr.Textbox()
+        out_gallery = gr.Gallery(columns=5)
+        btn_s.click(search_similar, img_q, [out_txt, out_gallery])
 
-demo.launch()
+app.launch()