Update app.py

app.py

@@ -1,334 +1,325 @@
 import os
-import zipfile
+import json
 import torch
 import clip
+import faiss
 import numpy as np
 from PIL import Image
 import gradio as gr
 import openai
+import requests
 from tqdm import tqdm
-from glob import glob
-import chromadb
-from chromadb.utils import embedding_functions
-import json
-import time
-from dotenv import load_dotenv
-
-# Load environment variables from .env file
-load_dotenv()
+from io import BytesIO
 
 # ─────────────────────────────────────────────
-# 📂 STEP 1: UNZIP TO CORRECT STRUCTURE
+# 🧠 STEP 1: LOAD CLIP MODEL
 # ─────────────────────────────────────────────
-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")
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model, preprocess = clip.load("ViT-B/32", device=device)
 
 # ─────────────────────────────────────────────
-# 🧠 STEP 2: LOAD CLIP MODEL
+# 📦 STEP 2: LOAD PROFILE DATA FROM JSON
 # ─────────────────────────────────────────────
-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}")
+def load_profile_data(json_file_path=None, json_data=None):
+    """Load profile data either from a file or directly from JSON data"""
+    if json_file_path and os.path.exists(json_file_path):
+        with open(json_file_path, 'r') as f:
+            profiles = json.load(f)
+    elif json_data:
+        profiles = json_data
+    else:
+        # Sample data structure as fallback
+        profiles = [
+            {
+                "Id": "sample-id",
+                "Name": "Sample Profile",
+                "Age": 25,
+                "Bio": "Sample bio",
+                "Photos": [
+                    "https://example.com/sample.jpg"
+                ]
+            }
+        ]
+    
+    return profiles
 
 # ─────────────────────────────────────────────
-# 🗄️ STEP 3: CHROMA DB SETUP & EMBEDDING FUNCTION
+# 🖼️ STEP 3: DOWNLOAD AND PROCESS IMAGES
 # ─────────────────────────────────────────────
-class ClipEmbeddingFunction:
-    """Custom embedding function for Chroma DB using CLIP"""
+def download_and_process_image(url):
+    """Download image from URL and return PIL Image"""
+    try:
+        response = requests.get(url, timeout=10)
+        response.raise_for_status()
+        img = Image.open(BytesIO(response.content)).convert("RGB")
+        return img
+    except Exception as e:
+        print(f"⚠️ Error downloading image from {url}: {e}")
+        return None
+
+def generate_embeddings(profiles, max_images=500):
+    """Generate CLIP embeddings for profile images"""
+    embeddings = []
+    image_urls = []
+    profile_info = []  # Store name, age, etc. for each image
     
-    def __init__(self, model, preprocess, device):
-        self.model = model
-        self.preprocess = preprocess
-        self.device = device
+    image_count = 0
     
-    def __call__(self, images):
-        """Generate embeddings for a list of image paths"""
-        embeddings = []
+    print(f"🧠 Generating CLIP embeddings for profile images...")
+    for profile in tqdm(profiles, desc="Processing profiles"):
+        name = profile.get("Name", "Unknown")
+        age = profile.get("Age", "?")
         
-        for image_path in images:
-            try:
-                # Check if the path is a string (for new additions from disk)
-                if isinstance(image_path, str) and os.path.exists(image_path):
-                    img = Image.open(image_path).convert("RGB")
-                else:
-                    # For query images that are already PIL images
-                    img = image_path.convert("RGB") if hasattr(image_path, 'convert') else image_path
+        for photo_url in profile.get("Photos", []):
+            if image_count >= max_images:
+                break
                 
-                img_input = self.preprocess(img).unsqueeze(0).to(self.device)
+            try:
+                img = download_and_process_image(photo_url)
+                if img is None:
+                    continue
+                    
+                img_input = preprocess(img).unsqueeze(0).to(device)
                 with torch.no_grad():
-                    emb = self.model.encode_image(img_input).cpu().numpy().flatten()
+                    emb = model.encode_image(img_input).cpu().numpy().flatten()
                     emb /= np.linalg.norm(emb)
-                embeddings.append(emb.tolist())
+                
+                embeddings.append(emb)
+                image_urls.append(photo_url)
+                profile_info.append({
+                    "Name": name,
+                    "Age": age,
+                    "Id": profile.get("Id", "Unknown"),
+                    "Bio": profile.get("Bio", "")
+                })
+                
+                image_count += 1
             except Exception as e:
-                print(f"⚠️ Error embedding image: {e}")
-                # Return a zero vector as fallback
-                embeddings.append([0] * 512)
+                print(f"⚠️ Error with {photo_url}: {e}")
         
