Create app.py

app.py

@@ -0,0 +1,339 @@
+import os
+import zipfile
+import torch
+import clip
+import numpy as np
+from PIL import Image
+import gradio as gr
+import openai
+from tqdm import tqdm
+from glob import glob
+import psycopg2
+from psycopg2.extras import execute_values
+import json
+import time
+
+# ─────────────────────────────────────────────
+# 📂 STEP 1: UNZIP TO CORRECT STRUCTURE
+# ─────────────────────────────────────────────
+zip_name = "lfw-faces.zip"
+unzip_dir = "lfw-faces"
+
+if not os.path.exists(unzip_dir):
+    print("🔓 Unzipping...")
+    with zipfile.ZipFile(zip_name, "r") as zip_ref:
+        zip_ref.extractall(unzip_dir)
+    print("✅ Unzipped into:", unzip_dir)
+
+# True image root after unzip
+img_root = os.path.join(unzip_dir, "lfw-deepfunneled")
+
+# ─────────────────────────────────────────────
+# 🗄️ STEP 2: DATABASE SETUP
+# ─────────────────────────────────────────────
+def setup_database():
+    """Setup PostgreSQL with pgvector extension"""
+    # Database configuration
+    DB_CONFIG = {
+        "dbname": "face_matcher",
+        "user": "postgres",
+        "password": "postgres",  # Change this to your actual password
+        "host": "localhost",
+        "port": "5432"
+    }
+    
+    try:
+        # Connect to PostgreSQL server to create database if it doesn't exist
+        conn = psycopg2.connect(
+            dbname="postgres", 
+            user=DB_CONFIG["user"], 
+            password=DB_CONFIG["password"], 
+            host=DB_CONFIG["host"]
+        )
+        conn.autocommit = True
+        cur = conn.cursor()
+        
+        # Create database if it doesn't exist
+        cur.execute(f"SELECT 1 FROM pg_catalog.pg_database WHERE datname = '{DB_CONFIG['dbname']}'")
+        exists = cur.fetchone()
+        if not exists:
+            cur.execute(f"CREATE DATABASE {DB_CONFIG['dbname']}")
+            print(f"Database {DB_CONFIG['dbname']} created.")
+        
+        cur.close()
+        conn.close()
+        
+        # Connect to the face_matcher database
+        conn = psycopg2.connect(**DB_CONFIG)
+        conn.autocommit = True
+        cur = conn.cursor()
+        
+        # Create pgvector extension if it doesn't exist
+        cur.execute("CREATE EXTENSION IF NOT EXISTS vector")
+        
+        # Create faces table if it doesn't exist
+        cur.execute("""
+            CREATE TABLE IF NOT EXISTS faces (
+                id SERIAL PRIMARY KEY,
+                path TEXT UNIQUE NOT NULL,
+                name TEXT NOT NULL,
+                embedding vector(512),
+                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
+            )
+        """)
+        
+        # Create index on the embedding column
+        cur.execute("CREATE INDEX IF NOT EXISTS faces_embedding_idx ON faces USING ivfflat (embedding vector_ip_ops)")
+        
+        print("✅ Database setup complete.")
+        return conn
+    except Exception as e:
+        print(f"❌ Database setup failed: {e}")
+        return None
+
+# ─────────────────────────────────────────────
+# 🧠 STEP 3: LOAD CLIP MODEL
+# ─────────────────────────────────────────────
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model, preprocess = clip.load("ViT-B/32", device=device)
+print(f"✅ CLIP model loaded on {device}")
+
+# ─────────────────────────────────────────────
+# 📊 STEP 4: EMBEDDING FUNCTIONS
+# ─────────────────────────────────────────────
+def embed_image(image_path):
+    """Generate CLIP embedding for a single image"""
+    try:
+        img = Image.open(image_path).convert("RGB")
+        img_input = preprocess(img).unsqueeze(0).to(device)
+        with torch.no_grad():
+            emb = model.encode_image(img_input).cpu().numpy().flatten()
+            emb /= np.linalg.norm(emb)
+        return emb
+    except Exception as e:
+        print(f"⚠️ Error embedding {image_path}: {e}")
+        return None
+
+def populate_database(conn, limit=500):
+    """Populate database with images and their embeddings"""
+    # Collect all .jpg files inside subfolders
+    all_images = sorted(glob(os.path.join(img_root, "*", "*.jpg")))
+    selected_images = all_images[:limit]
+    
+    if len(selected_images) == 0:
+        raise RuntimeError("❌ No image files found in unzipped structure!")
+    
+    cur = conn.cursor()
+    
+    # Check which images are already in the database
+    cur.execute("SELECT path FROM faces")
+    existing_paths = set(path[0] for path in cur.fetchall())
+    
+    # Filter out images that are already in the database
+    new_images = [path for path in selected_images if path not in existing_paths]
+    
+    if not new_images:
+        print("✅ All images are already in the database.")
+        return
+    
+    print(f"🧠 Generating CLIP embeddings for {len(new_images)} new images...")
+    
+    # Process images in batches to avoid memory issues
+    batch_size = 50
+    for i in range(0, len(new_images), batch_size):
+        batch = new_images[i:i+batch_size]
+        data_to_insert = []
+        
+        for fpath in tqdm(batch, desc=f"Embedding batch {i//batch_size + 1}"):
+            try:
+                emb = embed_image(fpath)
+                if emb is not None:
+                    name = os.path.splitext(os.path.basename(fpath))[0].replace("_", " ")
+                    data_to_insert.append((fpath, name, emb.tolist()))
+            except Exception as e:
+                print(f"⚠️ Error with {fpath}: {e}")
+        
+        # Insert batch into database
+        if data_to_insert:
+            execute_values(
+                cur,
+                "INSERT INTO faces (path, name, embedding) VALUES %s ON CONFLICT (path) DO NOTHING",
+                [(d[0], d[1], d[2]) for d in data_to_insert],
+                template="(%s, %s, %s::vector)"
+            )
+            conn.commit()
+    
+    # Count total faces in database
+    cur.execute("SELECT COUNT(*) FROM faces")
+    total_faces = cur.fetchone()[0]
+    print(f"✅ Database now contains {total_faces} faces.")
+
+# ─────────────────────────────────────────────
+# 🔐 STEP 5: LOAD OPENAI API KEY
+# ─────────────────────────────────────────────
+openai.api_key = os.getenv("OPENAI_API_KEY")
+
+# ─────────────────────────────────────────────
+# 🔍 STEP 6: FACE MATCHING FUNCTION
+# ─────────────────────────────────────────────
+def scan_face(user_image, conn):
+    """Scan a face image and find matches in the database"""
+    if user_image is None:
+        return [], "", "", "Please upload a face image."
+
+    try:
+        user_image = user_image.convert("RGB")
+        tensor = preprocess(user_image).unsqueeze(0).to(device)
+        with torch.no_grad():
+            query_emb = model.encode_image(tensor).cpu().numpy().flatten()
+            query_emb /= np.linalg.norm(query_emb)
+    except Exception as e:
+        return [], "", "", f"Image preprocessing failed: {e}"
+
+    # Query database for similar faces
+    cur = conn.cursor()
+    emb_list = query_emb.tolist()
+    cur.execute("""
+        SELECT path, name, embedding <-> %s::vector AS distance
+        FROM faces
+        ORDER BY distance
+        LIMIT 5
+    """, (emb_list,))
+    
+    results = cur.fetchall()
+    
+    gallery, captions, names = [], [], []
+    scores = []
+    
+    for path, name, distance in results:
+        try:
+            # Convert distance to similarity score (1 - distance)
+            similarity = 1 - distance
+            scores.append(similarity)
+            
+            img = Image.open(path)
+            gallery.append(img)
+            captions.append(f"{name} (Score: {similarity:.2f})")
+            names.append(name)
+        except Exception as e:
+            captions.append(f"⚠️ Error loading match image: {e}")
+    
+    risk_score = min(100, int(np.mean(scores) * 100)) if scores else 0
+
+    # 🧠 GPT-4 EXPLANATION
+    try:
+        prompt = (
+            f"The uploaded face matches closely with: {', '.join(names)}. "
+            f"Based on this, should the user be suspicious? Analyze like a funny but smart AI dating detective."
+        )
+        response = openai.chat.completions.create(
+            model="gpt-4",
+            messages=[
+                {"role": "system", "content": "You're a playful but intelligent AI face-matching analyst."},
+                {"role": "user", "content": prompt}
+            ]
+        )
+        explanation = response.choices[0].message.content
+    except Exception as e:
+        explanation = f"(OpenAI error): {e}"
+
+    return gallery, "\n".join(captions), f"{risk_score}/100", explanation
+
+# ─────────────────────────────────────────────
+# 🌱 STEP 7: ADD NEW FACE FUNCTION
+# ─────────────────────────────────────────────
+def add_new_face(image, name, conn):
+    """Add a new face to the database"""
+    if image is None or not name:
+        return "Please provide both an image and a name."
+    
+    try:
+        # Save image to a temporary file
+        timestamp = int(time.time())
+        os.makedirs("uploaded_faces", exist_ok=True)
+        path = f"uploaded_faces/{name.replace(' ', '_')}_{timestamp}.jpg"
+        image.save(path)
+        
+        # Generate embedding
+        emb = embed_image(path)
+        if emb is None:
+            return "Failed to generate embedding for the image."
+        
+        # Add to database
+        cur = conn.cursor()
+        cur.execute(
+            "INSERT INTO faces (path, name, embedding) VALUES (%s, %s, %s::vector)",
+            (path, name, emb.tolist())
+        )
+        conn.commit()
+        
+        return f"✅ Added {name} to the database successfully!"
+    except Exception as e:
+        return f"❌ Failed to add face: {e}"
+
+# ─────────────────────────────────────────────
+# 🎛️ STEP 8: GRADIO UI
+# ─────────────────────────────────────────────
+def create_ui():
+    """Create Gradio UI with both scan and add functionality"""
+    # Setup database connection
+    conn = setup_database()
+    if conn is None:
+        raise RuntimeError("❌ Database connection failed. Please check your PostgreSQL installation and pgvector extension.")
+    
+    # Populate database with initial images
+    populate_database(conn)
+    
+    # Wrapper functions for Gradio that use the database connection
+    def scan_face_wrapper(image):
+        return scan_face(image, conn)
+    
+    def add_face_wrapper(image, name):
+        return add_new_face(image, name, conn)
+    
+    with gr.Blocks(title="Tinder Scanner – Real Face Match Detector") as demo:
+        gr.Markdown("# Tinder Scanner – Real Face Match Detector")
+        gr.Markdown("Scan a face image to find visual matches using CLIP and PostgreSQL, and get a cheeky GPT-4 analysis.")
+        
+        with gr.Tab("Scan Face"):
+            with gr.Row():
+                with gr.Column():
+                    input_image = gr.Image(type="pil", label="Upload a Face Image")
+                    scan_button = gr.Button("🔍 Scan Face")
+                
+                with gr.Column():
+                    gallery = gr.Gallery(label="🔍 Top Matches", columns=[5], height="auto")
+                    captions = gr.Textbox(label="Match Names + Similarity Scores")
+                    risk_score = gr.Textbox(label="🚨 Cheating Risk Score")
+                    explanation = gr.Textbox(label="🧠 GPT-4 Explanation", lines=5)
+            
+            scan_button.click(
+                fn=scan_face_wrapper,
+                inputs=[input_image],
+                outputs=[gallery, captions, risk_score, explanation]
+            )
+        
+        with gr.Tab("Add New Face"):
+            with gr.Row():
+                with gr.Column():
+                    new_image = gr.Image(type="pil", label="Upload New Face Image")
+                    new_name = gr.Textbox(label="Person's Name")
+                    add_button = gr.Button("➕ Add to Database")
+                
+                with gr.Column():
+                    result = gr.Textbox(label="Result")
+            
+            add_button.click(
+                fn=add_face_wrapper,
+                inputs=[new_image, new_name],
+                outputs=result
+            )
+    
+    return demo
+
+# ─────────────────────────────────────────────
+# 🚀 MAIN EXECUTION
+# ─────────────────────────────────────────────
+if __name__ == "__main__":
+    demo = create_ui()
+    demo.launch()