depthmapflask / face_recog.py
srivatsavdamaraju's picture
Create face_recog.py
a82c6ba verified
raw
history blame
3.26 kB
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()