get_face_area.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cv2\n",
    "import numpy as np\n",
    "import os\n",
    "import csv\n",
    "import pandas as pd\n",
    "from batch_face import RetinaFace"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "class VideoCamera(object):\n",
    "    def __init__(self, path_video='', path_save_images='', label_df='', label_true='', conf=0.7):\n",
    "        self.path_video = path_video\n",
    "        self.path_save_images = path_save_images\n",
    "        self.label_df = label_df\n",
    "        self.label_true = label_true\n",
    "        self.conf = conf # порог доставерности\n",
    "        self.count_pie = 0 # подсчитываем количество кадров\n",
    "        self.video = None\n",
    "        self.h = None\n",
    "        self.w = None\n",
    "\n",
    "    def __del__(self):\n",
    "        self.video.release()\n",
    "        \n",
    "    def get_h_w(self):\n",
    "        try:\n",
    "            (self.h, self.w) = self.fr.shape[:2]\n",
    "        except AttributeError:\n",
    "            print('The end')\n",
    "            return False\n",
    "        else: return True\n",
    "            \n",
    "    def get_frame(self):\n",
    "        self.video = cv2.VideoCapture(self.path_video)\n",
    "        total_frame = int(self.video.get(cv2.CAP_PROP_FRAME_COUNT))\n",
    "        fps = np.round(self.video.get(cv2.CAP_PROP_FPS))\n",
    "        print('Number of frames: ', total_frame)\n",
    "        print('FPS: ', fps)\n",
    "        while True:\n",
    "            _, self.fr = self.video.read()\n",
    "            if self.get_h_w() is False: return None\n",
    "            self.count_pie += 1\n",
    "            name_img = str(self.count_pie).zfill(6)\n",
    "            faces = detector(self.fr, cv=False)\n",
    "            for f_id, box in enumerate(faces):\n",
    "                box, _, prob = box\n",
    "                if prob > self.conf:\n",
    "                    startX = int(box[0])\n",
    "                    startY = int(box[1])\n",
    "                    endX = int(box[2])\n",
    "                    endY = int(box[3])\n",
    "                    (startX, startY) = (max(0, startX), max(0, startY))\n",
    "                    (endX, endY) = (min(self.w - 1, endX), min(self.h - 1, endY))\n",
    "                    cur_fr = self.fr[startY: endY, startX: endX]\n",
    "                    \n",
    "                    if not os.path.exists(self.path_save_images + str(f_id)):\n",
    "                        os.makedirs(self.path_save_images + str(f_id))\n",
    "                        df = pd.DataFrame(columns=['name_folder','path_images','emotion'])\n",
    "                        df.to_csv(self.path_save_images + 'df_test.csv', index=False)\n",
    "                            \n",
    "                    cv2.imwrite(self.path_save_images + str(f_id) + \"/{}.jpg\".format(name_img), cur_fr)\n",
    "                    \n",
    "                    with open(self.path_save_images + 'df_test.csv', 'a+', newline='') as csvfile:\n",
    "                        writer = csv.writer(csvfile, delimiter=',')\n",
    "                        emo = self.label_true[self.label_df[os.path.basename(self.path_video).split('-')[2]]]\n",
    "                        writer.writerow([os.path.basename(self.path_video).split('.')[0], self.path_save_images + str(f_id) + \"/{}.jpg\".format(name_img), emo])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import warnings\n",
    "warnings.filterwarnings('ignore', category = UserWarning)\n",
    "detector = RetinaFace(gpu_id=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of frames:  119\n",
      "FPS:  30.0\n",
      "The end\n",
      "Wall time: 7.99 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "label_ravdess = {'01':'Neutral', '02':'Calm', '03':'Happiness', '04':'Sadness', '05':'Anger', '06':'Fear', '07':'Disgust', '08':'Surprise'}\n",
    "label_model = {'Neutral':0, 'Happiness':1, 'Sadness':2, 'Surprise':3, 'Fear':4, 'Disgust':5, 'Anger':6}\n",
    "\n",
    "detect = VideoCamera(path_video='video/01-01-07-02-02-02-06.mp4', path_save_images='img/01-01-07-02-02-02-06/', label_df=label_ravdess, label_true=label_model, conf=0.7)\n",
    "detect.get_frame()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}