Update app.py

app.py

@@ -1,55 +1,106 @@
+import os
 import gradio as gr
+from PIL import Image
 import torch
+import numpy as np
 from facenet_pytorch import MTCNN, InceptionResnetV1
-from PIL import Image
-from utils import fetch_image_urls, download_image
-import os
+from torch.nn.functional import cosine_similarity
+import requests
+from bs4 import BeautifulSoup
+from io import BytesIO
 
-# Initialize face detection and recognition models
+# Initialize models
 mtcnn = MTCNN(image_size=160, margin=20, keep_all=False)
 resnet = InceptionResnetV1(pretrained='vggface2').eval()
 
-def get_embedding(img):
-    face = mtcnn(img)
-    if face is not None:
-        with torch.no_grad():
-            embedding = resnet(face.unsqueeze(0))
-        return embedding
+# Directory for temporary images
+DB_DIR = "scraped_images"
+os.makedirs(DB_DIR, exist_ok=True)
+
+# Scrape image URLs from Bing
+def fetch_image_urls(query, max_images=3):
+    headers = {"User-Agent": "Mozilla/5.0"}
+    search_url = f"https://www.bing.com/images/search?q={query.replace(' ', '+')}"
+    response = requests.get(search_url, headers=headers)
+    soup = BeautifulSoup(response.text, 'html.parser')
+    image_elements = soup.find_all('a', class_='iusc')
+    urls = []
+    for elem in image_elements[:max_images]:
+        m = elem.get('m')
+        if m:
+            try:
+                m_json = eval(m)
+                urls.append(m_json['murl'])
+            except:
+                continue
+    return urls
+
+# Download and save images
+def download_images(image_urls):
+    for i, url in enumerate(image_urls):
+        try:
+            response = requests.get(url)
+            img = Image.open(BytesIO(response.content)).convert('RGB')
+            img.save(os.path.join(DB_DIR, f"scraped_{i}.jpg"))
+        except:
+            continue
+
+# Load embeddings for downloaded images
+def load_scraped_embeddings():
+    embeddings = {}
+    for file in os.listdir(DB_DIR):
+        if file.lower().endswith(('.jpg', '.png')):
+            img_path = os.path.join(DB_DIR, file)
+            img = Image.open(img_path).convert("RGB")
+            face = mtcnn(img)
+            if face is not None:
+                face = face.unsqueeze(0)
+                emb = resnet(face)
+                embeddings[file] = emb
+    return embeddings
+
+# Compare uploaded face with scraped images
+def identify_person(uploaded_img, search_query):
+    if uploaded_img is None:
+        return "Please upload an image."
+
+    # Step 1: Scrape and download images
+    for f in os.listdir(DB_DIR):
+        os.remove(os.path.join(DB_DIR, f))
+    image_urls = fetch_image_urls(search_query, max_images=5)
+    download_images(image_urls)
+
+    # Step 2: Load scraped embeddings
+    db_embeddings = load_scraped_embeddings()
+
+    # Step 3: Get uploaded face embedding
+    uploaded_face = mtcnn(uploaded_img.convert("RGB"))
+    if uploaded_face is None:
+        return "No face detected in the uploaded image."
+    uploaded_embedding = resnet(uploaded_face.unsqueeze(0))
+
+    # Step 4: Compare and find best match
+    best_match = None
+    best_score = -1
+    for name, emb in db_embeddings.items():
+        score = cosine_similarity(uploaded_embedding, emb).item()
+        if score > best_score:
+            best_score = score
+            best_match = name
+
+    if best_score > 0.7:
+        return f"Best Match: {best_match}\nSimilarity: {best_score:.2f}"
     else:
-        return None
-
-def compare_faces(img1, img2):
-    emb1 = get_embedding(img1)
-    emb2 = get_embedding(img2)
-    if emb1 is None or emb2 is None:
-        return "Face not detected in one or both images."
-    similarity = torch.nn.functional.cosine_similarity(emb1, emb2).item()
-    return f"Similarity Score: {similarity:.2f}"
-
-def fetch_and_display_images(query):
-    urls = fetch_image_urls(query)
-    images = []
-    for url in urls:
-        path = download_image(url)
-        if path:
-            images.append(Image.open(path))
-    return images
-
-with gr.Blocks() as demo:
-    gr.Markdown("# 🔍 Face Similarity Checker with Web Scraping")
-    
-    with gr.Tab("Compare Uploaded Images"):
-        with gr.Row():
-            img1 = gr.Image(label="Image 1")
-            img2 = gr.Image(label="Image 2")
-        compare_button = gr.Button("Compare Faces")
-        result = gr.Textbox(label="Result")
-        compare_button.click(fn=compare_faces, inputs=[img1, img2], outputs=result)
-    
-    with gr.Tab("Fetch Images from Web"):
-        query = gr.Textbox(label="Search Query")
-        fetch_button = gr.Button("Fetch Images")
-        gallery = gr.Gallery(label="Fetched Images")
-        fetch_button.click(fn=fetch_and_display_images, inputs=query, outputs=gallery)
-
-demo.launch()
+        return f"No confident match found.\nBest candidate: {best_match} (Score: {best_score:.2f})"
+
+# Gradio Interface
+iface = gr.Interface(
+    fn=identify_person,
+    inputs=[gr.Image(type="pil"), gr.Text(label="Search Query (e.g., Elon Musk)")],
+    outputs="text",
+    title="Reverse Face Search via Web Scraping",
+    description="Upload a face image and enter a name or query. The app scrapes the web for images and tries to identify the person."
+)
+
+if __name__ == "__main__":
+    iface.launch()