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)