From 4ab368b3cb9cae0bcaa70072bab516a4b27e8bb7 Mon Sep 17 00:00:00 2001 From: ak9250 Date: Tue, 6 Aug 2019 21:23:28 -0400 Subject: [PATCH] Created using Colaboratory --- Tecogan.ipynb | 439 ++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 439 insertions(+) create mode 100644 Tecogan.ipynb diff --git a/Tecogan.ipynb b/Tecogan.ipynb new file mode 100644 index 0000000..4dde1a7 --- /dev/null +++ b/Tecogan.ipynb @@ -0,0 +1,439 @@ +{ + "nbformat": 4, + "nbformat_minor": 0, + "metadata": { + "colab": { + "name": "Tecogan.ipynb", + "version": "0.3.2", + "provenance": [], + "include_colab_link": true + }, + "kernelspec": { + "name": "python3", + "display_name": "Python 3" + }, + "accelerator": "GPU" + }, + "cells": [ + { + "cell_type": "markdown", + "metadata": { + "id": "view-in-github", + "colab_type": "text" + }, + "source": [ + "\"Open" + ] + }, + { + "cell_type": "code", + "metadata": { + "id": "x2EfvnD4KF-B", + "colab_type": "code", + "colab": {} + }, + "source": [ + "!git clone https://github.com/thunil/TecoGAN.git" + ], + "execution_count": 0, + "outputs": [] + }, + { + "cell_type": "code", + "metadata": { + "id": "leqdK5tsNSbB", + "colab_type": "code", + "colab": {} + }, + "source": [ + "cd TecoGAN" + ], + "execution_count": 0, + "outputs": [] + }, + { + "cell_type": "code", + "metadata": { + "id": "gjzQDMzBPU-k", + "colab_type": "code", + "colab": {} + }, + "source": [ + "!pip3 install -r requirements.txt" + ], + "execution_count": 0, + "outputs": [] + }, + { + "cell_type": "code", + "metadata": { + "id": "32PrvHTyQj9i", + "colab_type": "code", + "colab": {} + }, + "source": [ + "!python3 runGan.py 0" + ], + "execution_count": 0, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "KEplQkE5q9-G", + "colab_type": "text" + }, + "source": [ + "change runcase 1 dataset name to yours" + ] + }, + { + "cell_type": "code", + "metadata": { + "id": "UfGhgrC6q3Gv", + "colab_type": "code", + "colab": {} + }, + "source": [ + "%%writefile runGan.py\n", + "'''\n", + "several running examples, run with\n", + "python3 runGan.py 1 # the last number is the run case number\n", + "\n", + "runcase == 1 inference a trained model\n", + "runcase == 2 calculate the metrics, and save the numbers in csv\n", + "runcase == 3 training TecoGAN\n", + "runcase == 4 training FRVSR\n", + "runcase == ... coming... data preparation and so on...\n", + "'''\n", + "import os, subprocess, sys, datetime, signal, shutil\n", + "\n", + "runcase = int(sys.argv[1])\n", + "print (\"Testing test case %d\" % runcase)\n", + "\n", + "def preexec(): # Don't forward signals.\n", + " os.setpgrp()\n", + " \n", + "def mycall(cmd, block=False):\n", + " if not block:\n", + " return subprocess.Popen(cmd)\n", + " else:\n", + " return subprocess.Popen(cmd, preexec_fn = preexec)\n", + " \n", + "def folder_check(path):\n", + " try_num = 1\n", + " oripath = path[:-1] if path.endswith('/') else path\n", + " while os.path.exists(path):\n", + " print(\"Delete existing folder \" + path + \"?(Y/N)\")\n", + " decision = input()\n", + " if decision == \"Y\":\n", + " shutil.rmtree(path, ignore_errors=True)\n", + " break\n", + " else:\n", + " path = oripath + \"_%d/\"%try_num\n", + " try_num += 1\n", + " print(path)\n", + " \n", + " return path\n", + "\n", + "if( runcase == 0 ): # download inference data, trained models\n", + " # download the trained model\n", + " if(not os.path.exists(\"./model/\")): os.mkdir(\"./model/\")\n", + " cmd1 = \"wget https://ge.in.tum.de/download/data/TecoGAN/model.zip -O model/model.zip;\"\n", + " cmd1 += \"unzip model/model.zip -d model; rm model/model.zip\"\n", + " subprocess.call(cmd1, shell=True)\n", + " \n", + " # download some test data\n", + " cmd2 = \"wget https://ge.in.tum.de/download/data/TecoGAN/vid3_LR.zip -O LR/vid3.zip;\"\n", + " cmd2 += \"unzip LR/vid3.zip -d LR; rm LR/vid3.zip\"\n", + " subprocess.call(cmd2, shell=True)\n", + " \n", + " cmd2 = \"wget https://ge.in.tum.de/download/data/TecoGAN/tos_LR.zip -O LR/tos.zip;\"\n", + " cmd2 += \"unzip LR/tos.zip -d LR; rm LR/tos.zip\"\n", + " subprocess.call(cmd2, shell=True)\n", + " \n", + " # download the ground-truth data\n", + " if(not os.path.exists(\"./HR/\")): os.mkdir(\"./HR/\")\n", + " cmd3 = \"wget https://ge.in.tum.de/download/data/TecoGAN/vid4_HR.zip -O HR/vid4.zip;\"\n", + " cmd3 += \"unzip HR/vid4.zip -d HR; rm HR/vid4.zip\"\n", + " subprocess.call(cmd3, shell=True)\n", + " \n", + " cmd3 = \"wget https://ge.in.tum.de/download/data/TecoGAN/tos_HR.zip -O HR/tos.zip;\"\n", + " cmd3 += \"unzip HR/tos.zip -d HR; rm HR/tos.zip\"\n", + " subprocess.call(cmd3, shell=True)\n", + " \n", + "elif( runcase == 1 ): # inference a trained model\n", + " \n", + " dirstr = './results/' # the place to save the results\n", + " testpre = ['calendar'] # the test cases\n", + "\n", + " if (not os.path.exists(dirstr)): os.mkdir(dirstr)\n", + " \n", + " # run these test cases one by one:\n", + " for nn in range(len(testpre)):\n", + " cmd1 = [\"python3\", \"main.py\",\n", + " \"--cudaID\", \"0\", # set the cudaID here to use only one GPU\n", + " \"--output_dir\", dirstr, # Set the place to put the results.\n", + " \"--summary_dir\", os.path.join(dirstr, 'log/'), # Set the place to put the log. \n", + " \"--mode\",\"inference\", \n", + " \"--input_dir_LR\", os.path.join(\"./LR/\", testpre[nn]), # the LR directory\n", + " #\"--input_dir_HR\", os.path.join(\"./HR/\", testpre[nn]), # the HR directory\n", + " # one of (input_dir_HR,input_dir_LR) should be given\n", + " \"--output_pre\", testpre[nn], # the subfolder to save current scene, optional\n", + " \"--num_resblock\", \"16\", # our model has 16 residual blocks, \n", + " # the pre-trained FRVSR and TecoGAN mini have 10 residual blocks\n", + " \"--checkpoint\", './model/TecoGAN', # the path of the trained model,\n", + " \"--output_ext\", \"png\" # png is more accurate, jpg is smaller\n", + " ]\n", + " mycall(cmd1).communicate()\n", + "\n", + "elif( runcase == 2 ): # calculate all metrics, and save the csv files, should use png\n", + "\n", + " testpre = [\"calendar\"] # just put more scenes to evaluate all of them\n", + " dirstr = './results/' # the outputs\n", + " tarstr = './HR/' # the GT\n", + "\n", + " tar_list = [(tarstr+_) for _ in testpre]\n", + " out_list = [(dirstr+_) for _ in testpre]\n", + " cmd1 = [\"python3\", \"metrics.py\",\n", + " \"--output\", dirstr+\"metric_log/\",\n", + " \"--results\", \",\".join(out_list),\n", + " \"--targets\", \",\".join(tar_list),\n", + " ]\n", + " mycall(cmd1).communicate()\n", + " \n", + "elif( runcase == 3 ): # Train TecoGAN\n", + " '''\n", + " In order to use the VGG as a perceptual loss,\n", + " we download from TensorFlow-Slim image classification model library:\n", + " https://github.com/tensorflow/models/tree/master/research/slim \n", + " '''\n", + " VGGPath = \"model/\" # the path for the VGG model, there should be a vgg_19.ckpt inside\n", + " VGGModelPath = os.path.join(VGGPath, \"vgg_19.ckpt\")\n", + " if(not os.path.exists(VGGPath)): os.mkdir(VGGPath)\n", + " if(not os.path.exists(VGGModelPath)):\n", + " # Download the VGG 19 model from \n", + " print(\"VGG model not found, downloading to %s\"%VGGPath)\n", + " cmd0 = \"wget http://download.tensorflow.org/models/vgg_19_2016_08_28.tar.gz -O \" + os.path.join(VGGPath, \"vgg19.tar.gz\")\n", + " cmd0 += \";tar -xvf \" + os.path.join(VGGPath,\"vgg19.tar.gz\") + \" -C \" + VGGPath + \"; rm \"+ os.path.join(VGGPath, \"vgg19.tar.gz\")\n", + " subprocess.call(cmd0, shell=True)\n", + " \n", + " '''\n", + " Use our pre-trained FRVSR model. If you want to train one, try runcase 4, and update this path by:\n", + " FRVSRModel = \"ex_FRVSRmm-dd-hh/model-500000\"\n", + " '''\n", + " FRVSRModel = \"model/ourFRVSR\" \n", + " if(not os.path.exists(FRVSRModel+\".data-00000-of-00001\")):\n", + " # Download our pre-trained FRVSR model\n", + " print(\"pre-trained FRVSR model not found, downloading\")\n", + " cmd0 = \"wget http://ge.in.tum.de/download/2019-TecoGAN/FRVSR_Ours.zip -O model/ofrvsr.zip;\"\n", + " cmd0 += \"unzip model/ofrvsr.zip -d model; rm model/ofrvsr.zip\"\n", + " subprocess.call(cmd0, shell=True)\n", + " \n", + " TrainingDataPath = \"/mnt/netdisk/video_data/\" \n", + " \n", + " '''Prepare Training Folder'''\n", + " # path appendix, manually define it, or use the current datetime, now_str = \"mm-dd-hh\"\n", + " now_str = datetime.datetime.now().strftime(\"%m-%d-%H\")\n", + " train_dir = folder_check(\"ex_TecoGAN%s/\"%now_str)\n", + " # train TecoGAN, loss = l2 + VGG54 loss + A spatio-temporal Discriminator\n", + " cmd1 = [\"python3\", \"main.py\",\n", + " \"--cudaID\", \"0\", # set the cudaID here to use only one GPU\n", + " \"--output_dir\", train_dir, # Set the place to save the models.\n", + " \"--summary_dir\", os.path.join(train_dir,\"log/\"), # Set the place to save the log. \n", + " \"--mode\",\"train\",\n", + " \"--batch_size\", \"4\" , # small, because GPU memory is not big\n", + " \"--RNN_N\", \"10\" , # train with a sequence of RNN_N frames, >6 is better, >10 is not necessary\n", + " \"--movingFirstFrame\", # a data augmentation\n", + " \"--random_crop\",\n", + " \"--crop_size\", \"32\",\n", + " \"--learning_rate\", \"0.00005\",\n", + " # -- learning_rate step decay, here it is not used --\n", + " \"--decay_step\", \"500000\", \n", + " \"--decay_rate\", \"1.0\", # 1.0 means no decay\n", + " \"--stair\",\n", + " \"--beta\", \"0.9\", # ADAM training parameter beta\n", + " \"--max_iter\", \"500000\", # 500k or more, the one we present is trained for 900k\n", + " \"--save_freq\", \"10000\", # the frequency we save models\n", + " # -- network architecture parameters --\n", + " \"--num_resblock\", \"16\", # FRVSR and TecoGANmini has num_resblock as 10. The TecoGAN has 16.\n", + " # -- VGG loss, disable with vgg_scaling < 0\n", + " \"--vgg_scaling\", \"0.2\",\n", + " \"--vgg_ckpt\", VGGModelPath, # necessary if vgg_scaling > 0\n", + " ]\n", + " '''Video Training data:\n", + " please udate the TrainingDataPath according to ReadMe.md\n", + " input_video_pre is hard coded as scene in dataPrepare.py at line 142\n", + " str_dir is the starting index for training data\n", + " end_dir is the ending index for training data\n", + " end_dir+1 is the starting index for validation data\n", + " end_dir_val is the ending index for validation data\n", + " max_frm should be duration (in dataPrepare.py) -1\n", + " queue_thread: how many cpu can be used for loading data when training\n", + " name_video_queue_capacity, video_queue_capacity: how much memory can be used\n", + " '''\n", + " cmd1 += [\n", + " \"--input_video_dir\", TrainingDataPath, \n", + " \"--input_video_pre\", \"scene\",\n", + " \"--str_dir\", \"2000\",\n", + " \"--end_dir\", \"2250\",\n", + " \"--end_dir_val\", \"2290\",\n", + " \"--max_frm\", \"119\",\n", + " # -- cpu memory for data loading --\n", + " \"--queue_thread\", \"12\",# Cpu threads for the data. >4 to speedup the training\n", + " \"--name_video_queue_capacity\", \"1024\",\n", + " \"--video_queue_capacity\", \"1024\",\n", + " ]\n", + " '''\n", + " loading the pre-trained model from FRVSR can make the training faster\n", + " --checkpoint, path of the model, here our pre-trained FRVSR is given\n", + " --pre_trained_model, to continue an old (maybe accidentally stopeed) training, \n", + " pre_trained_model should be false, and checkpoint should be the last model such as \n", + " ex_TecoGANmm-dd-hh/model-xxxxxxx\n", + " To start a new and different training, pre_trained_model is True. \n", + " The difference here is \n", + " whether to load the whole graph icluding ADAM training averages/momentums/ and so on\n", + " or just load existing pre-trained weights.\n", + " '''\n", + " cmd1 += [ # based on a pre-trained FRVSR model. Here we want to train a new adversarial training\n", + " \"--pre_trained_model\", # True\n", + " \"--checkpoint\", FRVSRModel,\n", + " ]\n", + " \n", + " # the following can be used to train TecoGAN continuously\n", + " # old_model = \"model/ex_TecoGANmm-dd-hh/model-xxxxxxx\" \n", + " # cmd1 += [ # Here we want to train continuously\n", + " # \"--nopre_trained_model\", # False\n", + " # \"--checkpoint\", old_model,\n", + " # ]\n", + " \n", + " ''' parameters for GAN training '''\n", + " cmd1 += [\n", + " \"--ratio\", \"0.01\", # the ratio for the adversarial loss from the Discriminator to the Generator\n", + " \"--Dt_mergeDs\", # if Dt_mergeDs == False, only use temporal inputs, so we have a temporal Discriminator\n", + " # else, use both temporal and spatial inputs, then we have a Dst, the spatial and temporal Discriminator\n", + " ]\n", + " ''' if the generator is pre-trained, to fade in the discriminator is usually more stable.\n", + " the weight of the adversarial loss will be weighed with a weight, started from Dt_ratio_0, \n", + " and increases until Dt_ratio_max, the increased value is Dt_ratio_add per training step\n", + " For example, fading Dst in smoothly in the first 4k steps is \n", + " \"--Dt_ratio_max\", \"1.0\", \"--Dt_ratio_0\", \"0.0\", \"--Dt_ratio_add\", \"0.00025\"\n", + " '''\n", + " cmd1 += [ # here, the fading in is disabled \n", + " \"--Dt_ratio_max\", \"1.0\",\n", + " \"--Dt_ratio_0\", \"1.0\", \n", + " \"--Dt_ratio_add\", \"0.0\", \n", + " ]\n", + " ''' Other Losses '''\n", + " cmd1 += [\n", + " \"--pingpang\", # our Ping-Pang loss\n", + " \"--pp_scaling\", \"0.5\", # the weight of the our bi-directional loss, 0.0~0.5\n", + " \"--D_LAYERLOSS\", # use feature layer losses from the discriminator\n", + " ]\n", + " \n", + " pid = mycall(cmd1, block=True) \n", + " try: # catch interruption for training\n", + " pid.communicate()\n", + " except KeyboardInterrupt: # Ctrl + C to stop current training try to save the last model \n", + " print(\"runGAN.py: sending SIGINT signal to the sub process...\")\n", + " pid.send_signal(signal.SIGINT)\n", + " # try to save the last model \n", + " pid.communicate()\n", + " print(\"runGAN.py: finished...\")\n", + " \n", + " \n", + "elif( runcase == 4 ): # Train FRVSR, loss = l2 warp + l2 content\n", + " now_str = datetime.datetime.now().strftime(\"%m-%d-%H\")\n", + " train_dir = folder_check(\"ex_FRVSR%s/\"%now_str)\n", + " cmd1 = [\"python3\", \"main.py\",\n", + " \"--cudaID\", \"0\", # set the cudaID here to use only one GPU\n", + " \"--output_dir\", train_dir, # Set the place to save the models.\n", + " \"--summary_dir\", os.path.join(train_dir,\"log/\"), # Set the place to save the log. \n", + " \"--mode\",\"train\",\n", + " \"--batch_size\", \"4\" , # small, because GPU memory is not big\n", + " \"--RNN_N\", \"10\" , # train with a sequence of RNN_N frames, >6 is better, >10 is not necessary\n", + " \"--movingFirstFrame\", # a data augmentation\n", + " \"--random_crop\",\n", + " \"--crop_size\", \"32\",\n", + " \"--learning_rate\", \"0.00005\",\n", + " # -- learning_rate step decay, here it is not used --\n", + " \"--decay_step\", \"500000\", \n", + " \"--decay_rate\", \"1.0\", # 1.0 means no decay\n", + " \"--stair\",\n", + " \"--beta\", \"0.9\", # ADAM training parameter beta\n", + " \"--max_iter\", \"500000\", # 500k is usually fine for FRVSR, GAN versions need more to be stable\n", + " \"--save_freq\", \"10000\", # the frequency we save models\n", + " # -- network architecture parameters --\n", + " \"--num_resblock\", \"10\", # a smaller model\n", + " \"--ratio\", \"-0.01\", # the ratio for the adversarial loss, negative means disabled\n", + " \"--nopingpang\",\n", + " ]\n", + " '''Video Training data... Same as runcase 3...'''\n", + " TrainingDataPath = \"/mnt/netdisk/video_data/\"\n", + " cmd1 += [\n", + " \"--input_video_dir\", TrainingDataPath, \n", + " \"--input_video_pre\", \"scene\",\n", + " \"--str_dir\", \"2000\",\n", + " \"--end_dir\", \"2250\",\n", + " \"--end_dir_val\", \"2290\",\n", + " \"--max_frm\", \"119\",\n", + " # -- cpu memory for data loading --\n", + " \"--queue_thread\", \"12\",# Cpu threads for the data. >4 to speedup the training\n", + " \"--name_video_queue_capacity\", \"1024\",\n", + " \"--video_queue_capacity\", \"1024\",\n", + " ]\n", + " \n", + " pid = mycall(cmd1, block=True)\n", + " try: # catch interruption for training\n", + " pid.communicate()\n", + " except KeyboardInterrupt: # Ctrl + C to stop current training try to save the last model \n", + " print(\"runGAN.py: sending SIGINT signal to the sub process...\")\n", + " pid.send_signal(signal.SIGINT)\n", + " # try to save the last model \n", + " pid.communicate()\n", + " print(\"runGAN.py: finished...\")" + ], + "execution_count": 0, + "outputs": [] + }, + { + "cell_type": "code", + "metadata": { + "id": "ZigwHEG4pi9a", + "colab_type": "code", + "colab": {} + }, + "source": [ + "cd .." + ], + "execution_count": 0, + "outputs": [] + }, + { + "cell_type": "code", + "metadata": { + "id": "PbHwtctaQsVg", + "colab_type": "code", + "colab": {} + }, + "source": [ + "!python3 runGan.py 1" + ], + "execution_count": 0, + "outputs": [] + }, + { + "cell_type": "code", + "metadata": { + "colab_type": "code", + "id": "iE6EpGtClUZK", + "colab": {} + }, + "source": [ + "!python3 runGan.py 2" + ], + "execution_count": 0, + "outputs": [] + } + ] +} \ No newline at end of file