main.py 2.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101
  1. import hashlib
  2. import json
  3. import pickle
  4. import cv2
  5. import dlib
  6. import numpy as np
  7. import requests
  8. from flask import Flask, request, jsonify
  9. app = Flask(__name__)
  10. app.config['JSON_SORT_KEYS'] = False
  11. detector = dlib.get_frontal_face_detector()
  12. sp = dlib.shape_predictor("models/shape_predictor_68_face_landmarks.dat")
  13. facerec = dlib.face_recognition_model_v1("models/dlib_face_recognition_resnet_model_v1.dat")
  14. def detect_one_file(data):
  15. data = np.frombuffer(data, np.uint8)
  16. if data is None:
  17. raise Exception('image is required.')
  18. zoom_ratio = 1
  19. # if data.size > 6 * 1024 * 1024:
  20. # img = cv2.imdecode(data, cv2.IMREAD_REDUCED_COLOR_4)
  21. # zoom_ratio = 4
  22. # elif data.size > 4 * 1024 * 124:
  23. # img = cv2.imdecode(data, cv2.IMREAD_REDUCED_COLOR_2)
  24. # zoom_ratio = 2
  25. # else:
  26. img = cv2.imdecode(data, cv2.IMREAD_COLOR)
  27. faces = []
  28. dets = detector(img, 1)
  29. for d in dets:
  30. if d.width() * zoom_ratio < 100:
  31. continue
  32. shape = sp(img, d)
  33. des_buffer = pickle.dumps(facerec.compute_face_descriptor(img, shape))
  34. h = hashlib.sha256()
  35. h.update(des_buffer)
  36. id = h.hexdigest()
  37. f = {
  38. 'id': id,
  39. 'des_buffer': des_buffer
  40. }
  41. res = requests.put('http://127.0.0.1:5001/sqlite/insert_face', data=pickle.dumps(f))
  42. if not res.ok:
  43. raise Exception("insert one face to the sqlite failed.")
  44. result = json.loads(res.content)
  45. if result is None or result['success'] is None or not result['success']:
  46. raise Exception("insert one face to the sqlite failed.")
  47. face = {
  48. 'token': id,
  49. 'rectangle': {
  50. 'width': d.width() * zoom_ratio,
  51. 'height': d.height() * zoom_ratio,
  52. 'left': d.left() * zoom_ratio,
  53. 'top': d.top() * zoom_ratio,
  54. },
  55. 'age': 0,
  56. 'quality': 0,
  57. 'gender': 0,
  58. }
  59. faces.append(face)
  60. return faces
  61. @app.route("/api/detect", methods=['PUT'])
  62. def face_detect():
  63. try:
  64. faces = detect_one_file(request.get_data())
  65. response = {
  66. 'success': True,
  67. 'count': len(faces),
  68. 'faces': faces,
  69. }
  70. resp = jsonify(response)
  71. resp.status_code = 200
  72. except Exception as e:
  73. response = {
  74. 'success': False,
  75. 'count': 0,
  76. 'faces': [],
  77. }
  78. resp = jsonify(response)
  79. resp.status_code = 400
  80. return resp
  81. if __name__ == '__main__':
  82. app.run(
  83. host='0.0.0.0',
  84. port=5000,
  85. debug=True
  86. )