sqlite_server.py 2.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110
  1. import pickle
  2. import sqlite3
  3. from datetime import datetime, timezone
  4. import numpy as np
  5. from flask import Flask, request
  6. from flask.json import loads, jsonify
  7. from dao import SQLiteUtil
  8. DBPATH = "face_recon.db"
  9. app = Flask(__name__)
  10. app.config['JSON_SORT_KEYS'] = False
  11. sqlite = SQLiteUtil(DBPATH)
  12. @app.route("/sqlite/insert_face", methods=['PUT'])
  13. def insert_face():
  14. data = request.get_data()
  15. face = None
  16. if data is not None:
  17. face = pickle.loads(data)
  18. if face is not None:
  19. sqlite.execute("INSERT INTO faces values(?,?,?)",
  20. [face['id'], sqlite3.Binary(face['des_buffer']), datetime.now(timezone.utc)])
  21. res = {
  22. 'success': True,
  23. 'message': 'successfully insert the face: ' + face['id']
  24. }
  25. else:
  26. res = {
  27. 'success': False,
  28. 'message': 'the input faceid is empty.'
  29. }
  30. response = jsonify(res)
  31. response.status_code = 200
  32. return response
  33. @app.route("/api/search", methods=['PUT'])
  34. def face_search():
  35. face_id_in = request.args.get('facetoken')
  36. if face_id_in is None:
  37. response = {
  38. 'success': False,
  39. 'messages': "the input face id is empty.",
  40. }
  41. resp = jsonify(response)
  42. resp.status_code = 400
  43. return resp
  44. max_distance = request.args.get('max-distance')
  45. if max_distance is None:
  46. max_distance = 0.3
  47. else:
  48. max_distance = float(max_distance)
  49. res = sqlite.execute_query("SELECT facetoken from faces where id=?;", [face_id_in])
  50. if len(res) != 1:
  51. response = {
  52. 'success': False,
  53. 'messages': "the input face id is not found in the database.",
  54. }
  55. resp = jsonify(response)
  56. resp.status_code = 400
  57. return resp
  58. face_token_in = res[0]['facetoken']
  59. if face_token_in is None:
  60. response = {
  61. 'success': False,
  62. 'messages': "the input face id has no face token found.",
  63. }
  64. resp = jsonify(response)
  65. resp.status_code = 400
  66. return resp
  67. face_in = pickle.loads(face_token_in)
  68. face_in = np.array(face_in)
  69. faces = []
  70. result = sqlite.execute_query("SELECT id,facetoken FROM faces where id != ?;", [face_id_in])
  71. for record in result:
  72. face_token = record['facetoken']
  73. face = pickle.loads(face_token)
  74. face = np.array(face)
  75. res = np.linalg.norm([face] - face_in, axis=1)
  76. if res[0] <= max_distance:
  77. face = {
  78. 'token_id': record['id'],
  79. 'distance': res[0],
  80. }
  81. faces.append(face)
  82. response = {
  83. 'success': True,
  84. 'count': len(faces),
  85. 'faces': faces,
  86. }
  87. resp = jsonify(response)
  88. resp.status_code = 200
  89. return resp
  90. if __name__ == '__main__':
  91. app.run(
  92. host='0.0.0.0',
  93. port=5001,
  94. debug=True
  95. )