Travis / app.py
vinay0123's picture
Update app.py
2cc9f40 verified
Raw
History Blame Contribute Delete
8.71 kB
import base64
import os
import tempfile
# Use temporary directory or current working directory for app data
# This avoids permission issues in containerized environments
# Check if running in container with pre-created directories
if os.path.exists('/tmp/app_data') and os.access('/tmp/app_data', os.W_OK):
app_data_path = '/tmp/app_data'
print("Using pre-created /tmp/app_data directory")
else:
try:
# Try to use /tmp first (usually writable in containers)
app_data_path = '/tmp/app_data'
os.makedirs(app_data_path, exist_ok=True)
except PermissionError:
# Fallback to current directory or temp directory
app_data_path = os.path.join(os.getcwd(), 'app_data')
try:
os.makedirs(app_data_path, exist_ok=True)
except PermissionError:
# Last resort: use system temp directory
app_data_path = tempfile.mkdtemp(prefix='app_data_')
# Make sure the necessary directories exist and have proper permissions
import cv2
import numpy as np
from flask import Flask, request, jsonify
from mtcnn.mtcnn import MTCNN
from keras_facenet import FaceNet
from sklearn.metrics.pairwise import cosine_similarity
from flask_cors import CORS
from pymongo import MongoClient
from pymongo.server_api import ServerApi
# Initialize MongoDB connection
client = MongoClient('mongodb+srv://nanduvinay719:76qqKRX4zC97yQun@travis.744fuyn.mongodb.net/?retryWrites=true&w=majority&appName=travis', server_api=ServerApi('1'))
if client:
print("Connected to MongoDB")
db = client["travis"]
mongo = db["travis_face_data"]
if "travis_face_data" not in db.list_collection_names():
db.create_collection("travis_face_data")
app = Flask(__name__)
CORS(app)
# Initialize MTCNN detector and FaceNet model
detector = MTCNN()
embedder = FaceNet()
def cosine(embedding1, embedding2):
dot_product = np.dot(embedding1, embedding2)
norm1 = np.linalg.norm(embedding1)
norm2 = np.linalg.norm(embedding2)
similarity = dot_product / (norm1 * norm2)
return similarity
def reload_embeddings():
global face_data, labels, names
face_data, labels, names = load_embeddings_from_db()
@app.route("/")
def home():
return {"message": "Travis Login API is running!"}
@app.route('/login', methods=['POST'])
def recognizeLogin():
try:
if 'image' not in request.files:
return jsonify({"error": "No image provided"}), 400
file = request.files['image']
if not file:
return jsonify({"error": "Invalid file"}), 400
image_data = file.read()
image_array = np.frombuffer(image_data, np.uint8)
image = cv2.imdecode(image_array, cv2.IMREAD_COLOR)
if image is None:
return jsonify({"error": "Invalid image"}), 400
results = recognize_faces_in_image(image)
print(results[0])
if results[0]['name'] != 'unknown':
return jsonify({"name": results[0]['name'], "probability": results[0]['probability']}), 200
else:
return jsonify({'name': "user not recognised"}), 401
except Exception as e:
print(f"Error in login: {str(e)}")
return jsonify({"error": "Internal server error"}), 500
@app.route('/register', methods=['POST'])
def register():
username = request.form['username']
phoneno = request.form["phoneNumber"]
email = request.form['email']
facenet_embeddings = []
stored_image = None # To store the first grayscale image
print(username)
# Check if user already exists
existing_user = mongo.find_one({"username": username})
if existing_user:
return jsonify({"error": f"User '{username}' already exists"}), 400
# Process uploaded images
for i in range(5): # Expecting 5 images
try:
image_file = request.files[f'image{i}']
except KeyError:
return jsonify({"error": f"Missing image{i} in the request"}), 400
image_data = image_file.read()
image_array = np.frombuffer(image_data, np.uint8)
image = cv2.imdecode(image_array, cv2.IMREAD_COLOR)
# Face detection using MTCNN for FaceNet
mtcnn_faces = detector.detect_faces(image)
if mtcnn_faces:
# Get the first detected face for FaceNet embedding
x, y, w, h = mtcnn_faces[0]['box']
x, y = max(0, x), max(0, y)
cropped_face = cv2.resize(image[y:y+h, x:x+w], (160, 160))
rgb_face = cv2.cvtColor(cropped_face, cv2.COLOR_BGR2RGB)
# Get FaceNet embedding
facenet_embedding = embedder.embeddings(np.expand_dims(rgb_face, axis=0)).flatten()
facenet_embeddings.append(facenet_embedding)
# Face detection using Haar Cascade for CNN
if stored_image is None:
_, buffer = cv2.imencode('.jpg', cv2.cvtColor(cropped_face, cv2.COLOR_BGR2GRAY))
stored_image = base64.b64encode(buffer).decode('utf-8')
if not facenet_embeddings:
return jsonify({"error": "No valid faces detected in the uploaded images"}), 400
# Calculate mean embeddings
mean_facenet_embedding = np.mean(facenet_embeddings, axis=0).astype(float).tolist()
# Save model weights with error handling
# Create user data
id = mongo.count_documents({}) + 1
user_data = {
'username': username,
'phoneNumber': phoneno,
'email': email,
'embeddings': mean_facenet_embedding,
'stored_image': stored_image,
'role': 'agent',
'id': id
}
# Insert into MongoDB
mongo.insert_one(user_data)
reload_embeddings()
return jsonify({"message": "User registered successfully!"}), 201
# Load embeddings from MongoDB for recognition
def load_embeddings_from_db():
try:
users = list(mongo.find({"role": "agent"})) # Only get agent users
face_data = [] # facenet embeddings
labels = [] # id 1,2,3,..
names = {} # dict of id and username
for user in users:
if 'embeddings' in user: # Only process users with embeddings
face_data.append(user["embeddings"])
labels.append(user['id'])
names[user['id']] = user['username']
print(f"Loaded {len(face_data)} user embeddings from database")
return (face_data, labels, names) if face_data else ([], [], {})
except Exception as e:
print(f"Error loading embeddings from database: {e}")
return [], [], {}
# Load face embeddings from MongoDB initially
face_data, labels, names = load_embeddings_from_db()
def recognize_faces_in_image(image):
if len(face_data) == 0:
return [{"name": "No registered faces", "probability": 0.0}]
faces = detector.detect_faces(image)
results = []
for face in faces:
x, y, width, height = face['box']
cropped_face = cv2.resize(image[y:y+height, x:x+width], (160, 160))
# Convert cropped face to RGB
rgb_face = cv2.cvtColor(cropped_face, cv2.COLOR_BGR2RGB)
embedding = embedder.embeddings(np.expand_dims(rgb_face, axis=0)).flatten() # Use RGB face here
# Compare with stored embeddings in MongoDB
similarities = cosine_similarity([embedding], face_data)
idx = np.argmax(similarities)
best_match = similarities[0][idx]
if best_match > 0.7:
recognized_id = labels[idx] # Get the ObjectId
recognized_name = names[recognized_id] # Use ObjectId to get the username
results.append({"name": recognized_name, "probability": float(best_match)})
else:
results.append({"name": "unknown", "probability": float(best_match)})
return results
@app.route('/admin-login', methods=['POST'])
def admin_login():
data = request.get_json()
username = data.get('username')
password = data.get('password')
if not username or not password:
return jsonify({"error": "Username and password are required"}), 400
# Find admin user in database
admin = mongo.find_one({
"username": username,
"role": "admin"
})
if not admin:
return jsonify({"error": "Invalid credentials"}), 401
# In a real application, you should hash passwords and compare hashes
# For now, we'll just check if the password matches
if admin.get('password') != password:
return jsonify({"error": "Invalid credentials"}), 401
return jsonify({
"success": True,
"username": username,
"role": "admin"
}), 200
if __name__ == '__main__':
app.run(host='0.0.0.0',port=7860,debug=True)