Spaces:
Runtime error
Runtime error
| import streamlit as st | |
| import os | |
| import shutil | |
| import numpy as np | |
| from PIL import Image | |
| import face_recognition | |
| from scipy.spatial import distance | |
| import tempfile | |
| class ImageToGroup(object): | |
| def __init__(self, filename, path): | |
| self.filename = filename | |
| self.path = path | |
| self.embeddings = self.extract_embeddings() | |
| def extract_embeddings(self): | |
| img = face_recognition.load_image_file(self.path) | |
| face_locations = face_recognition.face_locations(img) | |
| if len(face_locations) == 0: | |
| return [] # No face found in the image | |
| face_encodings = [face_recognition.face_encodings(img, [face_location])[0] for face_location in face_locations] | |
| return face_encodings | |
| def are_similar(self, other_embeddings, threshold=0.6): | |
| # Calculate the Euclidean distance between two embeddings | |
| for other_embedding in other_embeddings: | |
| for self_embedding in self.embeddings: | |
| dist = distance.euclidean(self_embedding, other_embedding) | |
| if dist < threshold: | |
| return True | |
| return False | |
| def main(images, output_dir): | |
| images_to_group = [ImageToGroup(os.path.basename(image_path), image_path) for image_path in images] | |
| # 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 embedding in image.embeddings: | |
| found_group = False | |
| for group_key, group_images in grouped_images.items(): | |
| if image.are_similar(group_images[0].embeddings): # Compare embeddings using are_similar method | |
| group_images.append(image) | |
| found_group = True | |
| break | |
| if not found_group: | |
| grouped_images[tuple(embedding)] = [image] # Convert numpy.ndarray to a hashable type using tuple() | |
| # Save grouped images | |
| for i, (embedding, group_images) in enumerate(grouped_images.items()): | |
| group_dir = os.path.join(output_dir, f"group_{i+1}") # Add +1 to the index to start from 1 | |
| os.makedirs(group_dir, exist_ok=True) | |
| for image in group_images: | |
| image_filename = os.path.basename(image.path) | |
| destination_path = os.path.join(group_dir, image_filename) | |
| shutil.copy(image.path, destination_path) | |
| def app(): | |
| st.title('Image Grouping based on Face Recognition') | |
| uploaded_files = st.file_uploader('Upload images:', type=['png', 'jpg', 'jpeg'], accept_multiple_files=True) | |
| if st.button('Group Images'): | |
| if uploaded_files: | |
| images = [] | |
| for uploaded_file in uploaded_files: | |
| tfile = tempfile.NamedTemporaryFile(delete=False) | |
| tfile.write(uploaded_file.read()) | |
| images.append(tfile.name) | |
| output_dir = tempfile.mkdtemp() | |
| main(images, output_dir) | |
| st.success('Images grouped successfully.') | |
| # Display the grouped images | |
| for group_dir in os.listdir(output_dir): | |
| st.header(f'Group: {group_dir}') | |
| for image_file in os.listdir(os.path.join(output_dir, group_dir)): | |
| image = Image.open(os.path.join(output_dir, group_dir, image_file)) | |
| st.image(image) | |
| else: | |
| st.error('Please upload images.') | |
| if __name__ == '__main__': | |
| app() |