Spaces:
Sleeping
Sleeping
| from agentlego.tools import BaseTool | |
| import os | |
| import cv2 | |
| import numpy as np | |
| from insightface.app import FaceAnalysis | |
| class FaceSimilarityTool(BaseTool): | |
| default_desc = 'Uses InsightFace to evaluate face similarity between two images.' | |
| def __init__(self, threshold=0.01, use_gpu=False): | |
| super().__init__() | |
| # Initialize InsightFace | |
| self.threshold = threshold | |
| # Use GPU if available (ctx_id=0) or CPU (ctx_id=-1) | |
| ctx_id = 0 if use_gpu else -1 | |
| # Initialize FaceAnalysis | |
| print("Initializing InsightFace model...") | |
| self.app = FaceAnalysis(name='buffalo_s') | |
| self.app.prepare(ctx_id=ctx_id, det_size=(640, 640)) | |
| print("✅ InsightFace model initialized successfully") | |
| def apply(self, img1_path: str, img2_path: str) -> str: | |
| try: | |
| # Validate file existence | |
| if not os.path.exists(img1_path): | |
| return f"Error: Image 1 not found at path: {img1_path}" | |
| if not os.path.exists(img2_path): | |
| return f"Error: Image 2 not found at path: {img2_path}" | |
| # Load images | |
| img1 = cv2.imread(img1_path) | |
| img2 = cv2.imread(img2_path) | |
| if img1 is None or img2 is None: | |
| return f"Error: Failed to load one or both images" | |
| # Get faces and embeddings | |
| faces1 = self.app.get(img1) | |
| faces2 = self.app.get(img2) | |
| # Check for no faces detected | |
| if len(faces1) == 0: | |
| return "Face similarity result: **No match** (No faces detected in first image)" | |
| if len(faces2) == 0: | |
| return "Face similarity result: **No match** (No faces detected in second image)" | |
| # Get first face from each image (use primary face if multiple detected) | |
| embedding1 = faces1[0].embedding | |
| embedding2 = faces2[0].embedding | |
| # Compute cosine similarity | |
| similarity = np.dot(embedding1, embedding2) / (np.linalg.norm(embedding1) * np.linalg.norm(embedding2)) | |
| # Interpret result | |
| if similarity > self.threshold: | |
| result = "Match detected" | |
| match_text = "True" | |
| else: | |
| result = "No match" | |
| match_text = "False" | |
| # Format in a way similar to your original implementation | |
| return f"Face similarity result: **{match_text}** ({result}, similarity score: {similarity:.4f})" | |
| except Exception as e: | |
| return f"Error during face similarity computation: {str(e)}" | |