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()