-        return embeddings
+        if image_count >= max_images:
+            break
+    
+    if embeddings:
+        embeddings = np.vstack(embeddings).astype("float32")
+    else:
+        embeddings = np.array([]).astype("float32")
+    
+    print(f"✅ Finished embedding {len(embeddings)} images.")
+    return embeddings, image_urls, profile_info
 
-def setup_database():
-    """Setup ChromaDB with CLIP embedding function"""
-    try:
-        # Create persistent client
-        client = chromadb.PersistentClient(path="./chroma_db")
-        
-        # Create custom embedding function
-        embedding_function = ClipEmbeddingFunction(model, preprocess, device)
-        
-        # Create or get existing collection
-        collection = client.get_or_create_collection(
-            name="faces",
-            embedding_function=embedding_function,
-            metadata={"hnsw:space": "cosine"}  # Use cosine similarity
-        )
+# ─────────────────────────────────────────────
+# ⚡ STEP 4: BUILD FAISS INDEX
+# ─────────────────────────────────────────────
+def build_faiss_index(embeddings):
+    """Build FAISS index from embeddings"""
+    if len(embeddings) == 0:
+        return None
         
-        print("✅ ChromaDB setup complete.")
-        return client, collection
-    except Exception as e:
-        print(f"❌ Database setup failed: {e}")
-        return None, None
+    dimension = embeddings.shape[1]
+    index = faiss.IndexFlatIP(dimension)
+    index.add(embeddings)
+    return index
 
