Data hosted with ♥ by Pastebin.com - Download Raw - See Original
  1. from MultiMsgSYncV2 import TwoStageHostSeqSync
  2. import blobconverter
  3. import cv2
  4. import depthai as dai
  5. import numpy as np
  6.  
  7. emotionsclass = ['neutral', 'happy', 'sad', 'surprise', 'anger']
  8.  
  9. def frame_norm(frame, bbox):
  10.     normVals = np.full(len(bbox), frame.shape[0])
  11.     normVals[::2] = frame.shape[1]
  12.     return (np.clip(np.array(bbox), 0, 1) * normVals).astype(int)
  13.  
  14. def create_pipeline(stereo):
  15.     pipeline = dai.Pipeline()
  16.  
  17.     cam = pipeline.create(dai.node.ColorCamera)
  18.     cam.setPreviewSize(640, 400)
  19.     cam.setResolution(dai.ColorCameraProperties.SensorResolution.THE_1080_P)
  20.     cam.setInterleaved(False)
  21.     cam.setBoardSocket(dai.CameraBoardSocket.RGB)
  22.  
  23.     manipRgb = pipeline.create(dai.node.ImageManip)
  24.     rgbRr = dai.RotatedRect()
  25.     rgbRr.center.x, rgbRr.center.y = cam.getPreviewWidth() // 2, cam.getPreviewHeight() // 2
  26.     rgbRr.size.width, rgbRr.size.height = cam.getPreviewHeight(), cam.getPreviewWidth()
  27.     rgbRr.angle = -90
  28.     manipRgb.initialConfig.setCropRotatedRect(rgbRr, False)
  29.  
  30.     manipRgbOut = pipeline.create(dai.node.XLinkOut)
  31.     manipRgbOut.setStreamName("color")
  32.     manipRgb.out.link(manipRgbOut.input)
  33.     cam.preview.link(manipRgb.inputImage)
  34.  
  35.     # ImageManip will resize the frame before sending it to the Face detection NN node
  36.     face_det_manip = pipeline.create(dai.node.ImageManip)
  37.     face_det_manip.initialConfig.setResize(300, 300)
  38.     face_det_manip.initialConfig.setFrameType(dai.RawImgFrame.Type.RGB888p)
  39.     rrFD = dai.RotatedRect()
  40.     rrFD.center.x, rrFD.center.y = cam.getPreviewWidth() // 2, cam.getPreviewHeight() // 2
  41.     rrFD.size.width, rrFD.size.height = cam.getPreviewHeight(), cam.getPreviewWidth()
  42.     rrFD.angle = -90
  43.     face_det_manip.initialConfig.setCropRotatedRect(rrFD, False)
  44.     cam.preview.link(face_det_manip.inputImage)
  45.  
  46.     face_det_nn = pipeline.create(dai.node.MobileNetDetectionNetwork)
  47.  
  48.     face_det_nn.setConfidenceThreshold(0.5)
  49.     face_det_nn.setBlobPath(blobconverter.from_zoo(name="face-detection-retail-0004", shaves=6))
  50.     face_det_manip.out.link(face_det_nn.input)
  51.  
  52.     # Send face detections to the host (for bounding boxes)
  53.     face_det_xout = pipeline.create(dai.node.XLinkOut)
  54.     face_det_xout.setStreamName("detection")
  55.     face_det_nn.out.link(face_det_xout.input)
  56.  
  57.     # Script node will take the output from the face detection NN as an input and set ImageManipConfig
  58.     # to the 'recognition_manip' to crop the initial frame
  59.     image_manip_script = pipeline.create(dai.node.Script)
  60.     face_det_nn.out.link(image_manip_script.inputs['face_det_in'])
  61.  
  62.     # Only send metadata, we are only interested in timestamp, so we can sync
  63.     # depth frames with NN output
  64.     face_det_nn.passthrough.link(image_manip_script.inputs['passthrough'])
  65.  
  66.     image_manip_script.setScript("""
  67.   l = [] # List of images
  68.   # So the correct frame will be the first in the list
  69.   # For this experiment this function is redundant, since everything
  70.   # runs in blocking mode, so no frames will get lost
  71.   def get_latest_frame(seq):
  72.       global l
  73.       for i, frame in enumerate(l):
  74.           #node.io['manip_frame'].send(frame)
  75.           if seq == frame.getSequenceNum():
  76.               # node.warn(f"List len {len(l)} Frame with same seq num: {i},seq {seq}")
  77.               l = l[i:]
  78.               break
  79.       return l[0]
  80.  
  81.   def correct_bb(bb):
  82.       if bb.xmin < 0: bb.xmin = 0.001
  83.       if bb.ymin < 0: bb.ymin = 0.001
  84.       if bb.xmax > 1: bb.xmax = 0.999
  85.       if bb.ymax > 1: bb.ymax = 0.999
  86.       return bb
  87.   while True:
  88.       preview = node.io['preview'].tryGet()
  89.       if preview is not None:
  90.           # node.warn(f"New frame {preview.getSequenceNum()}")
  91.           l.append(preview)
  92.  
  93.       face_dets = node.io['face_det_in'].tryGet()
  94.       # node.warn(f"Faces detected: {len(face_dets)}")
  95.       if face_dets is not None:
  96.           passthrough = node.io['passthrough'].get()
  97.           seq = passthrough.getSequenceNum()
  98.           # node.warn(f"New detection {seq}")
  99.           if len(l) == 0:
  100.               continue
  101.           img = get_latest_frame(seq)
  102.  
  103.           for i, det in enumerate(face_dets.detections):
  104.               cfg = ImageManipConfig()
  105.               correct_bb(det)
  106.               cfg.setCropRect(det.xmin, det.ymin, det.xmax, det.ymax)
  107.               # node.warn(f"Sending {i + 1}. det. Seq {seq}. Det {det.xmin}, {det.ymin}, {det.xmax}, {det.ymax}")
  108.               cfg.setResize(62, 62)
  109.               cfg.setKeepAspectRatio(False)
  110.               node.io['manip_cfg'].send(cfg)
  111.               node.io['manip_img'].send(img)
  112.  
  113.               cfg2 = ImageManipConfig()
  114.               cfg2.setCropRect(det.xmin, det.ymin, det.xmax, det.ymax)
  115.               # node.warn(f"Sending {i + 1}. det. Seq {seq}. Det {det.xmin}, {det.ymin}, {det.xmax}, {det.ymax}")
  116.               cfg2.setResize(64, 64)
  117.               cfg2.setKeepAspectRatio(False)
  118.               node.io['emanip_cfg'].send(cfg2)
  119.               node.io['emanip_img'].send(img)
  120.   """)
  121.     #cam.preview.link(image_manip_script.inputs['preview'])
  122.     manipRgb.out.link(image_manip_script.inputs['preview'])
  123.  
  124.     recognition_manip = pipeline.create(dai.node.ImageManip)
  125.     recognition_manip.initialConfig.setResize(62, 62)
  126.     #recognition_manip.setWaitForConfigInput(True)
  127.     image_manip_script.outputs['manip_cfg'].link(recognition_manip.inputConfig)
  128.     image_manip_script.outputs['manip_img'].link(recognition_manip.inputImage)
  129.  
  130.     # face_cropped_xout = pipeline.create(dai.node.XLinkOut)
  131.     # face_cropped_xout.setStreamName("face_cropped")
  132.     # recognition_manip.out.link(face_cropped_xout.input)
  133.  
  134.     # frame_xout = pipeline.create(dai.node.XLinkOut)
  135.     # frame_xout.setStreamName("frame_xout")
  136.     # image_manip_script.outputs['manip_frame'].link(frame_xout.input)
  137.  
  138.     # Second stange recognition NN
  139.     recognition_nn = pipeline.create(dai.node.NeuralNetwork)
  140.     recognition_nn.setBlobPath(blobconverter.from_zoo(name="age-gender-recognition-retail-0013", shaves=6))
  141.     recognition_manip.out.link(recognition_nn.input)
  142.  
  143.     recognition_xout = pipeline.create(dai.node.XLinkOut)
  144.     recognition_xout.setStreamName("recognition")
  145.     recognition_nn.out.link(recognition_xout.input)
  146.  
  147.     # third stange expression NN
  148.     expression_manip = pipeline.create(dai.node.ImageManip)
  149.     expression_manip.initialConfig.setResize(64, 64)
  150.     #expression_manip.setWaitForConfigInput(True)
  151.     image_manip_script.outputs['emanip_cfg'].link(expression_manip.inputConfig)
  152.     image_manip_script.outputs['emanip_img'].link(expression_manip.inputImage)
  153.  
  154.     expression_nn = pipeline.create(dai.node.NeuralNetwork)
  155.     expression_nn.setBlobPath(blobconverter.from_zoo(name="emotions-recognition-retail-0003", shaves=6))
  156.     expression_manip.out.link(expression_nn.input)
  157.  
  158.     expression_xout = pipeline.create(dai.node.XLinkOut)
  159.     expression_xout.setStreamName("emotions")
  160.     expression_nn.out.link(expression_xout.input)
  161.  
  162.     return pipeline
  163.  
  164. with dai.Device() as device:
  165.     stereo = False
  166.     device.startPipeline(create_pipeline(stereo))
  167.     # face_cropped_q = device.getOutputQueue("face_cropped", 4, False)
  168.     # frame_xout_q = device.getOutputQueue("frame_xout", 4, False)
  169.  
  170.     sync = TwoStageHostSeqSync()
  171.     queues = {}
  172.     # Create output queues
  173.     for name in ["color", "detection", "recognition", "emotions"]:
  174.         queues[name] = device.getOutputQueue(name)
  175.  
  176.     while True:
  177.         try:
  178.             for name, q in queues.items():
  179.                 # Add all msgs (color frames, object detections and recognitions) to the Sync class.
  180.                 if q.has():
  181.                     sync.add_msg(q.get(), name)
  182.  
  183.             # face_cropped_in = face_cropped_q.tryGet()
  184.             # if face_cropped_in is not None:
  185.             #     cv2.imshow("cropped", face_cropped_in.getCvFrame())
  186.  
  187.             # frame_in = frame_xout_q.tryGet()
  188.             # if frame_in is not None:
  189.             #     cv2.imshow("frame on host", frame_in.getCvFrame())
  190.            
  191.             msgs = sync.get_msgs()
  192.             resultList={"age":[],"gen_der":[],"dis_tance":[],"expression":[],"dtstamp":[]}
  193.  
  194.             if msgs is not None:
  195.                 print("SYnced frames")
  196.                
  197.                 frame = msgs["color"].getCvFrame()
  198.                 detections = msgs["detection"].detections
  199.                 recognitions = msgs["recognition"]
  200.                 expressions = msgs["emotions"]
  201.  
  202.                 for i, detection in enumerate(detections):
  203.                     print("New detection", i)
  204.                     bbox = frame_norm(frame, (detection.xmin, detection.ymin, detection.xmax, detection.ymax))
  205.  
  206.                     # Decoding of recognition results
  207.                     age = 0
  208.                     gender_str = ""
  209.                     try:
  210.                         rec = recognitions[i]
  211.                         age = int(float(np.squeeze(np.array(rec.getLayerFp16('age_conv3')))) * 100) + 10
  212.                         gender = np.squeeze(np.array(rec.getLayerFp16('prob')))
  213.                         gender_str = "female" if gender[0] > gender[1] else "male"
  214.                     except:
  215.                         pass
  216.  
  217.                     emotion_name = ""
  218.                     try:
  219.                         exp = expressions[i]
  220.                         emotion_results = np.array(exp.getFirstLayerFp16())
  221.                         emotion_name = emotionsclass[np.argmax(emotion_results)]
  222.                     except:
  223.                         pass
  224.  
  225.                     Distance = 0
  226.  
  227.                     det_frame = frame[bbox[1]:bbox[3], bbox[0]:bbox[2]]
  228.                     fh, fw, fc = det_frame.shape
  229.                     frame_h, frame_w, frame_c = frame.shape
  230.                    
  231.                     #cv2.rectangle(frame, (bbox[0], bbox[1]), (bbox[2], bbox[3]), (10, 245, 10), 2)
  232.                     print(bbox)
  233.                     frame = cv2.rectangle(frame,(bbox[0],bbox[1]), (bbox[2],bbox[3]), (0 ,127 , 255), 2)
  234.                     y = ((bbox[1] + bbox[3]) // 2) - 30
  235.                     cv2.putText(frame, str(age), (bbox[0], y), cv2.FONT_HERSHEY_SIMPLEX, 1.5, (0, 255, 255), 2)
  236.                     cv2.putText(frame, gender_str, (bbox[0], y + 30), cv2.FONT_HERSHEY_SIMPLEX, 1.5, (0, 255, 255), 2)
  237.                     cv2.putText(frame, emotion_name, (bbox[0], y + 60), cv2.FONT_HERSHEY_SIMPLEX, 1.5, (0, 255, 255), 2)
  238.  
  239.                 cv2.imshow("Camera", frame)
  240.                
  241.             if cv2.waitKey(1) == ord('q'):
  242.                 break
  243.  
  244.         except:
  245.             pass
  246.  
  247. cv2.destroyAllWindows()