AgeGender.py 4.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117
  1. # Import required modules
  2. import cv2 as cv
  3. import math
  4. import time
  5. import argparse
  6. def getFaceBox(net, frame, conf_threshold=0.7):
  7. frameOpencvDnn = frame.copy()
  8. frameHeight = frameOpencvDnn.shape[0]
  9. frameWidth = frameOpencvDnn.shape[1]
  10. blob = cv.dnn.blobFromImage(frameOpencvDnn, 1.0, (300, 300), [104, 117, 123], True, False)
  11. net.setInput(blob)
  12. detections = net.forward()
  13. bboxes = []
  14. for i in range(detections.shape[2]):
  15. confidence = detections[0, 0, i, 2]
  16. if confidence > conf_threshold:
  17. x1 = int(detections[0, 0, i, 3] * frameWidth)
  18. y1 = int(detections[0, 0, i, 4] * frameHeight)
  19. x2 = int(detections[0, 0, i, 5] * frameWidth)
  20. y2 = int(detections[0, 0, i, 6] * frameHeight)
  21. bboxes.append([x1, y1, x2, y2])
  22. cv.rectangle(frameOpencvDnn, (x1, y1), (x2, y2), (0, 255, 0), int(round(frameHeight/150)), 8)
  23. return frameOpencvDnn, bboxes
  24. parser = argparse.ArgumentParser(description='Use this script to run age and gender recognition using OpenCV.')
  25. parser.add_argument('--input', help='Path to input image or video file. Skip this argument to capture frames from a camera.')
  26. parser.add_argument("--device", default="cpu", help="Device to inference on")
  27. args = parser.parse_args()
  28. args = parser.parse_args()
  29. faceProto = "opencv_face_detector.pbtxt"
  30. faceModel = "opencv_face_detector_uint8.pb"
  31. ageProto = "age_deploy.prototxt"
  32. ageModel = "age_net.caffemodel"
  33. genderProto = "gender_deploy.prototxt"
  34. genderModel = "gender_net.caffemodel"
  35. MODEL_MEAN_VALUES = (78.4263377603, 87.7689143744, 114.895847746)
  36. ageList = ['(0-2)', '(4-6)', '(8-12)', '(15-20)', '(25-32)', '(38-43)', '(48-53)', '(60-100)']
  37. genderList = ['Male', 'Female']
  38. # Load network
  39. ageNet = cv.dnn.readNet(ageModel, ageProto)
  40. genderNet = cv.dnn.readNet(genderModel, genderProto)
  41. faceNet = cv.dnn.readNet(faceModel, faceProto)
  42. if args.device == "cpu":
  43. ageNet.setPreferableBackend(cv.dnn.DNN_TARGET_CPU)
  44. genderNet.setPreferableBackend(cv.dnn.DNN_TARGET_CPU)
  45. faceNet.setPreferableBackend(cv.dnn.DNN_TARGET_CPU)
  46. print("Using CPU device")
  47. elif args.device == "gpu":
  48. ageNet.setPreferableBackend(cv.dnn.DNN_BACKEND_CUDA)
  49. ageNet.setPreferableTarget(cv.dnn.DNN_TARGET_CUDA)
  50. genderNet.setPreferableBackend(cv.dnn.DNN_BACKEND_CUDA)
  51. genderNet.setPreferableTarget(cv.dnn.DNN_TARGET_CUDA)
  52. genderNet.setPreferableBackend(cv.dnn.DNN_BACKEND_CUDA)
  53. genderNet.setPreferableTarget(cv.dnn.DNN_TARGET_CUDA)
  54. print("Using GPU device")
  55. # Open a video file or an image file or a camera stream
  56. cap = cv.VideoCapture(args.input if args.input else 0)
  57. padding = 20
  58. while cv.waitKey(1) < 0:
  59. # Read frame
  60. t = time.time()
  61. hasFrame, frame = cap.read()
  62. if not hasFrame:
  63. cv.waitKey()
  64. break
  65. frameFace, bboxes = getFaceBox(faceNet, frame)
  66. if not bboxes:
  67. print("No face Detected, Checking next frame")
  68. continue
  69. for bbox in bboxes:
  70. # print(bbox)
  71. face = frame[max(0,bbox[1]-padding):min(bbox[3]+padding,frame.shape[0]-1),max(0,bbox[0]-padding):min(bbox[2]+padding, frame.shape[1]-1)]
  72. blob = cv.dnn.blobFromImage(face, 1.0, (227, 227), MODEL_MEAN_VALUES, swapRB=False)
  73. genderNet.setInput(blob)
  74. genderPreds = genderNet.forward()
  75. gender = genderList[genderPreds[0].argmax()]
  76. # print("Gender Output : {}".format(genderPreds))
  77. print("Gender : {}, conf = {:.3f}".format(gender, genderPreds[0].max()))
  78. ageNet.setInput(blob)
  79. agePreds = ageNet.forward()
  80. age = ageList[agePreds[0].argmax()]
  81. print("Age Output : {}".format(agePreds))
  82. print("Age : {}, conf = {:.3f}".format(age, agePreds[0].max()))
  83. label = "{},{}".format(gender, age)
  84. cv.putText(frameFace, label, (bbox[0], bbox[1]-10), cv.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 255), 2, cv.LINE_AA)
  85. cv.imshow("Age Gender Demo", frameFace)
  86. # cv.imwrite("age-gender-out-{}".format(args.input),frameFace)
  87. print("time : {:.3f}".format(time.time() - t))
  88. # cmake -DCMAKE_BUILD_TYPE=RELEASE -DCMAKE_INSTALL_PREFIX=~/opencv_gpu -DINSTALL_PYTHON_EXAMPLES=OFF -DINSTALL_C_EXAMPLES=OFF -DOPENCV_ENABLE_NONFREE=ON -DOPENCV_EXTRA_MODULES_PATH=~/cv2_gpu/opencv_contrib/modules -DPYTHON_EXECUTABLE=~/env/bin/python3 -DBUILD_EXAMPLES=ON -DWITH_CUDA=ON -DWITH_CUDNN=ON -DOPENCV_DNN_CUDA=ON -DENABLE_FAST_MATH=ON -DCUDA_FAST_MATH=ON -DWITH_CUBLAS=ON -DCUDA_TOOLKIT_ROOT_DIR=/usr/local/cuda-10.2 -DOpenCL_LIBRARY=/usr/local/cuda-10.2/lib64/libOpenCL.so -DOpenCL_INCLUDE_DIR=/usr/local/cuda-10.2/include/ ..