-def populate_database(collection, limit=500):
-    """Populate ChromaDB 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!")
+# ─────────────────────────────────────────────
+# 🔐 STEP 5: LOAD OPENAI API KEY
+# ─────────────────────────────────────────────
+def init_openai():
+    openai.api_key = os.getenv("OPENAI_API_KEY")
+    if not openai.api_key:
+        print("⚠️ Warning: OPENAI_API_KEY not found. GPT-4 analysis will not be available.")
+
+# ─────────────────────────────────────────────
+# 🔎 STEP 6: SEARCH FUNCTIONALITY
+# ─────────────────────────────────────────────
+def search_similar_faces(user_image, index, image_urls, profile_info, top_k=5):
+    """Search for similar faces using CLIP + FAISS"""
+    if index is None:
+        return [], [], 0, "No index available. Please load profile data first."
     
-    # Get existing IDs
-    existing_ids = set()
     try:
-        existing_count = collection.count()
-        if existing_count > 0:
-            results = collection.get(limit=existing_count)
-            existing_ids = set(results['ids'])
+        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().astype("float32")
+            query_emb /= np.linalg.norm(query_emb)
     except Exception as e:
-        print(f"Error getting existing IDs: {e}")
-    
-    # Filter out images that are already in the database
-    new_images = []
-    new_ids = []
-    new_metadatas = []
+        return [], [], 0, f"Image preprocessing failed: {e}"
     
-    for fpath in selected_images:
-        # Create ID from path
-        image_id = fpath.replace('/', '_')
-        if image_id not in existing_ids:
-            new_images.append(fpath)
-            new_ids.append(image_id)
-            name = os.path.splitext(os.path.basename(fpath))[0].replace("_", " ")
-            new_metadatas.append({
-                "path": fpath,
-                "name": name
-            })
+    scores, indices = index.search(query_emb, top_k)
+    scores, indices = scores.flatten(), indices.flatten()
     
-    if not new_images:
-        print("✅ All images are already in the database.")
-        return
+    matching_images = []
+    match_details = []
     
-    print(f"🧠 Adding {len(new_images)} new images to the database...")
-    
-    # Process images in batches to avoid memory issues
-    batch_size = 50
-    for i in range(0, len(new_images), batch_size):
-        batch_imgs = new_images[i:i+batch_size]
-        batch_ids = new_ids[i:i+batch_size]
-        batch_metadatas = new_metadatas[i:i+batch_size]
-        
-        print(f"Processing batch {i//batch_size + 1}/{(len(new_images)-1)//batch_size + 1}...")
+    for i in range(len(indices)):
+        idx = indices[i]
+        score = scores[i]
         
         try:
-            collection.add(
-                documents=batch_imgs,  # ChromaDB will call our embedding function on these
-                ids=batch_ids,
-                metadatas=batch_metadatas
-            )
+            url = image_urls[idx]
+            info = profile_info[idx]
+            
+            img = download_and_process_image(url)
+            if img:
+                matching_images.append(img)
+                match_details.append({
+                    "url": url,
+                    "score": score,
+                    "info": info
+                })
         except Exception as e:
-            print(f"⚠️ Error adding batch to database: {e}")
+            print(f"⚠️ Error processing match at index {idx}: {e}")
     
-    # Count total faces in database
-    total_faces = collection.count()
-    print(f"✅ Database now contains {total_faces} faces.")
-
-# ─────────────────────────────────────────────
-# 🔐 STEP 4: LOAD OPENAI API KEY
-# ─────────────────────────────────────────────
-openai.api_key = os.getenv("OPENAI_API_KEY")
-if not openai.api_key:
-    print("⚠️ OpenAI API key not found. GPT-4 analysis will not work.")
+    risk_score = min(100, int(np.mean(scores) * 100)) if scores.size > 0 else 0
+    
+    return matching_images, match_details, risk_score
 
 # ─────────────────────────────────────────────
-# 🔍 STEP 5: FACE MATCHING FUNCTION
+# 🧠 STEP 7: GPT-4 ANALYSIS
 # ─────────────────────────────────────────────
-def scan_face(user_image, collection):
-    """Scan a face image and find matches in the database"""
-    if user_image is None:
-        return [], "", "", "Please upload a face image."
-
+def generate_gpt4_analysis(match_details):
+    """Generate fun analysis using GPT-4"""
+    if not openai.api_key:
+        return "GPT-4 analysis not available (API key not configured)"
+    
+    if not match_details:
+        return "No matches found for analysis"
+    
     try:
-        # Query database for similar faces using the image directly
-        results = collection.query(
-            query_embeddings=None,  # Will be generated by our embedding function
-            query_images=[user_image],  # Pass the PIL image directly
-            n_results=5,
-            include=["metadatas", "distances"]
-        )
+        names = [f"{d['info']['Name']} ({d['info']['Age']})" for d in match_details]
         
-        metadatas = results.get("metadatas", [[]])[0]
-        distances = results.get("distances", [[]])[0]
-        
-        gallery, captions, names = [], [], []
-        scores = []
+        prompt = (
+            f"The uploaded face matches closely with: {', '.join(names)}. "
+            f"Based on this, should the user be suspicious? "
+            f"Analyze like a funny but smart AI dating detective. Keep it concise."
+        )
         
-        for i, metadata in enumerate(metadatas):
-            try:
-                path = metadata["path"]
-                name = metadata["name"]
-                
-                # Convert distance to similarity score (1 - normalized_distance)
-                # ChromaDB uses cosine distance, so 0 is most similar, 2 is most different
-                distance = distances[i]
-                similarity = 1 - (distance / 2)  # Convert to 0-1 scale
-                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}")
+        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}
+            ]
+        )
         
-        risk_score = min(100, int(np.mean(scores) * 100)) if scores else 0
-
-        # 🧠 GPT-4 EXPLANATION
-        explanation = ""
-        if openai.api_key and names:
-            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}"
-        else:
-            explanation = "OpenAI API key not set or no matches found."
-
-        return gallery, "\n".join(captions), f"{risk_score}/100", explanation
-    
+        return response.choices[0].message.content
     except Exception as e:
-        return [], "", "", f"Error scanning face: {e}"
+        return f"(OpenAI error): {e}"
 
 # ─────────────────────────────────────────────
-# 🌱 STEP 6: ADD NEW FACE FUNCTION
+# 🎛️ STEP 8: APPLICATION CLASS
 # ─────────────────────────────────────────────
-def add_new_face(image, name, collection):
-    """Add a new face to the database"""
-    if image is None or not name:
-        return "Please provide both an image and a name."
+class TinderScanner:
+    def __init__(self):
+        self.index = None
+        self.image_urls = []
+        self.profile_info = []
+        self.profiles = []
+        
+        # Initialize OpenAI
+        init_openai()
     
-    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)
+    def load_data(self, json_text=None, json_file=None):
+        """Load profile data and build index"""
+        try:
+            if json_text:
+                json_data = json.loads(json_text)
+                self.profiles = load_profile_data(json_data=json_data)
+            elif json_file:
+                self.profiles = load_profile_data(json_file_path=json_file)
+            else:
+                return "Please provide either JSON text or a JSON file"
+            
+            embeddings, self.image_urls, self.profile_info = generate_embeddings(self.profiles)
+            
+            if len(embeddings) > 0:
+                self.index = build_faiss_index(embeddings)
+                return f"✅ Successfully loaded {len(self.profiles)} profiles with {len(self.image_urls)} photos"
+            else:
+                return "⚠️ No valid images found in the provided data"
+        except Exception as e:
+            return f"❌ Error loading data: {e}"
+    
+    def scan_face(self, user_image, json_input=None):
+        """Process a user image and find matches"""
+        # Load data if provided and not already loaded
+        if json_input and not self.index:
+            load_result = self.load_data(json_text=json_input)
+            if "Successfully" not in load_result:
+                return [], "", "", load_result
+        
+        if not self.index:
+            return [], "", "", "Please load profile data first by providing JSON input"
         
-        # Add to ChromaDB
-        image_id = path.replace('/', '_')
-        collection.add(
-            documents=[path],
-            ids=[image_id],
-            metadatas=[{
-                "path": path,
-                "name": name
-            }]
+        if user_image is None:
+            return [], "", "", "Please upload a face image"
+        
+        images, match_details, risk_score = search_similar_faces(
+            user_image, self.index, self.image_urls, self.profile_info
         )
         
-        return f"✅ Added {name} to the database successfully!"
-    except Exception as e:
-        return f"❌ Failed to add face: {e}"
+        # Format match captions
+        captions = []
+        for detail in match_details:
+            info = detail["info"]
+            captions.append(f"{info['Name']} ({info['Age']}) - Score: {detail['score']:.2f}")
+        
+        # Generate GPT-4 analysis
+        explanation = generate_gpt4_analysis(match_details)
+        
+        return images, "\n".join(captions), f"{risk_score}/100", explanation
 
 # ─────────────────────────────────────────────
-# 🎛️ STEP 7: GRADIO UI
+# 🖥️ STEP 9: GRADIO UI
 # ─────────────────────────────────────────────
 def create_ui():
-    """Create Gradio UI with both scan and add functionality"""
-    # Setup database
-    client, collection = setup_database()
-    if collection is None:
-        raise RuntimeError("❌ Database setup failed.")
-    
-    # Populate database with initial images
-    populate_database(collection)
-    
-    # Wrapper functions for Gradio that use the database collection
-    def scan_face_wrapper(image):
-        return scan_face(image, collection)
-    
-    def add_face_wrapper(image, name):
-        return add_new_face(image, name, collection)
+    scanner = TinderScanner()
     
-    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 ChromaDB, and get a cheeky GPT-4 analysis.")
+    with gr.Blocks(title="Enhanced Tinder Scanner") as demo:
+        gr.Markdown("# 🔍 Tinder Scanner Pro – Face Match Detector")
+        gr.Markdown("Scan a face image to find visual matches in Tinder profiles 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.Tabs():
+            with gr.TabItem("Setup Data"):
+                with gr.Row():
+                    with gr.Column():
+                        json_input = gr.Textbox(
+                            label="JSON Profile Data",
+                            placeholder='Paste JSON data here. Format: [{"Id": "...", "Name": "...", "Age": 25, "Photos": ["url1", "url2"]}]',
+                            lines=10
+                        )
+                        load_btn = gr.Button("Load Profile Data", variant="primary")
+                        data_status = gr.Textbox(label="Status")
                 
-                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)
+                load_btn.click(
+                    fn=scanner.load_data,
+                    inputs=[json_input],
+                    outputs=[data_status]
+                )
             
-            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.TabItem("Scan Face"):
+                with gr.Row():
+                    with gr.Column():
+                        user_image = gr.Image(type="pil", label="Upload a Face Image")
+                        scan_btn = gr.Button("Scan Face", variant="primary")
+                    
+                    with gr.Column():
+                        matches_gallery = gr.Gallery(label="🔍 Top Matches", columns=[5], height="auto")
+                        match_details = gr.Textbox(label="Match Details")
+                        risk_score = gr.Textbox(label="🚨 Similarity Score")
+                        gpt_analysis = gr.Textbox(label="🧠 GPT-4 Analysis")
                 
-                with gr.Column():
-                    result = gr.Textbox(label="Result")
-            
-            add_button.click(
-                fn=add_face_wrapper,
-                inputs=[new_image, new_name],
-                outputs=result
-            )
+                scan_btn.click(
+                    fn=scanner.scan_face,
+                    inputs=[user_image, json_input],
+                    outputs=[matches_gallery, match_details, risk_score, gpt_analysis]
+                )
     
     return demo
 
 # ─────────────────────────────────────────────
-# 🚀 MAIN EXECUTION
+# 🚀 STEP 10: MAIN EXECUTION
 # ─────────────────────────────────────────────
 if __name__ == "__main__":
     demo = create_ui()