import os os.environ["INSIGHTFACE_HOME"] = "/tmp/.insightface" import json import tempfile import numpy as np from insightface.app import FaceAnalysis from scipy.spatial.distance import cosine import cv2 # OpenCV for image processing from typing import List, Dict, Any from datetime import datetime, timedelta import requests import base64 from io import BytesIO from PIL import Image from azure.storage.blob import BlobServiceClient, BlobClient, ContainerClient from config import get_config import time from concurrent.futures import ThreadPoolExecutor, as_completed class EndpointHandler: def __init__(self, model_dir=None): # Initialize FaceAnalysis with GPU/CPU fallback support print("\n" + "="*80) print("INITIALIZING FACEANALYSIS") print("="*80) try: # Try GPU first print("Attempting to initialize with GPU (CUDA)...") self.app = FaceAnalysis(root="/tmp/.insightface", providers=['CUDAExecutionProvider', 'CPUExecutionProvider']) self.app.prepare(ctx_id=0) # 0 = GPU print("✅ GPU initialization successful (ctx_id=0)") self.gpu_available = True except RuntimeError as e: # GPU not available, fall back to CPU print(f"⚠️ GPU initialization failed: {str(e)[:100]}...") print("Falling back to CPU (CPUExecutionProvider)...") try: self.app = FaceAnalysis(root="/tmp/.insightface", providers=['CPUExecutionProvider']) self.app.prepare(ctx_id=-1) # -1 = CPU print("✅ CPU initialization successful (ctx_id=-1)") self.gpu_available = False except Exception as cpu_error: print(f"❌ CPU initialization also failed: {cpu_error}") raise print("="*80 + "\n") print("=" * 80) print("InsightFace Providers:") for model in self.app.models: if hasattr(model, 'sess'): print(f" {model.__class__.__name__}: {model.sess.get_providers()}") print("=" * 80) # Get configuration config = get_config() azure_config = config.get_azure_config() storage_config = config.get_storage_config() # Initialize Azure Blob Storage client if azure_config['connection_string']: self.blob_service_client = BlobServiceClient.from_connection_string( azure_config['connection_string'] ) else: # Use account name and key if connection string not available account_url = f"https://{azure_config['account_name']}.blob.core.windows.net" self.blob_service_client = BlobServiceClient( account_url=account_url, credential=azure_config['account_key'] ) self.container_name = storage_config['container_name'] self.prefix = storage_config['prefix'] self.embeddings_folder = storage_config['embeddings_folder'] # Get container client self.container_client = self.blob_service_client.get_container_client(self.container_name) # Initialize caching self.embeddings_cache = None self.cache_timestamp = 0 self.cache_ttl = 86400 # 24 hours in seconds self.image_cache = {} # In-memory image cache (URL -> numpy array) self.image_cache_max_size = 100 # Max 100 images in memory self.thread_pool = ThreadPoolExecutor(max_workers=8) # Pre-warm GPU and compile models self._prewarm_models() def _prewarm_models(self): """Pre-warm GPU and compile ONNX models on startup to eliminate cold-start latency.""" try: print("\n" + "="*80) mode = "GPU" if self.gpu_available else "CPU" print(f"PRE-WARMING MODELS ({mode} MODE)") print("="*80) start = time.time() # Create a small dummy image (100x100 random RGB) dummy_img = np.random.randint(0, 255, (100, 100, 3), dtype=np.uint8) # Run inference on dummy image to trigger compilation _ = self.app.get(dummy_img) elapsed = time.time() - start print(f"✅ Models pre-warmed in {elapsed:.2f}s ({mode})") print("="*80 + "\n") except Exception as e: print(f"Warning: Model pre-warming failed (non-fatal): {e}") print("="*80 + "\n") def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]: try: if "inputs" in data: return self.process_hf_input(data) else: return self.process_json_input(data) except ValueError as e: return {"error": str(e)} except Exception as e: return {"error": str(e)} def process_hf_input(self, hf_data): """Process Hugging Face format input.""" if "inputs" in hf_data: actual_data = hf_data["inputs"] return self.process_json_input(actual_data) else: return {"error": "Invalid Hugging Face JSON structure."} def process_json_input(self, json_data): if "query_images" in json_data and "gender" in json_data: query_images = json_data["query_images"] gender = json_data["gender"] top_n = json_data.get("top_n", 5) similar_images = self.find_similar_images_aggregate(query_images, gender, top_n) return {"similar_images": similar_images} elif "extract_embeddings" in json_data and json_data["extract_embeddings"]: self.extract_and_save_embeddings() return {"status": "Embeddings extraction completed."} else: raise ValueError("Invalid JSON structure.") def load_embeddings_from_azure(self): """Load existing embeddings from Azure Blob Storage with caching.""" current_time = time.time() # Return cached embeddings if still valid if self.embeddings_cache is not None and (current_time - self.cache_timestamp) < self.cache_ttl: cache_age = int(current_time - self.cache_timestamp) print(f"[TIMING] Using cached embeddings (age: {cache_age}s)") return self.embeddings_cache print("\n" + "="*80) print("LOADING EMBEDDINGS FROM AZURE") print("="*80) blob_name = f'profile-media/embeddings/embeddings_db.json' print(f"Account: {self.blob_service_client.account_name}") print(f"Container: {self.container_name}") print(f"Blob Path: {blob_name}") print(f"Full URL: https://{self.blob_service_client.account_name}.blob.core.windows.net/{self.container_name}/{blob_name}") print("="*80 + "\n") try: blob_client = self.container_client.get_blob_client(blob_name) # Download the existing embeddings file if it exists temp_dir = tempfile.gettempdir() temp_file_path = os.path.join(temp_dir, 'embeddings_db.json') download_start = time.time() with open(temp_file_path, 'wb') as download_file: download_stream = blob_client.download_blob() download_file.write(download_stream.readall()) download_time = time.time() - download_start parse_start = time.time() with open(temp_file_path, 'r') as f: self.embeddings_cache = json.load(f) self.cache_timestamp = current_time parse_time = time.time() - parse_start print(f"[TIMING] Loaded {len(self.embeddings_cache)} embeddings: download={download_time:.3f}s, parse={parse_time:.3f}s") return self.embeddings_cache except Exception as e: print(f'Embeddings file not found in Azure, initializing a new one: {e}') self.embeddings_cache = [] self.cache_timestamp = current_time return [] def extract_and_save_embeddings(self): """Extract embeddings from images and save them to Azure Blob Storage.""" embeddings_db = self.load_embeddings_from_azure() now = datetime.utcnow() thirty_days_ago = now - timedelta(days=30) # Process images from both profile-media and ai-images folders folders_to_process = [ 'profile-media/', # profile-media folder (without container name) 'ai-images/men/', # ai-images/men folder (without container name) 'ai-images/women/' # ai-images/women folder (without container name) ] for folder_prefix in folders_to_process: try: print(f"Processing folder: {folder_prefix}") # List all blobs in the container with the current prefix blob_list = self.container_client.list_blobs(name_starts_with=folder_prefix) for blob in blob_list: blob_name = blob.name if blob_name.endswith(('.jpg', '.jpeg', '.png')): image_url = f'https://{self.blob_service_client.account_name}.blob.core.windows.net/{self.container_name}/{blob_name}' existing_entry = next((item for item in embeddings_db if item['image_url'] == image_url), None) if existing_entry: embedding_timestamp = datetime.fromisoformat(existing_entry['timestamp']) if (existing_entry.get('no_face_detected') or embedding_timestamp > thirty_days_ago) and blob.last_modified.replace(tzinfo=None) <= thirty_days_ago: continue print(f"Processing image: {blob_name}") try: # Create a unique temporary file with proper permissions temp_suffix = os.path.splitext(blob_name)[1] or '.jpg' with tempfile.NamedTemporaryFile(suffix=temp_suffix, delete=False) as temp_image_file: temp_file_path = temp_image_file.name # Download blob to temporary file blob_client = self.container_client.get_blob_client(blob_name) with open(temp_file_path, 'wb') as download_file: download_stream = blob_client.download_blob() download_file.write(download_stream.readall()) img = self.load_image_from_blob(blob_client) # Clean up temporary file immediately after reading try: os.unlink(temp_file_path) except: pass # Ignore cleanup errors if img is None: print(f"Failed to read image: {blob_name}") continue faces = self.app.get(img) if len(faces) == 0: print(f"No face detected in: {blob_name}") no_face_entry = { 'image_url': image_url, 'no_face_detected': True, 'timestamp': now.isoformat() } if existing_entry: existing_entry.update(no_face_entry) else: embeddings_db.append(no_face_entry) continue face = faces[0] embedding = face.embedding.tolist() gender = 'male' if face.gender == 1 else 'female' new_entry = { 'embedding': embedding, 'gender': gender, 'image_url': image_url, 'timestamp': now.isoformat() } if existing_entry: existing_entry.update(new_entry) else: embeddings_db.append(new_entry) print(f"Successfully processed: {blob_name} (gender: {gender})") except Exception as e: print(f"Error processing image {blob_name}: {e}") continue except Exception as e: print(f"Error processing folder {folder_prefix}: {e}") continue print(f"Total embeddings in database: {len(embeddings_db)}") # Save embeddings back to Azure try: temp_json_path = os.path.join(tempfile.gettempdir(), f'embeddings_db_{int(time.time())}.json') with open(temp_json_path, 'w') as temp_json_file: json.dump(embeddings_db, temp_json_file) # Upload to Azure Blob Storage - save in profile-media/embeddings/ blob_name = f'profile-media/embeddings/embeddings_db.json' blob_client = self.container_client.get_blob_client(blob_name) with open(temp_json_path, 'rb') as data: blob_client.upload_blob(data, overwrite=True) print(f"Embeddings saved to Azure: {blob_name}") # Clean up temporary file try: os.unlink(temp_json_path) except: pass # Ignore cleanup errors except Exception as e: print(f"Error saving embeddings: {e}") def find_similar_images_aggregate(self, query_images: List[str], gender: str, top_n: int = 5) -> List[str]: start_time = time.time() print(f"Debug: Starting similarity search with {len(query_images)} query images") print(f"Debug: Looking for gender: {gender}, top_n: {top_n}") # Load embeddings database once (cached) embeddings_db = self.load_embeddings_from_azure() print(f"Debug: Total embeddings in database: {len(embeddings_db)}") # Filter to only include images from profile-media folder structure profile_media_db = [item for item in embeddings_db if 'image_url' in item and 'profile-media' in item['image_url']] print(f"Debug: Profile-media embeddings: {len(profile_media_db)}") # Filter by gender: if 'all', include all items with gender field; otherwise filter by specific gender if gender == 'all': filtered_db = [item for item in profile_media_db if 'gender' in item and 'embedding' in item] else: filtered_db = [item for item in profile_media_db if 'gender' in item and item['gender'] == gender and 'embedding' in item] print(f"Debug: Filtered by gender '{gender}': {len(filtered_db)}") if len(filtered_db) == 0: print(f"Debug: No embeddings found for gender '{gender}' in profile-media folder") return [] # Process query images in parallel similarities = {} futures = {} for i, image_input in enumerate(query_images): future = self.thread_pool.submit(self._extract_query_embedding, image_input, i) futures[future] = i # Collect results from parallel processing query_embeddings = [] for future in as_completed(futures): i = futures[future] try: query_embedding = future.result() if query_embedding is not None: query_embeddings.append(query_embedding) print(f"Debug: Successfully extracted face embedding from query image {i+1}") except Exception as e: print(f"Debug: Error processing query image {i+1}: {e}") if not query_embeddings: print("Debug: No valid query embeddings extracted") return [] # Compute similarities for all query embeddings against filtered database similarity_start = time.time() for query_embedding in query_embeddings: for item in filtered_db: similarity = 1 - cosine(query_embedding, np.array(item['embedding'])) if item['image_url'] in similarities: similarities[item['image_url']].append(similarity) else: similarities[item['image_url']] = [similarity] similarity_time = time.time() - similarity_start # Aggregate similarities print(f"[TIMING] Similarity computation: {similarity_time:.3f}s") print(f"Debug: Total similarities found: {len(similarities)}") aggregated_similarities = [(np.mean(scores), url) for url, scores in similarities.items()] aggregated_similarities.sort(reverse=True, key=lambda x: x[0]) result = [url for _, url in aggregated_similarities[:top_n]] elapsed = time.time() - start_time print(f"\n[TIMING] REQUEST SUMMARY:") print(f" Total time: {elapsed:.3f}s") print(f" Query embeddings extracted: {len(query_embeddings)} images") print(f" Similarity computations: {len(similarities)} results") print(f"Debug: Returning {len(result)} recommendations\n") return result def _extract_query_embedding(self, image_input: str, index: int) -> Any: """Extract embedding from a single query image (for parallel processing).""" try: print(f"\nDebug: Processing query image {index+1}: {image_input}") request_start = time.time() # Load image load_start = time.time() if image_input.startswith('http'): # It's a URL img = self.load_image_from_url(image_input) elif image_input.startswith('data:image/'): # It's a base64-encoded image img = self.load_image_from_base64(image_input) elif image_input.startswith('ai-images/'): # Local filesystem (ai-images baked into container) img = self.load_image_from_local(image_input) else: # It's a local file path reference - convert to full Azure blob URL blob_url = f"https://koottuprod.blob.core.windows.net/koottu-media/{image_input}" img = self.load_image_from_url(blob_url) load_time = time.time() - load_start if img is None: print(f"Failed to load image: {image_input}") return None # Face detection detect_start = time.time() faces = self.app.get(img) detect_time = time.time() - detect_start if len(faces) == 0: elapsed = time.time() - request_start print(f" [TIMING] No faces detected: load={load_time:.3f}s, detect={detect_time:.3f}s, total={elapsed:.3f}s") return None embedding_extraction_time = time.time() - request_start print(f" [TIMING] Face found: load={load_time:.3f}s, detect={detect_time:.3f}s, total={embedding_extraction_time:.3f}s") return faces[0].embedding except Exception as e: print(f"Error extracting embedding from image {index+1}: {e}") return None def find_similar_images_by_embedding(self, query_embedding: np.ndarray, gender: str = 'all', top_n: int = 10, excluded_images: List[str] = None) -> List[str]: """Find similar images based on a given embedding vector.""" try: # Load embeddings database from Azure embeddings_db = self.load_embeddings_from_azure() # Filter to only include images from profile-media folder structure profile_media_db = [item for item in embeddings_db if 'image_url' in item and 'profile-media' in item['image_url']] # Filter by gender if specified if gender != 'all': filtered_db = [item for item in profile_media_db if 'gender' in item and item['gender'] == gender] else: filtered_db = [item for item in profile_media_db if 'embedding' in item] # Filter out excluded images if excluded_images is not None: filtered_db = [item for item in filtered_db if item['image_url'] not in excluded_images] similarities = [] for item in filtered_db: if 'embedding' in item and not item.get('no_face_detected', False): similarity = 1 - cosine(query_embedding, np.array(item['embedding'])) similarities.append((similarity, item['image_url'])) # Sort by similarity and return top matches similarities.sort(reverse=True, key=lambda x: x[0]) return [url for _, url in similarities[:top_n]] except Exception as e: print(f"Error in find_similar_images_by_embedding: {e}") return [] def load_image_from_url(self, url): try: # Check cache first if url in self.image_cache: print(f" [CACHE HIT] {url}") return self.image_cache[url] start = time.time() response = requests.get(url, timeout=30) download_time = time.time() - start response.raise_for_status() start = time.time() image = Image.open(BytesIO(response.content)).convert('RGB') image = np.array(image) parse_time = time.time() - start start = time.time() result = cv2.cvtColor(image, cv2.COLOR_RGB2BGR) convert_time = time.time() - start # Cache the result if len(self.image_cache) >= self.image_cache_max_size: # Remove oldest entry (simple FIFO) self.image_cache.pop(next(iter(self.image_cache))) self.image_cache[url] = result print(f" [TIMING] Image load: download={download_time:.3f}s, parse={parse_time:.3f}s, convert={convert_time:.3f}s [CACHED]") return result except Exception as e: print(f"Error loading image from URL {url}: {e}") return None def load_image_from_blob(self, blob_client): try: blob_bytes = blob_client.download_blob().readall() image = Image.open(BytesIO(blob_bytes)).convert('RGB') image = np.array(image) return cv2.cvtColor(image, cv2.COLOR_RGB2BGR) except Exception as e: print(f"Error loading image from blob: {e}") return None def load_image_from_local(self, local_path): """Load image from local filesystem (for ai-images baked into container).""" try: start = time.time() full_path = os.path.join('/app', local_path) if not os.path.exists(full_path): print(f"Local image not found: {full_path}") return None image = cv2.imread(full_path) if image is None: print(f"Failed to read image: {full_path}") return None load_time = time.time() - start print(f" [TIMING] Image load (local): {load_time:.3f}s [LOCAL FILESYSTEM]") return image except Exception as e: print(f"Error loading local image {local_path}: {e}") return None def load_image_from_base64(self, base64_string): header, encoded = base64_string.split(',', 1) data = base64.b64decode(encoded) np_arr = np.frombuffer(data, np.uint8) img = cv2.imdecode(np_arr, cv2.IMREAD_COLOR) return img # Returns BGR image as expected by OpenCV # Instantiate the handler handler = EndpointHandler()