demo.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "87122079",
   "metadata": {},
   "source": [
    "### install face_alignment (https://github.com/1adrianb/face-alignment), we use it to do face detection and landmark detection here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "44071ea0",
   "metadata": {},
   "outputs": [],
   "source": [
    "import cv2\n",
    "import face_alignment\n",
    "import numpy as np\n",
    "from PIL import Image\n",
    "from skimage import io\n",
    "from skimage import transform as trans\n",
    "from torchvision import transforms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "11682e8b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_scale_center(bb, scale_=220.0):\n",
    "    center = np.array([bb[2] - (bb[2] - bb[0]) / 2,\n",
    "                      bb[3] - (bb[3] - bb[1]) / 2])\n",
    "    scale = (bb[2] - bb[0] + bb[3] - bb[1]) / scale_\n",
    "\n",
    "    return scale, center"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "80a4e650",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_transform(center, scale, res, rot=0):\n",
    "    # Generate transformation matrix\n",
    "    h = 200 * scale\n",
    "    t = np.zeros((3, 3))\n",
    "    t[0, 0] = float(res[1]) / h\n",
    "    t[1, 1] = float(res[0]) / h\n",
    "    t[0, 2] = res[1] * (-float(center[0]) / h + .5)\n",
    "    t[1, 2] = res[0] * (-float(center[1]) / h + .5)\n",
    "    t[2, 2] = 1\n",
    "\n",
    "    if not rot == 0:\n",
    "        rot = -rot  # To match direction of rotation from cropping\n",
    "        rot_mat = np.zeros((3, 3))\n",
    "        rot_rad = rot * np.pi / 180\n",
    "        sn, cs = np.sin(rot_rad), np.cos(rot_rad)\n",
    "        rot_mat[0, :2] = [cs, -sn]\n",
    "        rot_mat[1, :2] = [sn, cs]\n",
    "        rot_mat[2, 2] = 1\n",
    "        # Need to rotate around center\n",
    "        t_mat = np.eye(3)\n",
    "        t_mat[0, 2] = -res[1] / 2\n",
    "        t_mat[1, 2] = -res[0] / 2\n",
    "        t_inv = t_mat.copy()\n",
    "        t_inv[:2, 2] *= -1\n",
    "        t = np.dot(t_inv, np.dot(rot_mat, np.dot(t_mat, t)))\n",
    "    return t"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "63af7368",
   "metadata": {},
   "outputs": [],
   "source": [
    "def crop(img,rects):\n",
    "    im_w = 256\n",
    "    bb = rects[:4]\n",
    "    scale, center = get_scale_center(bb, scale_=260)\n",
    "    aug_rot = 0\n",
    "    dx, dy = 0, 0\n",
    "    center[0] += dx * center[0]\n",
    "    center[1] += dy * center[1]\n",
    "    mat = get_transform(center, scale, (im_w, im_w), aug_rot)[:2]\n",
    "    img = cv2.warpAffine(img.copy(), mat, (im_w, im_w))\n",
    "    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n",
    "    return bb,img"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "202b41cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# video to img\n",
    "def video2img(src,dst)\n",
    "    import os\n",
    "    vid = cv2.VideoCapture(src)\n",
    "    save_folder = dst\n",
    "    os.makedirs(save_folder,exist_ok=True)\n",
    "    index = 0\n",
    "    while True:\n",
    "        _, frame = vid.read()\n",
    "        if frame is None:\n",
    "            break\n",
    "        else:\n",
    "            cv2.imwrite(f'{save_folder}/%06d.jpg'%index,frame)\n",
    "            index = index+1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "a3598e44",
   "metadata": {},
   "outputs": [],
   "source": [
    "def caculate_pts5(pred):\n",
    "    eye_left = np.mean(pred[36:42,:],axis=0)\n",
    "    eye_right = np.mean(pred[42:48,:],axis=0)\n",
    "    return np.array([eye_left,eye_right,pred[33,:],pred[48,:],pred[54,:]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 78,
   "id": "3e61b0ad",
   "metadata": {},
   "outputs": [],
   "source": [
    "def AU_detection(model,src,dst):\n",
    "    import os\n",
    "    import matplotlib.pyplot as plt\n",
    "    tform = trans.SimilarityTransform()\n",
    "    sr_folder = './video_2/'\n",
    "    save_folder = './res_video_2/' \n",
    "    os.makedirs(save_folder,exist_ok=True)\n",
    "    au_indices = (1,2,4,6,7,9,10,12,14,15,17,23,24,25,26)\n",
    "    fa = face_alignment.FaceAlignment(face_alignment.LandmarksType._2D, face_detector='sfd')\n",
    "    for idx in range(664) :\n",
    "        # Detect face\n",
    "        try:\n",
    "            im_0 = cv2.imread(f'./{sr_folder}/%06d.jpg'%idx-10)\n",
    "        except:\n",
    "            im_0 = cv2.imread(f'./{sr_folder}/%06d.jpg'%idx)\n",
    "\n",
    "\n",
    "        im_1 = cv2.imread(f'./{sr_folder}/%06d.jpg'%idx)\n",
    "        frame = im_1.copy()\n",
    "\n",
    "        try:\n",
    "            im_2 = cv2.imread(f'./{sr_folder}/%06d.jpg'%idx+10)\n",
    "        except:\n",
    "            im_2 = cv2.imread(f'./{sr_folder}/%06d.jpg'%idx)\n",
    "\n",
    "        #detect landmarks \n",
    "        lmks_0 = fa.get_landmarks(im_0)\n",
    "        lmks_0 = caculate_pts5(lmks_0[0])\n",
    "\n",
    "        #detect landmarks and bounding box\n",
    "        lmks_1, _, rects_1 = fa.get_landmarks(im_1,return_bboxes=True)\n",
    "        lmks_1 = caculate_pts5(lmks_1[0])\n",
    "\n",
    "        #detect landmarks\n",
    "        lmks_2 = fa.get_landmarks(im_2)\n",
    "        lmks_2 = caculate_pts5(lmks_2[0])\n",
    "\n",
    "        #crop the face area based on bounding box\n",
    "        bb = rects_1[0][:4]\n",
    "        bbox = bb.copy()\n",
    "        scale, center = get_scale_center(bb, scale_=260)\n",
    "        aug_rot = 0\n",
    "        dx, dy = 0,0\n",
    "        im_w = 256\n",
    "        center[0] += dx * center[0]\n",
    "        center[1] += dy * center[1]\n",
    "        mat = get_transform(center, scale, (im_w, im_w), aug_rot)[:2]\n",
    "        im_1 = cv2.warpAffine(im_1.copy(), mat, (im_w, im_w))\n",
    "        im_1 = cv2.cvtColor(im_1, cv2.COLOR_BGR2RGB)\n",
    "\n",
    "        # calculate the transformed landmarks\n",
    "        lmks_1 = np.dot(np.concatenate((lmks_1, lmks_1[:, 0:1]*0+1), axis=1), mat.T)\n",
    "\n",
    "        im_1 = Image.fromarray(im_1)\n",
    "        im_1 = transform_val(im_1)\n",
    "        im_1 = im_1.unsqueeze(0)\n",
    "        im_1 = im_1.cuda()\n",
    "\n",
    "        #align im_0 to im_1\n",
    "        tform.estimate(lmks_0, lmks_1)\n",
    "        M = tform.params[0:2, :]\n",
    "        im_0 = cv2.warpAffine(im_0.copy(), M, (im_w, im_w), borderValue=0.0)\n",
    "        im_0 = cv2.cvtColor(im_0, cv2.COLOR_BGR2RGB)\n",
    "        im_0 = Image.fromarray(im_0)\n",
    "        im_0 = transform_val(im_0)\n",
    "        im_0 = im_0.unsqueeze(0)\n",
    "        im_0 = im_0.cuda()\n",
    "\n",
    "        #align im_2 to im_1\n",
    "        tform.estimate(lmks_2, lmks_1)\n",
    "        M = tform.params[0:2, :]\n",
    "        im_2 = cv2.warpAffine(im_2.copy(), M, (im_w, im_w), borderValue=0.0)\n",
    "        im_2 = cv2.cvtColor(im_2, cv2.COLOR_BGR2RGB)\n",
    "\n",
    "\n",
    "        im_2 = Image.fromarray(im_2)\n",
    "        im_2 = transform_val(im_2)\n",
    "        im_2 = im_2.unsqueeze(0)\n",
    "        im_2 = im_2.cuda()\n",
    "\n",
    "        #AU detection\n",
    "        with torch.no_grad():\n",
    "            pred = model(im_1,im_0,im_2)\n",
    "\n",
    "        probs = list(np.array(torch.sigmoid(pred[0]).cpu().data))   \n",
    "        bbox = np.array(bbox,dtype=int)\n",
    "        cv2.rectangle(frame, (bbox[0], bbox[1]), (bbox[2], bbox[3]), color=(0, 0, 255), thickness=2)\n",
    "        for loc, (au_idx, au_prob) in enumerate(zip(au_indices, probs)):\n",
    "            colour = (0, round(255 * au_prob), round(255 * (1 - au_prob)))\n",
    "            cv2.circle(frame, (bbox[2] + 10, bbox[1] + 15 * loc + 5),\n",
    "                       radius=5, thickness=-1, color=colour, lineType=cv2.LINE_AA)\n",
    "            cv2.putText(frame, f'AU {au_idx}', (bbox[2] + 20, bbox[1] + 15 * loc + 10),\n",
    "                    cv2.FONT_HERSHEY_DUPLEX, 0.45, colour, lineType=cv2.LINE_AA)\n",
    "        cv2.imwrite(f'./{save_folder}/%06d.png'%idx,frame)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "eea32022",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(1280, 720, 3)\n",
      "Moviepy - Building video ./res_video2.mp4.\n",
      "Moviepy - Writing video ./res_video2.mp4\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                              \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Moviepy - Done !\n",
      "Moviepy - video ready ./res_video2.mp4\n"
     ]
    }
   ],
   "source": [
    "def img2video(src):\n",
    "    import moviepy\n",
    "    import glob\n",
    "    import moviepy.video.io.ImageSequenceClip\n",
    "    im_list = glob.glob(f'{src}/*.png')\n",
    "    im_list.sort()\n",
    "    fps = 10\n",
    "    clip = moviepy.video.io.ImageSequenceClip.ImageSequenceClip(im_list, fps=fps)\n",
    "    print(clip.get_frame(3).shape)\n",
    "    clip.write_videofile('./res_video2.mp4')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "34e2a065",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from models.aunet import AU_NET\n",
    "transform_val = transforms.Compose([transforms.ToTensor()])\n",
    "model = AU_NET(alpha=0.9, beta=0.1, n_classes=15)\n",
    "pre_trained = torch.load('./cross_model.pth', 'cpu')\n",
    "pretrained_items = list(pre_trained.items())\n",
    "current_items = model.predictor.state_dict()\n",
    "count = 0\n",
    "for key, value in current_items.items():\n",
    "    layer_name, weights = pretrained_items[count]\n",
    "    current_items[key] = weights\n",
    "    count = count + 1\n",
    "model.predictor.load_state_dict(current_items, strict=True)\n",
    "model = model.cuda()\n",
    "model = model.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "312898bd",
   "metadata": {},
   "outputs": [],
   "source": [
    "src = 'demo.mp4'\n",
    "dst = './demos_imgs/'\n",
    "video2img(src,dst)\n",
    "res = './demos_res'\n",
    "AU_detection(model,dst,res)\n",
    "img2video(res)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.7.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}