Danmu-MaskCNN.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 准备工作\n",
    "\n",
    "#### 1. 用 https://www.kanbilibili.com/ 来下载 B站视频和弹幕\n",
    "例如：https://www.kanbilibili.com/video/av54944922?from=search&seid=13855581700550241007#download\n",
    "\n",
    "#### 2. 然后用 DanmakuFactory_1.11_release 来将 xml 文件转化为 ass 字幕文件\n",
    "\n",
    "#### 3. 使用格式工厂将视频和ass进行合成，产生原视频和字幕视频\n",
    "\n",
    "#### 4. 下载 https://github.com/matterport/Mask_RCNN 源代码，在sample目录下运行该notebook\n",
    "\n",
    "#### 5. 修改配置文件"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import logging\n",
    "import os\n",
    "import sys\n",
    "from time import time as timer\n",
    "import cv2\n",
    "import numpy as np\n",
    "\n",
    "logging.basicConfig(level=logging.INFO)\n",
    "logger = logging.getLogger(__name__)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_video_shape(originVideo, danmuVideo):\n",
    "    '''\n",
    "    检查原视频和弹幕视频的shape是否相同\n",
    "    :param danmu:\n",
    "    :param origin:\n",
    "    :return: shape\n",
    "    '''\n",
    "    # 读取一帧的视频\n",
    "    _, frame_ori = originVideo.read()\n",
    "    _, frame_dan = danmuVideo.read()\n",
    "    assert frame_dan.shape == frame_ori.shape, \\\n",
    "        \"The height and width are not the same! \\n {} with {}\".format(frame_ori.shape, frame_dan.shape)\n",
    "    return frame_dan.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_video_frame(danmuVideo, originVideo):\n",
    "    '''\n",
    "    检查视频帧率\n",
    "    :param danmu:\n",
    "    :param origin:\n",
    "    :return:\n",
    "    '''\n",
    "    if originVideo.get(cv2.CAP_PROP_FPS) != danmuVideo.get(cv2.CAP_PROP_FPS):\n",
    "        logger.warning('frame not same!')\n",
    "        logger.warning('origin:{}'.format(originVideo.get(cv2.CAP_PROP_FPS)))\n",
    "        logger.warning('danmu:{}'.format(danmuVideo.get(cv2.CAP_PROP_FPS)))\n",
    "    return originVideo.get(cv2.CAP_PROP_FPS)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def set_mask(originVideo, danmuVideo, mask):\n",
    "    '''\n",
    "    将 原帧的 蒙版区域放到 弹幕视频帧 之上\n",
    "    :param origin:\n",
    "    :param danmu:\n",
    "    :param mask:\n",
    "    :return:\n",
    "    '''\n",
    "    for c in range(3):\n",
    "        danmuVideo[:, :, c] = np.where(mask == 1,\n",
    "                                  originVideo[:, :, c],\n",
    "                                  danmuVideo[:, :, c])\n",
    "    return danmuVideo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def display_instances(originVideo, danmuVideo, boxes, masks, ids, scores):\n",
    "    '''\n",
    "    返回处理蒙版弹幕后的图片\n",
    "    :param origin: 原始图片\n",
    "    :param danmu: 带弹幕的图片\n",
    "    :param boxes: 框, [num_instances, ... ...]\n",
    "    :param masks: 掩膜，[h, w, num_instances]\n",
    "    :param ids: 各个mask所属的class id, [num_instances]\n",
    "    :param scores:  置信度\n",
    "    :return:\n",
    "    '''\n",
    "    n_instances = boxes.shape[0]\n",
    "    h, w, _ = originVideo.shape\n",
    "\n",
    "    if not n_instances:\n",
    "        print('NO INSTANCES TO DISPLAY')\n",
    "    else:\n",
    "        assert boxes.shape[0] == masks.shape[-1] == ids.shape[0]\n",
    "\n",
    "    for i in range(n_instances):\n",
    "        if not np.any(boxes[i]):\n",
    "            continue\n",
    "        if ids[i] != 1:\n",
    "        \t# 1：person\n",
    "            continue\n",
    "        mask = masks[:, :, i]\n",
    "\n",
    "        danmuVideo = set_mask(originVideo, danmuVideo, mask)\n",
    "\n",
    "    return danmuVideo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def processVideo(model, params):\n",
    "    '''\n",
    "    返回人物防遮挡蒙版弹幕视频\n",
    "    :param model: Mask-RCNN模型\n",
    "    :param params: 配置dict\n",
    "    :return:\n",
    "    '''\n",
    "    if not os.path.exists(params.get('originVideo')):\n",
    "        logger.warning('originVideo video is not exists!')\n",
    "        exit(0)\n",
    "\n",
    "    if not os.path.exists(params.get('danmuVideo')):\n",
    "        logger.warning('danmuVideo video is not exists!')\n",
    "        exit(0)\n",
    "\n",
    "    # opencv打开视频\n",
    "    originVideo = cv2.VideoCapture(params.get('originVideo'))\n",
    "    danmuVideo = cv2.VideoCapture(params.get('danmuVideo'))\n",
    "\n",
    "    saveVideo = params.get('saveVideo')\n",
    "\n",
    "    # 检查两个视频的大小是否吻合，返回shape\n",
    "    height, width, _ = get_video_shape(danmuVideo, originVideo)\n",
    "\n",
    "    # 检查帧率是否吻合\n",
    "    fps = get_video_frame(danmuVideo, originVideo)\n",
    "\n",
    "    if saveVideo:\n",
    "        if not os.path.exists(params.get('output').split('/')[0]):\n",
    "            os.makedirs(params.get('output').split('/')[0])\n",
    "        fourcc = cv2.VideoWriter_fourcc(*'XVID')\n",
    "        fps = round(danmuVideo.get(cv2.CAP_PROP_FPS))\n",
    "        videoWriter = cv2.VideoWriter(\n",
    "            params.get('output'), fourcc, fps, (width, height))\n",
    "\n",
    "    elapsed = int()  # 记录处理了多少帧\n",
    "    start = timer()\n",
    "    stop = fps * params.get('keep')\n",
    "\n",
    "    while originVideo.isOpened():\n",
    "        elapsed += 1\n",
    "        _, frame_ori = originVideo.read()\n",
    "        _, frame_dan = danmuVideo.read()\n",
    "        if frame_ori is None or frame_dan is None:\n",
    "            logger.info('\\nEnd of Video')\n",
    "            break\n",
    "        results = model.detect([frame_ori], verbose=0)\n",
    "        r = results[0]\n",
    "        frame = display_instances(frame_ori, frame_dan, r['rois'], r['masks'], r['class_ids'], r['scores'])\n",
    "\n",
    "        if saveVideo:\n",
    "            videoWriter.write(frame)\n",
    "\n",
    "        # 输出fps\n",
    "        if elapsed % 5 == 0:\n",
    "            sys.stdout.write('\\r')\n",
    "            sys.stdout.write('{0:3.3f} FPS'.format(\n",
    "                elapsed / (timer() - start)))\n",
    "            sys.stdout.flush()\n",
    "\n",
    "        if elapsed == stop:\n",
    "            logger.info('Made a video of {} seconds.'.format(params.get('keep')))\n",
    "            break\n",
    "\n",
    "    sys.stdout.write('\\n')\n",
    "    if saveVideo:\n",
    "        videoWriter.release()\n",
    "    originVideo.release()\n",
    "    danmuVideo.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.7.3"
  },
  "latex_envs": {
   "LaTeX_envs_menu_present": true,
   "autoclose": false,
   "autocomplete": true,
   "bibliofile": "biblio.bib",
   "cite_by": "apalike",
   "current_citInitial": 1,
   "eqLabelWithNumbers": true,
   "eqNumInitial": 1,
   "hotkeys": {
    "equation": "Ctrl-E",
    "itemize": "Ctrl-I"
   },
   "labels_anchors": false,
   "latex_user_defs": false,
   "report_style_numbering": false,
   "user_envs_cfg": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}