from flask import Flask, request, jsonify
import face_recognition
import cv2
import numpy as np
import os
from pathlib import Path
import base64

# Initialize Flask app
app = Flask(__name__)

# Set the maximum file upload size (50MB)
app.config['MAX_CONTENT_LENGTH'] = 50 * 1024 * 1024  # 50MB

# Define the path to dataset and test image
dataset_folder_path = '/content/sample_data/dataset'

# Load and encode faces from the dataset folder
def load_face_encodings(dataset_folder_path):
    known_face_encodings = []
    known_face_names = []

    for img_path in Path(dataset_folder_path).glob("*.jpg"):
        image = face_recognition.load_image_file(img_path)
        encodings = face_recognition.face_encodings(image)

        if encodings:
            known_face_encodings.append(encodings[0])
            known_face_names.append(img_path.stem)

    return known_face_encodings, known_face_names

# Load and process the uploaded test image
def load_test_image(test_image):
    test_image_rgb = cv2.imdecode(np.fromstring(test_image.read(), np.uint8), cv2.IMREAD_COLOR)
    test_image_rgb = cv2.cvtColor(test_image_rgb, cv2.COLOR_BGR2RGB)
    face_locations = face_recognition.face_locations(test_image_rgb)
    face_encodings = face_recognition.face_encodings(test_image_rgb, face_locations)
    return test_image_rgb, face_locations, face_encodings

# Compare faces from the test image with the dataset faces
def compare_faces(known_face_encodings, known_face_names, test_image_encodings):
    matches = []
    for test_encoding in test_image_encodings:
        results = face_recognition.compare_faces(known_face_encodings, test_encoding)
        if True in results:
            first_match_index = results.index(True)
            matches.append(known_face_names[first_match_index])
        else:
            matches.append("Unknown")
    return matches

# Convert image to base64
def convert_image_to_base64(image):
    _, buffer = cv2.imencode('.jpg', image)
    img_bytes = buffer.tobytes()
    img_base64 = base64.b64encode(img_bytes).decode('utf-8')
    return img_base64

# Flask route for face recognition
@app.route('/recognize_faces', methods=['POST'])
def recognize_faces():
    if 'image' not in request.files:
        return jsonify({"error": "No image file provided!"}), 400

    # Load known face encodings
    known_face_encodings, known_face_names = load_face_encodings(dataset_folder_path)
    
    image_file = request.files['image']
    test_image_rgb, face_locations, test_image_encodings = load_test_image(image_file)
    
    matches = compare_faces(known_face_encodings, known_face_names, test_image_encodings)

    for (top, right, bottom, left), name in zip(face_locations, matches):
        cv2.rectangle(test_image_rgb, (left, top), (right, bottom), (0, 0, 255), 2)
        font = cv2.FONT_HERSHEY_DUPLEX
        cv2.putText(test_image_rgb, name, (left + 6, bottom - 6), font, 0.5, (255, 255, 255), 1)

    result_image_base64 = convert_image_to_base64(test_image_rgb)
    return jsonify({'recognized_faces': matches, 'image': result_image_base64})

@app.route('/')
def hello_world():
    return 'Welcome to the Face Recognition API! Upload an image for recognition.'

# Run Flask app
if __name__ == '__main__':
    app.run()