Spaces:
Running
Running
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 | |
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}) | |
def hello_world(): | |
return 'Welcome to the Face Recognition API! Upload an image for recognition.' | |
# Run Flask app | |
if __name__ == '__main__': | |
app.run() | |