Spaces:
Runtime error
Runtime error
| from fastapi import FastAPI, HTTPException | |
| from pydantic import BaseModel | |
| import os | |
| import shutil | |
| import numpy as np | |
| from PIL import Image | |
| import face_recognition | |
| from scipy.spatial import distance | |
| app = FastAPI() | |
| 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 | |
| # 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 | |
| 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(input_dir, output_dir): | |
| filenames = os.listdir(input_dir) | |
| images_to_group = [ImageToGroup(filename, os.path.join(input_dir, filename)) for filename in filenames] | |
| # 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) | |
| class GroupImagesRequest(BaseModel): | |
| input_dir: str | |
| output_dir: str | |
| async def group_images(request: GroupImagesRequest): | |
| if not os.path.isdir(request.input_dir): | |
| raise HTTPException(status_code=400, detail="Input directory does not exist") | |
| if not os.path.isdir(request.output_dir): | |
| os.makedirs(request.output_dir) | |
| main(request.input_dir, request.output_dir) | |
| return {"message": "Images grouped successfully"} |