Spaces:
Runtime error
Runtime error
| import os | |
| import numpy as np | |
| from PIL import Image | |
| import face_recognition | |
| from scipy.spatial import distance | |
| import gradio as gr | |
| class ImageToGroup(object): | |
| def __init__(self, filename, path): | |
| self.filename = filename | |
| self.path = path | |
| self.embeddings = self.extract_embeddings() | |
| def extract_embeddings(self): | |
| try: | |
| img = Image.open(self.path) | |
| img = img.convert("RGB") # Convert image to RGB | |
| img = img.resize((800, 800)) # Resize the image to improve performance | |
| img = np.array(img) | |
| face_locations = face_recognition.face_locations(img) | |
| if len(face_locations) == 0: | |
| return [] # No face found in the image | |
| # Generate multiple face encodings with jitter | |
| face_encodings = [face_recognition.face_encodings(img, [face_location], num_jitters=10)[0] for face_location in face_locations] | |
| return face_encodings | |
| except Exception as e: | |
| print(f"Error extracting embeddings from {self.path}: {e}") | |
| return [] | |
| def are_similar(self, other_embeddings, threshold=0.6): | |
| # Calculate the Euclidean distance between two embeddings | |
| for other_embedding in other_embeddings: | |
| for _ in self.embeddings: | |
| dist = distance.euclidean(_, other_embedding) | |
| if dist < threshold: | |
| return True | |
| return False | |
| def group_images(input_files): | |
| images_to_group = [ImageToGroup(os.path.basename(file), file) for file in input_files] | |
| # Group images into clusters based on face embeddings | |
| grouped_images = {} | |
| for image in images_to_group: | |
| if not image.embeddings: | |
| continue # Skip images with no faces | |
| for group_images in grouped_images.values(): | |
| if image.are_similar(group_images[0].embeddings): # Compare embeddings using are_similar method | |
| group_images.append(image) | |
| break | |
| else: | |
| # Use the first embedding of the current image as the key for the new group | |
| grouped_images[tuple(image.embeddings[0])] = [image] | |
| # Convert grouped images to PIL Image objects | |
| group_images_pil = [] | |
| for i, (_, group_images) in enumerate(grouped_images.items()): | |
| for image in group_images: | |
| pil_image = Image.open(image.path) | |
| group_images_pil.append((pil_image, f"Group {i+1}")) | |
| return group_images_pil | |
| # Interface for Gradio | |
| input_directory = gr.File(label="Input Directory",file_count="multiple") | |
| gallery = gr.Gallery(label="Grouped Images") | |
| gr.Interface( | |
| fn=group_images, | |
| inputs=[input_directory], | |
| outputs=gallery | |
| ).launch(share=True) |