check_backbone_models_by_webcam.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "06ff260a-e226-45ab-830c-8152149f3d35",
   "metadata": {},
   "outputs": [],
   "source": [
    "import cv2\n",
    "import mediapipe as mp #face detector\n",
    "import math\n",
    "import numpy as np\n",
    "import time\n",
    "\n",
    "# torch\n",
    "import torch\n",
    "from PIL import Image\n",
    "from torchvision import transforms\n",
    "\n",
    "# tf\n",
    "# import tensorflow as tf"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "06982039-7586-4bd5-bebe-e17cf44349e5",
   "metadata": {},
   "source": [
    "#### Sub functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "3f8b869f-7c10-4b12-96fc-94e8b35953b0",
   "metadata": {},
   "outputs": [],
   "source": [
    "def pth_processing(fp):\n",
    "    class PreprocessInput(torch.nn.Module):\n",
    "        def init(self):\n",
    "            super(PreprocessInput, self).init()\n",
    "\n",
    "        def forward(self, x):\n",
    "            x = x.to(torch.float32)\n",
    "            x = torch.flip(x, dims=(0,))\n",
    "            x[0, :, :] -= 91.4953\n",
    "            x[1, :, :] -= 103.8827\n",
    "            x[2, :, :] -= 131.0912\n",
    "            return x\n",
    "\n",
    "    def get_img_torch(img):\n",
    "        \n",
    "        ttransform = transforms.Compose([\n",
    "            transforms.PILToTensor(),\n",
    "            PreprocessInput()\n",
    "        ])\n",
    "        img = img.resize((224, 224), Image.Resampling.NEAREST)\n",
    "        img = ttransform(img)\n",
    "        img = torch.unsqueeze(img, 0).to('cuda')\n",
    "        return img\n",
    "    return get_img_torch(fp)\n",
    "\n",
    "def tf_processing(fp):\n",
    "    def preprocess_input(x):\n",
    "        x_temp = np.copy(x)\n",
    "        x_temp = x_temp[..., ::-1]\n",
    "        x_temp[..., 0] -= 91.4953\n",
    "        x_temp[..., 1] -= 103.8827\n",
    "        x_temp[..., 2] -= 131.0912\n",
    "        return x_temp\n",
    "\n",
    "    def get_img_tf(img):\n",
    "        img = cv2.resize(img, (224,224), interpolation=cv2.INTER_NEAREST)\n",
    "        img = tf.keras.utils.img_to_array(img)\n",
    "        img = preprocess_input(img)\n",
    "        img = np.array([img])\n",
    "        return img\n",
    "\n",
    "    return get_img_tf(fp)\n",
    "\n",
    "def norm_coordinates(normalized_x, normalized_y, image_width, image_height):\n",
    "    \n",
    "    x_px = min(math.floor(normalized_x * image_width), image_width - 1)\n",
    "    y_px = min(math.floor(normalized_y * image_height), image_height - 1)\n",
    "    \n",
    "    return x_px, y_px\n",
    "\n",
    "def get_box(fl, w, h):\n",
    "    idx_to_coors = {}\n",
    "    for idx, landmark in enumerate(fl.landmark):\n",
    "        landmark_px = norm_coordinates(landmark.x, landmark.y, w, h)\n",
    "\n",
    "        if landmark_px:\n",
    "            idx_to_coors[idx] = landmark_px\n",
    "\n",
    "    x_min = np.min(np.asarray(list(idx_to_coors.values()))[:,0])\n",
    "    y_min = np.min(np.asarray(list(idx_to_coors.values()))[:,1])\n",
    "    endX = np.max(np.asarray(list(idx_to_coors.values()))[:,0])\n",
    "    endY = np.max(np.asarray(list(idx_to_coors.values()))[:,1])\n",
    "\n",
    "    (startX, startY) = (max(0, x_min), max(0, y_min))\n",
    "    (endX, endY) = (min(w - 1, endX), min(h - 1, endY))\n",
    "    \n",
    "    return startX, startY, endX, endY\n",
    "\n",
    "def display_EMO_PRED(img, box, label='', color=(128, 128, 128), txt_color=(255, 255, 255), line_width=2, ):\n",
    "    lw = line_width or max(round(sum(img.shape) / 2 * 0.003), 2)\n",
    "    text2_color = (255, 0, 255)\n",
    "    p1, p2 = (int(box[0]), int(box[1])), (int(box[2]), int(box[3]))\n",
    "    cv2.rectangle(img, p1, p2, text2_color, thickness=lw, lineType=cv2.LINE_AA)\n",
    "    font = cv2.FONT_HERSHEY_SIMPLEX\n",
    "\n",
    "    tf = max(lw - 1, 1)\n",
    "    text_fond = (0, 0, 0)\n",
    "    text_width_2, text_height_2 = cv2.getTextSize(label, font, lw / 3, tf)\n",
    "    text_width_2 = text_width_2[0] + round(((p2[0] - p1[0]) * 10) / 360)\n",
    "    center_face = p1[0] + round((p2[0] - p1[0]) / 2)\n",
    "\n",
    "    cv2.putText(img, label,\n",
    "                (center_face - round(text_width_2 / 2), p1[1] - round(((p2[0] - p1[0]) * 20) / 360)), font,\n",
    "                lw / 3, text_fond, thickness=tf, lineType=cv2.LINE_AA)\n",
    "    cv2.putText(img, label,\n",
    "                (center_face - round(text_width_2 / 2), p1[1] - round(((p2[0] - p1[0]) * 20) / 360)), font,\n",
    "                lw / 3, text2_color, thickness=tf, lineType=cv2.LINE_AA)\n",
    "    return img\n",
    "\n",
    "def display_FPS(img, text, margin=1.0, box_scale=1.0):\n",
    "    img_h, img_w, _ = img.shape\n",
    "    line_width = int(min(img_h, img_w) * 0.001)  # line width\n",
    "    thickness = max(int(line_width / 3), 1)  # font thickness\n",
    "\n",
    "    font_face = cv2.FONT_HERSHEY_SIMPLEX\n",
    "    font_color = (0, 0, 0)\n",
    "    font_scale = thickness / 1.5\n",
    "\n",
    "    t_w, t_h = cv2.getTextSize(text, font_face, font_scale, None)[0]\n",
    "\n",
    "    margin_n = int(t_h * margin)\n",
    "    sub_img = img[0 + margin_n: 0 + margin_n + t_h + int(2 * t_h * box_scale),\n",
    "              img_w - t_w - margin_n - int(2 * t_h * box_scale): img_w - margin_n]\n",
    "\n",
    "    white_rect = np.ones(sub_img.shape, dtype=np.uint8) * 255\n",
    "\n",
    "    img[0 + margin_n: 0 + margin_n + t_h + int(2 * t_h * box_scale),\n",
    "    img_w - t_w - margin_n - int(2 * t_h * box_scale):img_w - margin_n] = cv2.addWeighted(sub_img, 0.5, white_rect, .5,\n",
    "                                                                                          1.0)\n",
    "\n",
    "    cv2.putText(img=img,\n",
    "                text=text,\n",
    "                org=(img_w - t_w - margin_n - int(2 * t_h * box_scale) // 2,\n",
    "                     0 + margin_n + t_h + int(2 * t_h * box_scale) // 2),\n",
    "                fontFace=font_face,\n",
    "                fontScale=font_scale,\n",
    "                color=font_color,\n",
    "                thickness=thickness,\n",
    "                lineType=cv2.LINE_AA,\n",
    "                bottomLeftOrigin=False)\n",
    "\n",
    "    return img"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e6c5eb65-7e6b-4919-8354-c0594f1ddd8f",
   "metadata": {},
   "source": [
    "#### Testing models by webcam"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "8fffbfbd-5f4e-4fa0-8c67-4a8e71c47f29",
   "metadata": {},
   "outputs": [],
   "source": [
    "mp_face_mesh = mp.solutions.face_mesh\n",
    "\n",
    "name = '0_66_49_wo_gl'\n",
    "\n",
    "# torch\n",
    "pth_model = torch.jit.load('models_EmoAffectnet/torchscript_model_{0}.pth'.format(name)).to('cuda')\n",
    "pth_model.eval()\n",
    "\n",
    "\n",
    "\n",
    "# tf\n",
    "# tf_model = tf.keras.models.load_model('models_EmoAffectnet/weights_{0}.h5'.format(name))\n",
    "\n",
    "DICT_EMO = {0: 'Neutral', 1: 'Happiness', 2: 'Sadness', 3: 'Surprise', 4: 'Fear', 5: 'Disgust', 6: 'Anger'}\n",
    "\n",
    "cap = cv2.VideoCapture(0)\n",
    "w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))\n",
    "h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))\n",
    "fps = np.round(cap.get(cv2.CAP_PROP_FPS))\n",
    "\n",
    "path_save_video = 'result.mp4'\n",
    "vid_writer = cv2.VideoWriter(path_save_video, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))\n",
    "    \n",
    "with mp_face_mesh.FaceMesh(\n",
    "max_num_faces=1,\n",
    "refine_landmarks=False,\n",
    "min_detection_confidence=0.5,\n",
    "min_tracking_confidence=0.5) as face_mesh:\n",
    "\n",
    "    while cap.isOpened():\n",
    "        t1 = time.time()\n",
    "        success, frame = cap.read()\n",
    "        if frame is None: break\n",
    "\n",
    "        frame_copy = frame.copy()\n",
    "        frame_copy.flags.writeable = False\n",
    "        frame_copy = cv2.cvtColor(frame_copy, cv2.COLOR_BGR2RGB)\n",
    "        results = face_mesh.process(frame_copy)\n",
    "        frame_copy.flags.writeable = True\n",
    "        \n",
    "        if results.multi_face_landmarks:\n",
    "            for fl in results.multi_face_landmarks:\n",
    "                startX, startY, endX, endY  = get_box(fl, w, h)\n",
    "                cur_face = frame_copy[startY:endY, startX: endX]\n",
    "\n",
    "                # torch\n",
    "                cur_face = pth_processing(Image.fromarray(cur_face))\n",
    "                output = torch.nn.functional.softmax(pth_model(cur_face), dim=1).cpu().detach().numpy()\n",
    "                \n",
    "                # tf\n",
    "                # cur_face = tf_processing(cur_face)\n",
    "                # output = tf_model(cur_face, training=False).numpy()\n",
    "                \n",
    "                cl = np.argmax(output)\n",
    "                label = DICT_EMO[cl]\n",
    "                frame = display_EMO_PRED(frame, (startX, startY, endX, endY), label, line_width=3)\n",
    "\n",
    "        t2 = time.time()\n",
    "\n",
    "        frame = display_FPS(frame, 'FPS: {0:.1f}'.format(1 / (t2 - t1)), box_scale=.5)\n",
    "\n",
    "        vid_writer.write(frame)\n",
    "        \n",
    "        cv2.imshow('Webcam', frame)\n",
    "        if cv2.waitKey(1) & 0xFF == ord('q'):\n",
    "            break\n",
    "\n",
    "    vid_writer.release()\n",
    "    cap.release()\n",
    "    cv2.destroyAllWindows()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "1628e231-20c0-4f33-86ab-bc9f147542c7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MoviePy - Building file result.gif with imageio.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                                                       \r"
     ]
    }
   ],
   "source": [
    "from moviepy.editor import VideoFileClip\n",
    "\n",
    "def convert_mp4_to_gif(input_path, output_path, fps=10):\n",
    "    clip = VideoFileClip(input_path)\n",
    "    clip.write_gif(output_path, fps=fps)\n",
    "\n",
    "# Пример использования\n",
    "input_video_path = \"result.mp4\"\n",
    "output_gif_path = \"result.gif\"\n",
    "\n",
    "convert_mp4_to_gif(input_video_path, output_gif_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3d3b8ca4-7aea-41c1-80eb-d2d437c4ee71",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}