tutorials/mouth_cropping.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "556076dd-8537-4356-8ac3-f4a838b2e896",
   "metadata": {},
   "source": [
    "# MOUTH CROPPING"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6b21296a-9b12-40c3-b1fc-c3e316120075",
   "metadata": {},
   "source": [
    "This tutorial shows how to use our `AVSRDataLoader` to extract cropped mouth."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5e067f7d-cc60-43a0-b65a-7272c9cecc81",
   "metadata": {},
   "source": [
    "**Note** This tutorial requires `mediapipe` or `retinaface` detector. Please refer to [preparation](../preparation#setup) for installation."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c97575fc-bb77-43d0-8eb7-f94baf4cdc89",
   "metadata": {},
   "source": [
    "**Note** To run this tutorial, please make sure you are in tutorials folder."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa81fb4f-6e89-4238-9f47-3f66eedd3c22",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "sys.path.insert(0, \"../\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d2fb0de3-af7f-477e-baef-dd595effeef8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import torch\n",
    "import torchvision\n",
    "\n",
    "from preparation.data.data_module import AVSRDataLoader"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2875bb66-4e96-4c51-85d6-a129748d197a",
   "metadata": {},
   "source": [
    "## 1. Initialize AVSR Dataloader.\n",
    "    \n",
    "The AVSRDataLoader class is a pre-process pipeline from raw video to cropped mouth, which includes face detection, landmarks tracking, face transformation and mouth cropping. We provide `mediapipe` and `retinaface` face detector, `mediapipe` is faster with relatively lower precision compared with `retinaface`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ac03a1e1-3082-43ca-bb44-e58f96464bc3",
   "metadata": {},
   "outputs": [],
   "source": [
    "class AVSRDataLoader(torch.nn.Module):\n",
    "    def __init__(self, detector=\"retinaface\"):\n",
    "        super().__init__()\n",
    "        if detector == \"mediapipe\":\n",
    "            from preparation.detectors.mediapipe.detector import LandmarksDetector\n",
    "            from preparation.detectors.mediapipe.video_process import VideoProcess\n",
    "            self.landmarks_detector = LandmarksDetector()\n",
    "            self.video_process = VideoProcess(convert_gray=False)\n",
    "        elif detector == \"retinaface\":\n",
    "            from preparation.detectors.retinaface.detector import LandmarksDetector\n",
    "            from preparation.detectors.retinaface.video_process import VideoProcess\n",
    "            self.landmarks_detector = LandmarksDetector(device=\"cuda:0\")\n",
    "            self.video_process = VideoProcess(convert_gray=False)\n",
    "\n",
    "    def forward(self, data_filename):\n",
    "        video = self.load_video(data_filename)\n",
    "        landmarks = self.landmarks_detector(video)\n",
    "        video = self.video_process(video, landmarks)\n",
    "        video = torch.tensor(video)\n",
    "        return video\n",
    "\n",
    "    def load_video(self, data_filename):\n",
    "        return torchvision.io.read_video(data_filename, pts_unit=\"sec\")[0].numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4fa2b52b-8dd6-42ea-a4cc-bc8c30253fa0",
   "metadata": {},
   "outputs": [],
   "source": [
    "video_dataloader = AVSRDataLoader(detector=\"retinaface\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7dfa0be0-788d-42b4-8873-0cd7ce804eb3",
   "metadata": {},
   "source": [
    "## 2. Download a video from web"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ea438ae5-2c1b-420b-a7f2-0a23e3c22ed3",
   "metadata": {},
   "outputs": [],
   "source": [
    "!wget --content-disposition http://www.doc.ic.ac.uk/~pm4115/autoAVSR/autoavsr_demo_video.mp4 -O ./input.mp4\n",
    "data_filename = \"input.mp4\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "326d12c2-c051-4e5b-b39b-9c24a2e9fd73",
   "metadata": {},
   "outputs": [],
   "source": [
    "preprocessed_video = video_dataloader(data_filename)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7a302c17-78f8-49b6-af40-6210c83193ef",
   "metadata": {},
   "source": [
    "## 3. Saved the cropped video"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "401c84d6-5910-4959-8761-34cbd1d18c56",
   "metadata": {},
   "outputs": [],
   "source": [
    "def save2vid(filename, vid, frames_per_second):\n",
    "    os.makedirs(os.path.dirname(filename), exist_ok=True)\n",
    "    torchvision.io.write_video(filename, vid, frames_per_second)\n",
    "save2vid(\"./output.mp4\", preprocessed_video, frames_per_second=30)"
   ]
  }
 ],
 "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.8.17"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}