Face_embedding / app.py
projectimx's picture
Update app.py
7883d65 verified
import numpy as np
import tensorflow as tf
from flask import Flask, request, jsonify
from PIL import Image
import io
import base64
# Initialize Flask app
app = Flask(__name__)
# Load TensorFlow Lite model
interpreter = tf.lite.Interpreter(model_path="facenet.tflite")
interpreter.allocate_tensors()
# Get input and output details
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()
def preprocess_image(image_data):
"""
Preprocess the input image for the FaceNet model.
"""
image = Image.open(io.BytesIO(image_data)).convert('RGB')
image = image.resize((160, 160)) # Resize to the model's input size
image_array = np.asarray(image).astype(np.float32)
image_array = (image_array - 127.5) / 127.5 # Normalize to [-1, 1]
image_array = np.expand_dims(image_array, axis=0) # Add batch dimension
return image_array
def create_face_embedding(image_data):
"""
Generate a face embedding for the given image data.
"""
processed_image = preprocess_image(image_data)
# Run the model
interpreter.set_tensor(input_details[0]['index'], processed_image)
interpreter.invoke()
# Extract the embedding
embedding = interpreter.get_tensor(output_details[0]['index'])
return embedding.flatten().tolist()
@app.route('/generate-embedding', methods=['POST'])
def generate_embedding():
"""
Endpoint to process an image and return its face embedding.
"""
try:
# Parse incoming JSON with base64-encoded image
data = request.json
if "image" not in data:
return jsonify({"error": "Image data not provided"}), 400
# Decode base64 image
image_data = base64.b64decode(data["image"])
# Generate embedding
embedding = create_face_embedding(image_data)
return jsonify({"embedding": embedding}), 200
except Exception as e:
return jsonify({"error": str(e)}), 500
@app.route('/')
def home():
return "Face Embedding Generator API is running!"
if __name__ == '__main__':
app.run(debug=True, host='0.0.0.0', port=7860)