【F2】图像转译-CycleGAN-视频预测.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "325f4eb2-b241-484b-a7b6-21c3c55f30cf",
   "metadata": {},
   "source": [
    "# 图像翻译（Image Translation） 之 CycleGAN"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "44148d35-55c0-47c1-bc31-12e7890dc61d",
   "metadata": {},
   "source": [
    "在 MMGeneration中调用 Image Translation 预训练模型\n",
    "\n",
    "> 作者：[同济子豪兄](https://space.bilibili.com/1900783)、杨逸飞 2022-4-10"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "62832831-58a0-4681-9b74-66b8a44191ce",
   "metadata": {},
   "source": [
    "## 进入MMGeneration主目录"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "01fbcfcc-e626-4926-93fb-10954715e503",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.chdir('mmgeneration')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "92bd845d-d076-403e-8241-6faa44565fbb",
   "metadata": {
    "tags": []
   },
   "source": [
    "## 导入工具包"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "798a948b-aa97-48e3-9d34-96767c20a395",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 导入mmcv和mmgeneration\n",
    "import mmcv\n",
    "from mmgen.apis import init_model, sample_img2img_model\n",
    "\n",
    "# 导入 opencv\n",
    "import cv2\n",
    "\n",
    "# 导入numpy和matplotlib\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# 定义可视化图像函数，输入图像 array，可视化图像\n",
    "def show_img_from_path(img_path):\n",
    "    '''输入 array，matplotlib 可视化格式为 RGB，因此需将 BGR 转 RGB，最后可视化出来'''\n",
    "    img_BGR = cv2.imread(img_path)\n",
    "    img_RGB = cv2.cvtColor(img_BGR, cv2.COLOR_BGR2RGB)\n",
    "    plt.imshow(img_RGB)\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "932a9fe9-be21-4706-abf0-c1d442926a98",
   "metadata": {},
   "source": [
    "## 封装函数，输入图像路径，保存输出图像并可视化"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "1ad0288f-da4e-440e-89d3-aa615416270d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def gen_vis_fake_img(input_path, output_path='outputs/F1/', target_domain='winter', figsize=15, frame_id=0):\n",
    "    \n",
    "    # 读入输入图像，获取高宽尺寸\n",
    "    input_img = cv2.imread(input_path)\n",
    "\n",
    "    # 生成另外一个图像域的图像，注意 target_domain 要设置正确\n",
    "    fake_imgs = sample_img2img_model(model, input_path, target_domain=target_domain)\n",
    "    # 获取生成图像尺寸\n",
    "    img_size = fake_imgs.shape[2]\n",
    "\n",
    "    # 分别抽取RGB三通道图像\n",
    "    RGB = np.zeros((img_size, img_size, 3))\n",
    "    RGB[:,:,0] = fake_imgs[0][2]\n",
    "    RGB[:,:,1] = fake_imgs[0][1]\n",
    "    RGB[:,:,2] = fake_imgs[0][0]\n",
    "    # 将生成图转为输入图像大小\n",
    "    RGB = cv2.resize(RGB, dsize=(input_img.shape[1],input_img.shape[0]))\n",
    "    # 像素值归一化为0-255的uint8自然图像\n",
    "    RGB = 255 * (RGB - RGB.min()) / (RGB.max()-RGB.min())\n",
    "    # 像素值转为整数\n",
    "    RGB = RGB.astype('uint8')\n",
    "    # 导出生成的图像文件\n",
    "    # cv2.imwrite(output_path+target_domain+'_'+input_path.split('/')[-1], cv2.cvtColor(RGB, cv2.COLOR_BGR2RGB))\n",
    "\n",
    "    plt.figure(figsize=(figsize, figsize))\n",
    "    # 展示原始输入图像\n",
    "    plt.subplot(1,2,1)\n",
    "    plt.title('input')\n",
    "    input_RGB = cv2.cvtColor(input_img, cv2.COLOR_BGR2RGB)\n",
    "    plt.xticks([])\n",
    "    plt.yticks([])\n",
    "    plt.imshow(input_RGB)\n",
    "\n",
    "    # 展示生成图\n",
    "    plt.subplot(1,2,2)\n",
    "    plt.title(target_domain)\n",
    "    plt.xticks([])\n",
    "    plt.yticks([])\n",
    "    plt.imshow(RGB)\n",
    "\n",
    "    # plt.show()\n",
    "    plt.savefig(f\"{temp_dir}/{frame_id:06d}.jpg\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "faf52561-ef99-4d5a-9d73-99053922e348",
   "metadata": {},
   "source": [
    "## 冬天和夏天互转，导入模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "1a27b7ab-3271-418e-b774-69f19c45e404",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/environment/miniconda3/lib/python3.7/site-packages/mmcv/cnn/bricks/conv_module.py:151: UserWarning: Unnecessary conv bias before batch/instance norm\n",
      "  'Unnecessary conv bias before batch/instance norm')\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "load checkpoint from http path: https://download.openmmlab.com/mmgen/cyclegan/refactor/cyclegan_lsgan_resnet_in_1x1_246200_summer2winter_convert-bgr_20210902_165932-fcf08dc1.pth\n"
     ]
    }
   ],
   "source": [
    "# 指定config文件路径\n",
    "config_file = 'configs/cyclegan/cyclegan_lsgan_resnet_in_summer2winter_b1x1_250k.py'\n",
    "# 指定预训练模型权重文件路径\n",
    "checkpoint_file = 'https://download.openmmlab.com/mmgen/cyclegan/refactor/cyclegan_lsgan_resnet_in_1x1_246200_summer2winter_convert-bgr_20210902_165932-fcf08dc1.pth'\n",
    "\n",
    "# 输出图像宽高像素尺寸\n",
    "img_size = 256\n",
    "\n",
    "model = init_model(config_file, checkpoint_file, device='cuda:0')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e6901616-9380-42ee-b67e-a64e2cf73218",
   "metadata": {},
   "source": [
    "## 生成视频"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ca45656b-7c5a-4351-8a0e-6928e8038ed7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "创建文件夹 20220508154035 用于存放每帧图像\n",
      "[>>                             ] 20/295, 1.1 task/s, elapsed: 17s, ETA:   240s"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/environment/miniconda3/lib/python3.7/site-packages/ipykernel_launcher.py:25: RuntimeWarning: More than 20 figures have been opened. Figures created through the pyplot interface (`matplotlib.pyplot.figure`) are retained until explicitly closed and may consume too much memory. (To control this warning, see the rcParam `figure.max_open_warning`).\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[>>>>>>>>>>>>>>>>>>>>>>>>>>>> ] 289/295, 1.2 task/s, elapsed: 250s, ETA:     5s"
     ]
    }
   ],
   "source": [
    "# 创建存放图像的临时文件夹\n",
    "import time\n",
    "temp_dir = time.strftime('%Y%m%d%H%M%S')\n",
    "\n",
    "if os.path.exists(temp_dir):\n",
    "    pass\n",
    "else:\n",
    "    os.mkdir(temp_dir)\n",
    "    print('创建文件夹 {} 用于存放每帧图像'.format(temp_dir))\n",
    "    \n",
    "input_video = 'data/tongji_video.mp4'\n",
    "\n",
    "# 读入待预测视频\n",
    "imgs = mmcv.VideoReader(input_video)\n",
    "prog_bar = mmcv.ProgressBar(len(imgs))\n",
    "\n",
    "# 对视频逐帧处理\n",
    "for frame_id, img in enumerate(imgs):\n",
    "    cv2.imwrite('temp.jpg', img)\n",
    "    gen_vis_fake_img('temp.jpg', target_domain='winter', figsize=15, frame_id=frame_id)\n",
    "    # 更新进度条\n",
    "    prog_bar.update()\n",
    "    \n",
    "# 由每一帧的图像文件，生成视频\n",
    "mmcv.frames2video(temp_dir, 'outputs/F2/F2_output.mp4', fps=imgs.fps, fourcc='mp4v')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5967a3f6-9640-4c11-814a-29fc97c3cf51",
   "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.7.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}