facenet-pytorch

import torch from facenet_pytorch import MTCNN, InceptionResnetV1 from PIL import Image import os # Set device for GPU if available, else CPU device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') print(f'Running on device: {device}') # Initialize MTCNN for face detection mtcnn = MTCNN( image_size=160, margin=0, min_face_size=20, thresholds=[0.6, 0.7, 0.7], factor=0.709, post_process=True, device=device ) # Initialize InceptionResnetV1 for face recognition resnet = InceptionResnetV1(pretrained='vggface2').eval().to(device) # Create a dummy image for demonstration (replace with your image path) # In a real scenario, load an image from disk or URL # Example: img = Image.open('path/to/your/image.jpg').convert('RGB') # For a runnable example, we create a blank image try: # Attempt to load a real image for better demo, if it exists dummy_image_path = os.path.join(os.path.dirname(__file__), 'dummy_face.jpg') if os.path.exists(dummy_image_path): img = Image.open(dummy_image_path).convert('RGB') else: # Create a blank image if no dummy_face.jpg is found img = Image.new('RGB', (250, 250), color = 'red') print("No 'dummy_face.jpg' found. Using a blank red image. Face detection will likely fail.") except Exception as e: img = Image.new('RGB', (250, 250), color = 'red') print(f"Could not load image, creating a blank red image. Error: {e}") # Detect faces img_cropped = mtcnn(img) if img_cropped is not None: # Calculate face embedding img_embedding = resnet(img_cropped.unsqueeze(0)).detach().cpu() print("Face detected and embedding calculated.") print(f"Embedding shape: {img_embedding.shape}") else: print("No face detected.")