main_pipeline_webcam.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cv2\n",
    "import numpy as np\n",
    "\n",
    "import pycuda.driver as cuda\n",
    "import pycuda.autoinit\n",
    "import tensorrt as trt\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.patches as patches\n",
    "\n",
    "from itertools import product"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# G-Streamer Pipelines"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Opening in BLOCKING MODE \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "NvMMLiteOpen : Block : BlockType = 261 \n",
      "NVMEDIA: Reading vendor.tegra.display-size : status: 6 \n",
      "NvMMLiteBlockCreate : Block : BlockType = 261 \n",
      "[ WARN:0] global /tmp/build_opencv/opencv/modules/videoio/src/cap_gstreamer.cpp (1100) open OpenCV | GStreamer warning: Cannot query video position: status=1, value=12, duration=-1\n"
     ]
    }
   ],
   "source": [
    "def gstr_in():\n",
    "    return (\n",
    "        \"v4l2src device=/dev/video0 ! \"\n",
    "        \"video/x-h264, width=640, height=480, framerate=20/1 ! \"\n",
    "        \"h264parse ! \"\n",
    "        \"nvv4l2decoder ! \"\n",
    "        \"nvvidconv ! \"\n",
    "        \"video/x-raw, format=(string)BGRx ! \"\n",
    "        \"videoconvert ! \"\n",
    "        \"video/x-raw, format=(string)BGR ! \"\n",
    "        \"queue max-size-buffers=10 max-size-time=0 max-size-bytes=0 leaky=downstream drop=True ! \"\n",
    "        \"appsink\"\n",
    "    )\n",
    "\n",
    "pipe_in  = gstr_in()\n",
    "pipe_in  = cv2.VideoCapture(pipe_in, cv2.CAP_GSTREAMER)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Using winsys: x11 \n"
     ]
    }
   ],
   "source": [
    "width  = int(pipe_in.get(cv2.CAP_PROP_FRAME_HEIGHT))\n",
    "height = int(pipe_in.get(cv2.CAP_PROP_FRAME_WIDTH))\n",
    "fps    = int(pipe_in.get(cv2.CAP_PROP_FPS))\n",
    "\n",
    "def gstr_out():\n",
    "    return (\n",
    "        \"appsrc ! \"\n",
    "        \"videoconvert ! \"\n",
    "        \"video/x-raw, format=(string)BGRx ! \"\n",
    "        \"nvvidconv ! \"\n",
    "        \"video/x-raw(memory:NVMM), format=(string)NV12 ! \"\n",
    "        \"nvegltransform ! \"\n",
    "        \"nveglglessink sync=false\"\n",
    "    )\n",
    "\n",
    "pipe_out = gstr_out()\n",
    "pipe_out = cv2.VideoWriter(pipe_out, cv2.CAP_GSTREAMER, 0, fps, (height, width), True)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Cuda Needs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_engine(engine_file_path=None):\n",
    "    with open(engine_file_path, \"rb\") as f, trt.Runtime(trt.Logger(trt.Logger.WARNING)) as runtime:\n",
    "        return runtime.deserialize_cuda_engine(f.read())\n",
    "    \n",
    "def allocate_buffers_um_pinned(engine):\n",
    "    \n",
    "    inputs, outputs, bindings = [], [], []\n",
    "    \n",
    "    for binding in engine:\n",
    "\n",
    "        size  = trt.volume(engine.get_binding_shape(binding)) #* engine.max_batch_size\n",
    "        dtype = trt.nptype(engine.get_binding_dtype(binding))\n",
    "\n",
    "        # Allocate SINGLE pinned memory buffer\n",
    "        mem = cuda.pagelocked_empty(size, dtype)\n",
    "        \n",
    "        # Append the device buffer to device bindings.\n",
    "        bindings.append(int(mem.base.get_device_pointer()))\n",
    "        \n",
    "        # Append to the appropriate list.\n",
    "        if engine.binding_is_input(binding): inputs.append(mem)\n",
    "        else:                                outputs.append(mem)\n",
    "    \n",
    "    return inputs, outputs, bindings\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Preprocess Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "#DexiNed\n",
    "#This implementation is based on the DexiNed from X. Soria and E. Riba and A. Sappa\n",
    "#Source: https://github.com/xavysp/DexiNed\n",
    "\n",
    "def dexi_pre(frame):\n",
    "\n",
    "    frame = np.array(frame, dtype=np.float32)\n",
    "    frame -= [103.939, 116.779, 123.68] #Mean subtraction\n",
    "    frame = np.transpose(frame, (2, 0, 1)) #Move channels to the first index\n",
    "    frame = np.expand_dims(frame, axis=0) \n",
    "    frame = frame.ravel()\n",
    "\n",
    "    return frame\n",
    "\n",
    "def sig(x): return 1 / (1 + np.exp(-x))\n",
    "\n",
    "def dexi_post(frame):\n",
    "\n",
    "    frame = 1 / (1 + np.exp(-frame))\n",
    "    frame = frame * 255 #Normalization also possible\n",
    "    frame = frame.reshape((480, 640)) \n",
    "    frame = cv2.cvtColor(frame, cv2.COLOR_GRAY2BGR)\n",
    "    #frame = frame.astype(np.uint8)\n",
    "\n",
    "    return frame\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "#YuNet\n",
    "#This implementation is based on the YuNet from the OpenCV Zoo\n",
    "#Source: https://github.com/opencv/opencv_zoo/blob/bfac311b2b30de4648307d9939d2f9e33e012007/models/face_detection_yunet/yunet.py\n",
    "\n",
    "def facenet_pre(image):\n",
    "    image = cv2.resize(image, (160, 120))\n",
    "    image = np.asarray(image, dtype=\"float32\")\n",
    "    image = image.transpose(2, 0, 1)\n",
    "    image = image.ravel()\n",
    "    \n",
    "    return image\n",
    "\n",
    "min_sizes = [[10, 16, 24], [32, 48], [64, 96], [128, 192, 256]]\n",
    "steps     = [8, 16, 32, 64]\n",
    "variance  = [0.1, 0.2]\n",
    "\n",
    "def priorGen(min_sizes, steps):\n",
    "    \n",
    "    w, h = (160, 120)\n",
    "    feature_map_2th = [int(int((h + 1) / 2) / 2), int(int((w + 1) / 2) / 2)]\n",
    "    feature_map_3th = [int(feature_map_2th[0] / 2), int(feature_map_2th[1] / 2)]\n",
    "    feature_map_4th = [int(feature_map_3th[0] / 2), int(feature_map_3th[1] / 2)]\n",
    "    feature_map_5th = [int(feature_map_4th[0] / 2), int(feature_map_4th[1] / 2)]\n",
    "    feature_map_6th = [int(feature_map_5th[0] / 2), int(feature_map_5th[1] / 2)]\n",
    "\n",
    "    feature_maps = [\n",
    "        feature_map_3th,\n",
    "        feature_map_4th,\n",
    "        feature_map_5th,\n",
    "        feature_map_6th,\n",
    "    ]\n",
    "\n",
    "    priors = []\n",
    "    for k, f in enumerate(feature_maps):\n",
    "        _min_sizes = min_sizes[k]\n",
    "        \n",
    "        for i, j in product(range(f[0]), range(f[1])):  # i->h, j->w\n",
    "            for min_size in _min_sizes:\n",
    "                s_kx = min_size / w\n",
    "                s_ky = min_size / h\n",
    "\n",
    "                cx = (j + 0.5) * steps[k] / w\n",
    "                cy = (i + 0.5) * steps[k] / h\n",
    "\n",
    "                priors.append([cx, cy, s_kx, s_ky])\n",
    "\n",
    "    return np.array(priors, dtype=np.float32) #priors\n",
    "\n",
    "def decode(outputBlob, priors, variance):\n",
    "\n",
    "    loc  = np.array(outputBlob[0]).reshape([-1, 14])\n",
    "    conf = np.array(outputBlob[1]).reshape([-1, 2])\n",
    "    iou  = np.array(outputBlob[2]).reshape([-1, 1])\n",
    "\n",
    "    # get score\n",
    "    cls_scores = conf[:, 1]\n",
    "    iou_scores = iou[:, 0]\n",
    "\n",
    "    # clamp\n",
    "    _idx = np.where(iou_scores < 0.0)\n",
    "    iou_scores[_idx] = 0.0\n",
    "\n",
    "    _idx = np.where(iou_scores > 1.0)\n",
    "    iou_scores[_idx] = 1.0\n",
    "    \n",
    "    scores = np.sqrt(cls_scores * iou_scores)\n",
    "    scores = scores[:, np.newaxis]\n",
    "    \n",
    "    scale = np.array((160, 120))\n",
    "\n",
    "    # get bboxes\n",
    "    bboxes = np.hstack(((priors[:, 0:2]+ loc[:, 0:2] * variance[0] * priors[:, 2:4]) * scale, (priors[:, 2:4] * np.exp(loc[:, 2:4] * variance)) * scale,))\n",
    "    # (x_c, y_c, w, h) -> (x1, y1, w, h)\n",
    "    bboxes[:, 0:2] -= bboxes[:, 2:4] / 2\n",
    "    \n",
    "    # get landmarks\n",
    "    landmarks = np.hstack(\n",
    "        (\n",
    "            (priors[:, 0:2] + loc[:, 4:6] * variance[0] * priors[:, 2:4]) * scale,\n",
    "            (priors[:, 0:2] + loc[:, 6:8] * variance[0] * priors[:, 2:4]) * scale,\n",
    "            (priors[:, 0:2] + loc[:, 8:10] * variance[0] * priors[:, 2:4]) * scale,\n",
    "            (priors[:, 0:2] + loc[:, 10:12] * variance[0] * priors[:, 2:4]) * scale,\n",
    "            (priors[:, 0:2] + loc[:, 12:14] * variance[0] * priors[:, 2:4]) * scale,\n",
    "        )\n",
    "    )\n",
    "\n",
    "    return np.hstack((bboxes, landmarks, scores))\n",
    "\n",
    "priors = priorGen(min_sizes, steps)\n",
    "\n",
    "def facedet_post(output_blob, priors=priors, variance=variance):\n",
    "    # Decode\n",
    "    dets = decode(output_blob, priors, variance)\n",
    "\n",
    "    # NMS\n",
    "    keepIdx = cv2.dnn.NMSBoxes(\n",
    "        bboxes          =dets[:, 0:4].tolist(),\n",
    "        scores          =dets[:, -1].tolist(),\n",
    "        score_threshold =0.6, #self._conf_threshold,\n",
    "        nms_threshold   =0.3, #self._nms_threshold,\n",
    "        top_k           =5000,#self._top_k,\n",
    "    )  # box_num x class_num\n",
    "\n",
    "    if len(keepIdx) > 0:\n",
    "        dets = dets[keepIdx]\n",
    "        #dets = np.squeeze(dets, axis=1)\n",
    "        return dets[: 5000] #self._top_k\n",
    "    \n",
    "    else:\n",
    "        return np.empty(shape=(0, 15))\n",
    "\n",
    "#Calculate the scaling factors\n",
    "old_size = (160, 120)\n",
    "new_size = (640, 480)\n",
    "scale_w = new_size[0] / old_size[0]\n",
    "scale_h = new_size[1] / old_size[1]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Dynamic Patameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tkinter as tk\n",
    "\n",
    "# Initial value\n",
    "param, param1 = 0, 0\n",
    "face_b = False\n",
    "\n",
    "def set_param():\n",
    "    global param\n",
    "    param = int(entry.get())\n",
    "\n",
    "def set_param1():\n",
    "    global param1\n",
    "    param1 = int(entry.get())\n",
    "\n",
    "def update_label():\n",
    "    global face_b\n",
    "    face_b = not face_b\n",
    "\n",
    "root = tk.Tk()\n",
    "\n",
    "# Create an Entry widget\n",
    "entry = tk.Entry(root)\n",
    "entry.pack()\n",
    "\n",
    "# Create a Button that will call set_param when clicked\n",
    "button = tk.Button(root, text='Set Param', command=set_param)\n",
    "button.pack()\n",
    "\n",
    "# Create a Button that will call set_param when clicked\n",
    "button1 = tk.Button(root, text='Set Param1', command=set_param1)\n",
    "button1.pack()\n",
    "\n",
    "checkbutton = tk.Checkbutton(root, text=\"Toggle me\", command=update_label)\n",
    "checkbutton.pack()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "lmks_face       = np.arange(0,60)\n",
    "lmks_face_outer = np.append(np.arange(0,17), np.arange(17,27)[::-1])\n",
    "lmks_eyes       = np.arange(36,48)\n",
    "lmks_nose       = np.arange(27,36)\n",
    "lmks_mouth      = np.arange(48,60)\n",
    "\n",
    "lmks_left_eye  = np.arange(36,42)\n",
    "lmks_right_eye = np.arange(42,48)\n",
    "\n",
    "lmks = [lmks_face, lmks_eyes, lmks_nose, lmks_mouth]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Main Loop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[05/24/2023-14:17:00] [TRT] [W] Using an engine plan file across different models of devices is not recommended and is likely to affect performance or even cause errors.\n",
      "[05/24/2023-14:17:01] [TRT] [W] Using an engine plan file across different models of devices is not recommended and is likely to affect performance or even cause errors.\n",
      "[05/24/2023-14:17:01] [TRT] [W] Using an engine plan file across different models of devices is not recommended and is likely to affect performance or even cause errors.\n"
     ]
    }
   ],
   "source": [
    "dir_engine_dexi    = '/home/gzguevara/laser/engines/dexi_block4_640x480_fp16.trt'\n",
    "dir_engine_facedet = '/home/gzguevara/laser/engines/face_detection_yunet_2022mar.trt'\n",
    "dir_engine_lmk     = '/home/gzguevara/laser/engines/landmarks.trt'\n",
    "\n",
    "#Build an engine\n",
    "engine_dexi    = get_engine(dir_engine_dexi)\n",
    "engine_facedet = get_engine(dir_engine_facedet)\n",
    "engine_lmk     = get_engine(dir_engine_lmk)\n",
    "\n",
    "#Allocate buffers and create a stream\n",
    "stream_dexi    = cuda.Stream()\n",
    "stream_facedet = cuda.Stream()\n",
    "stream_lmk     = cuda.Stream()\n",
    "\n",
    "#Get locations\n",
    "in_dexi, out_dexi, bind_dexi          = allocate_buffers_um_pinned(engine_dexi)\n",
    "in_facedet, out_facedet, bind_facedet = allocate_buffers_um_pinned(engine_facedet)\n",
    "in_lmk, out_lmk, bind_lmk             = allocate_buffers_um_pinned(engine_lmk)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with engine_dexi.create_execution_context() as context_dexi, \\\n",
    "     engine_facedet.create_execution_context() as context_facedet, \\\n",
    "     engine_lmk.create_execution_context() as context_lmk:\n",
    "\n",
    "    while pipe_in.isOpened():\n",
    "\n",
    "        ret, frame = pipe_in.read()\n",
    "        if not ret: break\n",
    "        \n",
    "        np.copyto(in_dexi[0], dexi_pre(frame))\n",
    "        np.copyto(in_facedet[0], facenet_pre(frame))\n",
    "\n",
    "        context_dexi.execute_async_v2(bindings=bind_dexi, stream_handle=stream_dexi.handle)\n",
    "        context_facedet.execute_async_v2(bindings=bind_facedet, stream_handle=stream_facedet.handle)\n",
    "\n",
    "        stream_dexi.synchronize()\n",
    "        stream_facedet.synchronize()\n",
    "\n",
    "        dexi_frame = dexi_post(out_dexi[0])\n",
    "        bboxs      = facedet_post(out_facedet)\n",
    "\n",
    "        for det in (bboxs if bboxs is not None else []):\n",
    "            \n",
    "            bbox = det[0:4].astype(np.int32)\n",
    "            \n",
    "            # Scale the bounding box to the new size\n",
    "            #if bbox[0] < 0: bbox[0] = 0\n",
    "            x_new = int(bbox[0] * scale_w)\n",
    "            y_new = int(bbox[1] * scale_h)\n",
    "            w_new = int(bbox[2] * scale_w)\n",
    "            h_new = int(bbox[3] * scale_h)\n",
    "\n",
    "            # Draw the bounding box on the larger image\n",
    "            #cv2.rectangle(dexi_frame, (x_new, y_new), (x_new + w_new, y_new + h_new), (255, 0, 0), 2)\n",
    "\n",
    "            #Crop out face\n",
    "            face = frame[y_new:y_new+h_new, x_new:x_new+w_new]\n",
    "            if 0 in face.shape: continue\n",
    "            face = cv2.resize(face, (80,80))\n",
    "            face = cv2.cvtColor(face, cv2.COLOR_BGR2GRAY)\n",
    "\n",
    "            np.copyto(in_lmk[0], face.ravel())\n",
    "            context_lmk.execute_async_v2(bindings=bind_lmk, stream_handle=stream_lmk.handle)\n",
    "            stream_lmk.synchronize()\n",
    "\n",
    "            # Denormalize the landmarks to the 120x160 size (bbox space)\n",
    "            landmarks_denorm = out_lmk[1].reshape((80,2)) / 80 * [w_new, h_new] + [x_new, y_new] \n",
    "\n",
    "            #### Draw on edge map ###\n",
    "\n",
    "            # Now, landmarks_rescaled_to_512_512 contains the coordinates of the landmarks in the 512x512 image.\n",
    "            for (x, y) in landmarks_denorm[lmks[0]]: \n",
    "                \n",
    "                cv2.circle(dexi_frame, (int(x), int(y)), 2, (255, 0, 0), -1)\n",
    "\n",
    "\n",
    "        dexi_frame = dexi_frame.astype(np.uint8)\n",
    "    \n",
    "        pipe_out.write(dexi_frame)\n",
    "        root.update() \n",
    "    \n",
    "pipe_in.release()\n",
    "pipe_out.release()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}