├── nowcasting ├── __init__.py ├── helpers │ ├── __init__.py │ ├── gifmaker.py │ ├── ordered_easydict.py │ ├── log_analysis.py │ ├── visualization.py │ └── msssim.py ├── operators │ ├── __init__.py │ ├── transformations.py │ └── traj_rnn.py ├── mask.py ├── prediction_base_factory.py ├── utils.py ├── hko_factory.py ├── image.py ├── numba_accelerated.py └── my_module.py ├── VarFlow ├── varflow │ ├── __init__.py │ └── varflow.py ├── build.sh ├── Data │ ├── yos_img_08.jpg │ └── yos_img_09.jpg ├── windows_cmake_command.txt ├── CMakeLists.txt ├── c_api.h ├── setup.py ├── README.md ├── Readme.txt ├── c_api.cpp ├── ProfTimer.h ├── VarFlow.h └── example.cpp ├── hko_data ├── uncertain_days.txt ├── mask_dat.npz ├── intensity_day.pkl ├── benchmark_stat │ ├── hko7_all_15_in5_out20_stride5_fixed.json │ ├── hko7_rainy_test_in5_out20_stride5_fixed.json │ ├── README.md │ ├── hko7_rainy_valid_in5_out20_stride5_fixed.json │ ├── hko7_all_15_out20_stride5_online.json │ ├── hko7_rainy_valid_out20_stride5_online.json │ └── hko7_rainy_test_out20_stride5_online.json ├── hko7_rainy_valid_days.txt ├── README.md └── hko7_rainy_test_days.txt ├── mnist_data ├── ebrnn1_link_sample.gif ├── ebrnn2_link_sample.gif ├── ebrnn3_link_sample.gif ├── fbrnn1_link_sample.gif ├── fbrnn2_link_sample.gif ├── fbrnn3_link_sample.gif └── movingmnist_10000_nodistr.npz ├── HKO-7_Dataset_Undertaking_fillable.pdf ├── experiments ├── hko │ ├── download_pretrained.sh │ ├── configurations │ │ ├── conv2d_3d │ │ │ ├── conv2d.yml │ │ │ └── conv3d.yml │ │ ├── convgru_55_55_33_1_64_1_192_1_192_b4.yml │ │ ├── convgru_55_55_33_1_64_1_192_1_192_nobal_b4.yml │ │ └── trajgru_55_55_33_1_64_1_192_1_192_13_13_9_b4.yml │ ├── last_frame_prediction.py │ ├── README.md │ ├── hko_rnn_test.py │ ├── hko_factory.py │ ├── rover.py │ └── deconvolution.py └── movingmnist │ ├── download_pretrained.sh │ ├── configurations │ ├── conv2d_3d │ │ ├── conv2d.yml │ │ └── conv3d.yml │ ├── convgru_1_64_1_96_1_96_K5.yml │ ├── convgru_1_64_1_96_1_96_K7.yml │ ├── convgru_1_64_1_96_1_96_K3D2.yml │ ├── convgru_1_64_1_96_1_96_K5_DFN.yml │ ├── trajgru_1_64_1_96_1_96_L13.yml │ ├── trajgru_1_64_1_96_1_96_L17.yml │ ├── trajgru_1_64_1_96_1_96_L5.yml │ └── trajgru_1_64_1_96_1_96_L9.yml │ ├── mnist_rnn_factory.py │ ├── analyze_learned_structure.py │ ├── README.md │ ├── mnist_rnn_test.py │ ├── deconvolution.py │ └── mnist_rnn_main.py ├── LICENSE ├── setup.py ├── download_all.py ├── .gitignore └── README.md /nowcasting/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /nowcasting/helpers/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /VarFlow/varflow/__init__.py: -------------------------------------------------------------------------------- 1 | from .varflow import * -------------------------------------------------------------------------------- /VarFlow/build.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | mkdir build 3 | cd build 4 | cmake .. 5 | make 6 | -------------------------------------------------------------------------------- /hko_data/uncertain_days.txt: -------------------------------------------------------------------------------- 1 | 20140502 last several frames 2 | 20140612 last several frames 3 | -------------------------------------------------------------------------------- /hko_data/mask_dat.npz: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/sxjscience/HKO-7/HEAD/hko_data/mask_dat.npz -------------------------------------------------------------------------------- /VarFlow/Data/yos_img_08.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/sxjscience/HKO-7/HEAD/VarFlow/Data/yos_img_08.jpg -------------------------------------------------------------------------------- /VarFlow/Data/yos_img_09.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/sxjscience/HKO-7/HEAD/VarFlow/Data/yos_img_09.jpg -------------------------------------------------------------------------------- /hko_data/intensity_day.pkl: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/sxjscience/HKO-7/HEAD/hko_data/intensity_day.pkl -------------------------------------------------------------------------------- /mnist_data/ebrnn1_link_sample.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/sxjscience/HKO-7/HEAD/mnist_data/ebrnn1_link_sample.gif -------------------------------------------------------------------------------- /mnist_data/ebrnn2_link_sample.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/sxjscience/HKO-7/HEAD/mnist_data/ebrnn2_link_sample.gif -------------------------------------------------------------------------------- /mnist_data/ebrnn3_link_sample.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/sxjscience/HKO-7/HEAD/mnist_data/ebrnn3_link_sample.gif -------------------------------------------------------------------------------- /mnist_data/fbrnn1_link_sample.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/sxjscience/HKO-7/HEAD/mnist_data/fbrnn1_link_sample.gif -------------------------------------------------------------------------------- /mnist_data/fbrnn2_link_sample.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/sxjscience/HKO-7/HEAD/mnist_data/fbrnn2_link_sample.gif -------------------------------------------------------------------------------- /mnist_data/fbrnn3_link_sample.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/sxjscience/HKO-7/HEAD/mnist_data/fbrnn3_link_sample.gif -------------------------------------------------------------------------------- /HKO-7_Dataset_Undertaking_fillable.pdf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/sxjscience/HKO-7/HEAD/HKO-7_Dataset_Undertaking_fillable.pdf -------------------------------------------------------------------------------- /mnist_data/movingmnist_10000_nodistr.npz: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/sxjscience/HKO-7/HEAD/mnist_data/movingmnist_10000_nodistr.npz -------------------------------------------------------------------------------- /VarFlow/windows_cmake_command.txt: -------------------------------------------------------------------------------- 1 | cmake -G "Visual Studio 14 2015 Win64" ^ 2 | -DCMAKE_BUILD_TYPE=Release ^ 3 | -DCMAKE_CONFIGURATION_TYPES="Release" ^ 4 | .. -------------------------------------------------------------------------------- /hko_data/benchmark_stat/hko7_all_15_in5_out20_stride5_fixed.json: -------------------------------------------------------------------------------- 1 | { 2 | "pred_seq_num": 17482, 3 | "episode_num": 17482, 4 | "episode_start_datetime": [] 5 | } -------------------------------------------------------------------------------- /hko_data/benchmark_stat/hko7_rainy_test_in5_out20_stride5_fixed.json: -------------------------------------------------------------------------------- 1 | { 2 | "pred_seq_num": 6053, 3 | "episode_num": 6053, 4 | "episode_start_datetime": [] 5 | } -------------------------------------------------------------------------------- /hko_data/benchmark_stat/README.md: -------------------------------------------------------------------------------- 1 | HKO-7 Benchmark Statistics 2 | --------------------------- 3 | Stores the necessary statistics for running the benchmark environmente 4 | -------------------------------------------------------------------------------- /hko_data/benchmark_stat/hko7_rainy_valid_in5_out20_stride5_fixed.json: -------------------------------------------------------------------------------- 1 | { 2 | "pred_seq_num": 2166, 3 | "episode_num": 2166, 4 | "episode_start_datetime": [] 5 | } -------------------------------------------------------------------------------- /nowcasting/operators/__init__.py: -------------------------------------------------------------------------------- 1 | from .base_rnn import * 2 | from .conv_rnn import * 3 | from .traj_rnn import * 4 | from .transformations import * 5 | from .common import * 6 | -------------------------------------------------------------------------------- /experiments/hko/download_pretrained.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | wget -O nips2017_pretrained.zip https://www.dropbox.com/sh/cp8zpi08umfiyha/AAAS6HJSsDQPjpKlxnBHtHvga?dl=1 --no-check-certificate 3 | unzip nips2017_pretrained.zip 4 | -------------------------------------------------------------------------------- /experiments/movingmnist/download_pretrained.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | wget -O nips2017_pretrained_movingmnist.zip https://www.dropbox.com/sh/n7gxfdd1pdasoio/AAC8uC4yto5Uam_7f3BEl-3La?dl=1 --no-check-certificate 3 | unzip nips2017_pretrained_movingmnist.zip 4 | -------------------------------------------------------------------------------- /hko_data/benchmark_stat/hko7_all_15_out20_stride5_online.json: -------------------------------------------------------------------------------- 1 | { 2 | "pred_seq_num": 17482, 3 | "episode_num": 6, 4 | "episode_start_datetime": [ 5 | "201501010000", 6 | "201503262200", 7 | "201506151242", 8 | "201508151724", 9 | "201511191148", 10 | "201511191530" 11 | ] 12 | } -------------------------------------------------------------------------------- /VarFlow/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | cmake_minimum_required(VERSION 2.8) 2 | project( Varflow ) 3 | find_package( OpenCV REQUIRED ) 4 | include_directories(SYSTEM ${OpenCV_INCLUDE_DIRS}) 5 | add_library(varflow SHARED c_api.cpp VarFlow.cpp) 6 | add_executable( FlowExample example.cpp VarFlow.cpp) 7 | target_link_libraries( FlowExample ${OpenCV_LIBS} ) 8 | target_link_libraries( varflow ${OpenCV_LIBS} ) -------------------------------------------------------------------------------- /VarFlow/c_api.h: -------------------------------------------------------------------------------- 1 | /*! 2 | * Copyright (c) 2017 by Contributors 3 | * \file c_api.h 4 | * \brief C API of VarFlow 5 | */ 6 | #ifndef VARFLOW_C_API_H_ 7 | #define VARFLOW_C_API_H_ 8 | 9 | 10 | #ifdef __cplusplus 11 | extern "C" { 12 | #endif // __cplusplus 13 | 14 | #include 15 | #include 16 | #include 17 | #ifdef _WIN32 18 | #define VARFLOW_DLL __declspec(dllexport) 19 | #else 20 | #define VARFLOW_DLL 21 | #endif 22 | 23 | VARFLOW_DLL void varflow(int width, int height, int max_level, int start_level, int n1, int n2, 24 | float rho, float alpha, float sigma, void* U, void* V, 25 | void* I1, void* I2); 26 | 27 | #ifdef __cplusplus 28 | } 29 | #endif // __cplusplus 30 | #endif -------------------------------------------------------------------------------- /nowcasting/helpers/gifmaker.py: -------------------------------------------------------------------------------- 1 | from PIL import Image 2 | import numpy as np 3 | 4 | def save_gif(single_seq, fname): 5 | """Save a single gif consisting of image sequence in single_seq to fname.""" 6 | img_seq = [Image.fromarray(img.astype(np.float32) * 255, 'F').convert("L") for img in single_seq] 7 | img = img_seq[0] 8 | img.save(fname, save_all=True, append_images=img_seq[1:]) 9 | 10 | def save_gifs(seq, prefix): 11 | """Save several gifs. 12 | 13 | Args: 14 | seq: Shape (num_gifs, IMG_SIZE, IMG_SIZE) 15 | prefix: prefix-idx.gif will be the final filename. 16 | """ 17 | 18 | for idx, single_seq in enumerate(seq): 19 | save_gif(single_seq, "{}-{}.gif".format(prefix, idx)) 20 | -------------------------------------------------------------------------------- /experiments/movingmnist/configurations/conv2d_3d/conv2d.yml: -------------------------------------------------------------------------------- 1 | SEED: 2 2 | MODEL: 3 | VALID_ITER: 10000 4 | TRAIN: 5 | LR_DECAY_FACTOR: 0.7 6 | BETA1: 0.5 7 | GAMMA1: 0.9 8 | BATCH_SIZE: 32 9 | MAX_ITER: 80000 10 | LR: 0.0001 11 | LR_DECAY_ITER: 20000 12 | WD: 0 13 | GRAD_CLIP: 50.0 14 | MIN_LR: 1.0e-06 15 | OPTIMIZER: adam 16 | EPS: 1.0e-08 17 | DECONVBASELINE: 18 | BN: true 19 | ENCODER: concat 20 | COMPAT: 21 | CONV_INSTEADOF_FC_IN_ENCODER: false 22 | USE_3D: false 23 | BASE_NUM_FILTER: 70 24 | SAVE_ITER: 10000 25 | L2_LAMBDA: 1.0 26 | MOVINGMNIST: 27 | OUT_LEN: 10 28 | VELOCITY_UPPER: 3.6 29 | ILLUMINATION_LOWER: 0.6 30 | SCALE_VARIATION_LOWER: 0.9 31 | SCALE_VARIATION_UPPER: 1.1 32 | ILLUMINATION_UPPER: 1.0 33 | DISTRACTOR_NUM: 0 34 | ROTATION_LOWER: -30 35 | ROTATION_UPPER: 30 36 | VELOCITY_LOWER: 0.0 -------------------------------------------------------------------------------- /experiments/movingmnist/configurations/conv2d_3d/conv3d.yml: -------------------------------------------------------------------------------- 1 | SEED: 2 2 | MODEL: 3 | VALID_ITER: 10000 4 | TRAIN: 5 | LR_DECAY_FACTOR: 0.7 6 | BETA1: 0.5 7 | GAMMA1: 0.9 8 | BATCH_SIZE: 32 9 | MAX_ITER: 80000 10 | LR: 0.0001 11 | LR_DECAY_ITER: 20000 12 | WD: 0 13 | GRAD_CLIP: 50.0 14 | MIN_LR: 1.0e-06 15 | OPTIMIZER: adam 16 | EPS: 1.0e-08 17 | DECONVBASELINE: 18 | BN: true 19 | ENCODER: concat 20 | COMPAT: 21 | CONV_INSTEADOF_FC_IN_ENCODER: false 22 | USE_3D: true 23 | BASE_NUM_FILTER: 70 24 | SAVE_ITER: 10000 25 | L2_LAMBDA: 1.0 26 | MOVINGMNIST: 27 | OUT_LEN: 10 28 | VELOCITY_UPPER: 3.6 29 | ILLUMINATION_LOWER: 0.6 30 | SCALE_VARIATION_LOWER: 0.9 31 | SCALE_VARIATION_UPPER: 1.1 32 | ILLUMINATION_UPPER: 1.0 33 | DISTRACTOR_NUM: 0 34 | ROTATION_LOWER: -30 35 | ROTATION_UPPER: 30 36 | VELOCITY_LOWER: 0.0 -------------------------------------------------------------------------------- /VarFlow/setup.py: -------------------------------------------------------------------------------- 1 | import os 2 | import setuptools 3 | 4 | 5 | setuptools.setup( 6 | name='varflow', 7 | version="1.0", 8 | author="Xingjian Shi", 9 | author_email="xshiab@cse.ust.hk", 10 | packages=setuptools.find_packages(), 11 | description='Python Wrapper of VarFlow', 12 | long_description=open(os.path.join(os.path.dirname( 13 | os.path.abspath(__file__)), 'README.md')).read(), 14 | license='MIT', 15 | url='https://github.com/sxjscience/HKO-7/VarFlow', 16 | install_requires=['numpy'], 17 | classifiers=['Development Status :: 2 - Pre-Alpha', 18 | 'Intended Audience :: Science/Research', 19 | 'License :: OSI Approved :: MIT License', 20 | 'Natural Language :: English', 21 | 'Operating System :: OS Independent', 22 | 'Topic :: Scientific/Engineering :: Artificial Intelligence'], 23 | ) 24 | -------------------------------------------------------------------------------- /VarFlow/README.md: -------------------------------------------------------------------------------- 1 | Python Wrapper of VarFlow 2 | ------------------------- 3 | A python wrapper of VarFlow (http://sourceforge.net/projects/varflow). 4 | 5 | To install, first create a directory "build" and then run cmake. 6 | 7 | On windows, follows the following command: 8 | ```bash 9 | mkdir build 10 | cmake -G "Visual Studio 14 2015 Win64" ^ 11 | -DCMAKE_BUILD_TYPE=Release ^ 12 | -DCMAKE_CONFIGURATION_TYPES="Release" ^ 13 | .. 14 | ``` 15 | 16 | On linux, follows the following command 17 | ```bash 18 | mkdir build 19 | cd build 20 | cmake .. 21 | ``` 22 | 23 | If OpenCV is not found, you can try adding the `-DOpenCV_DIR` flag. The following is an example: 24 | ```bash 25 | mkdir build 26 | cd build 27 | cmake -DOpenCV_DIR=/usr/local/software/opencv/share/OpenCV .. 28 | make 29 | ``` 30 | 31 | Run the example in `varflow/varflow.py` to validate the installation. 32 | 33 | After that, you can install the package via 34 | ```bash 35 | python setup.py develop 36 | ``` 37 | -------------------------------------------------------------------------------- /hko_data/hko7_rainy_valid_days.txt: -------------------------------------------------------------------------------- 1 | 20100105, 0.8 2 | 20120124, 4.5 3 | 20100222, 0.1 4 | 20120208, 0.7 5 | 20120228, 18 6 | 20140208, 0.3 7 | 20140219, 3.8 8 | 20090301, 0.8 9 | 20100325, 8.9 10 | 20130330, 58.2 11 | 20140314, 0.4 12 | 20100418, 3.1 13 | 20120429, 21.2 14 | 20130420, 12.2 15 | 20140414, 0.4 16 | 20090519, 0.3 17 | 20090520, 10.9 18 | 20120519, 7.4 19 | 20120520, 49.6 20 | 20130503, 33.8 21 | 20140508, 106.3 22 | 20090615, 17.3 23 | 20090616, 6.1 24 | 20100601, 16.1 25 | 20110618, 1.2 26 | 20110630, 5.3 27 | 20140625, 18.5 28 | 20090725, 8.3 29 | 20110701, 11.2 30 | 20110717, 0.2 31 | 20120731, 9.5 32 | 20130710, 14.2 33 | 20140722, 35.7 34 | 20090817, 2 35 | 20100820, 22.1 36 | 20110831, 0.5 37 | 20120810, 7.7 38 | 20120816, 15.4 39 | 20140820, 88.8 40 | 20090916, 20.5 41 | 20120901, 11.1 42 | 20130922, 30.6 43 | 20140917, 7.7 44 | 20140930, 11.4 45 | 20101005, 2.2 46 | 20141026, 0.1 47 | 20131112, 33.4 48 | 20141129, 0.2 49 | 20091214, 1 50 | 20121218, 2.3 51 | -------------------------------------------------------------------------------- /experiments/hko/configurations/conv2d_3d/conv2d.yml: -------------------------------------------------------------------------------- 1 | SEED: 3 2 | MODEL: 3 | CNN_ACT_TYPE: leaky 4 | RNN_ACT_TYPE: leaky 5 | FRAME_STACK: 1 6 | FRAME_SKIP: 1 7 | IN_LEN: 5 8 | OUT_LEN: 20 9 | OUT_TYPE: direct 10 | NORMAL_LOSS_GLOBAL_SCALE: 5.0e-05 11 | USE_BALANCED_LOSS: true 12 | TEMPORAL_WEIGHT_TYPE: same 13 | TEMPORAL_WEIGHT_UPPER: 5 14 | L1_LAMBDA: 1.0 15 | L2_LAMBDA: 1.0 16 | GDL_LAMBDA: 0.0 17 | USE_SEASONALITY: false 18 | DECONVBASELINE: 19 | BASE_NUM_FILTER: 70 20 | USE_3D: false 21 | ENCODER: concat 22 | BN: true 23 | COMPAT: 24 | CONV_INSTEADOF_FC_IN_ENCODER: false 25 | TRAIN: 26 | BATCH_SIZE: 8 27 | TBPTT: false 28 | OPTIMIZER: adam 29 | LR: 0.0001 30 | GAMMA1: 0.9 31 | BETA1: 0.5 32 | EPS: 1.0e-08 33 | MIN_LR: 1.0e-06 34 | GRAD_CLIP: 50.0 35 | WD: 0 36 | MAX_ITER: 80000 37 | LR_DECAY_ITER: 20000 38 | LR_DECAY_FACTOR: 0.7 39 | VALID_ITER: 10000 40 | SAVE_ITER: 10000 -------------------------------------------------------------------------------- /experiments/hko/configurations/conv2d_3d/conv3d.yml: -------------------------------------------------------------------------------- 1 | SEED: 3 2 | MODEL: 3 | CNN_ACT_TYPE: leaky 4 | RNN_ACT_TYPE: leaky 5 | FRAME_STACK: 1 6 | FRAME_SKIP: 1 7 | IN_LEN: 5 8 | OUT_LEN: 20 9 | OUT_TYPE: direct 10 | NORMAL_LOSS_GLOBAL_SCALE: 5.0e-05 11 | USE_BALANCED_LOSS: true 12 | TEMPORAL_WEIGHT_TYPE: same 13 | TEMPORAL_WEIGHT_UPPER: 5 14 | L1_LAMBDA: 1.0 15 | L2_LAMBDA: 1.0 16 | GDL_LAMBDA: 0.0 17 | USE_SEASONALITY: false 18 | DECONVBASELINE: 19 | BASE_NUM_FILTER: 128 20 | USE_3D: true 21 | ENCODER: concat 22 | BN: true 23 | COMPAT: 24 | CONV_INSTEADOF_FC_IN_ENCODER: false 25 | TRAIN: 26 | BATCH_SIZE: 8 27 | TBPTT: false 28 | OPTIMIZER: adam 29 | LR: 0.0001 30 | GAMMA1: 0.9 31 | BETA1: 0.5 32 | EPS: 1.0e-08 33 | MIN_LR: 1.0e-06 34 | GRAD_CLIP: 50.0 35 | WD: 0 36 | MAX_ITER: 80000 37 | LR_DECAY_ITER: 20000 38 | LR_DECAY_FACTOR: 0.7 39 | VALID_ITER: 10000 40 | SAVE_ITER: 10000 -------------------------------------------------------------------------------- /VarFlow/Readme.txt: -------------------------------------------------------------------------------- 1 | //------------------------------------------------------------------ 2 | // Released under the BDS License 3 | // 4 | // Located at http://sourceforge.net/projects/varflow 5 | // 6 | // Author: Adam Harmat 7 | // 8 | // Feb 2010 9 | //------------------------------------------------------------------ 10 | 11 | WINDOWS VERSION 12 | 13 | The VarFlow class is an implementation of Bruhn et al's variational optical flow algorithm. An example of how to use the class can be found in example.cpp. Make sure your compiler environment is set up to link with the OpenCV libraries correctly. The VarFlow class was developed with OpenCV 1.1, but it should work with 2.0 as well. 14 | 15 | Included in this archive: 16 | 17 | VarFlow.cpp - Implementation of VarFlow 18 | VarFlow.h - Definition of VarFlow 19 | example.cpp - Example of usage 20 | Data\yos_img_08.jpg - Sample data from the Yosemite flyby sequence 21 | Data\yos_img_09.jpg - Sample data from the Yosemite flyby sequence 22 | Readme.txt - This file -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | Copyright (c) 2017-2027 Xingjian Shi and others 2 | 3 | Permission is hereby granted, free of charge, to any person obtaining 4 | a copy of this software and associated documentation files (the 5 | "Software"), to deal in the Software without restriction, including 6 | without limitation the rights to use, copy, modify, merge, publish, 7 | distribute, sublicense, and/or sell copies of the Software, and to 8 | permit persons to whom the Software is furnished to do so, subject to 9 | the following conditions: 10 | 11 | The above copyright notice and this permission notice shall be 12 | included in all copies or substantial portions of the Software. 13 | 14 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, 15 | EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF 16 | MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND 17 | NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE 18 | LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION 19 | OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION 20 | WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 21 | -------------------------------------------------------------------------------- /VarFlow/c_api.cpp: -------------------------------------------------------------------------------- 1 | #include "c_api.h" 2 | #include "VarFlow.h" 3 | #include 4 | #include 5 | 6 | 7 | void varflow(int width, int height, int max_level, int start_level, int n1, int n2, 8 | float rho, float alpha, float sigma, void* U, void* V, 9 | void* I1, void* I2) { 10 | VarFlow OpticalFlow(width, height, max_level, start_level, n1, n2, rho, alpha, sigma); 11 | IplImage* imgU = cvCreateImageHeader(cvSize(width, height), IPL_DEPTH_32F, 1); 12 | IplImage* imgV = cvCreateImageHeader(cvSize(width, height), IPL_DEPTH_32F, 1); 13 | IplImage* imgI1 = cvCreateImageHeader(cvSize(width, height), IPL_DEPTH_8U, 1); 14 | IplImage* imgI2 = cvCreateImageHeader(cvSize(width, height), IPL_DEPTH_8U, 1); 15 | cvSetData(imgU, static_cast(U), sizeof(float)*width); 16 | cvSetData(imgV, static_cast(V), sizeof(float)*width); 17 | cvSetData(imgI1, static_cast(I1), sizeof(unsigned char)*width); 18 | cvSetData(imgI2, static_cast(I2), sizeof(unsigned char)*width); 19 | OpticalFlow.CalcFlow(imgI1, imgI2, imgU, imgV, 0); 20 | } -------------------------------------------------------------------------------- /experiments/hko/last_frame_prediction.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from nowcasting.utils import logging_config 3 | from nowcasting.hko_benchmark import HKOBenchmarkEnv 4 | from nowcasting.config import cfg 5 | import os 6 | 7 | def run(pd_path=cfg.HKO_PD.RAINY_TEST, mode="fixed"): 8 | base_dir = os.path.join('hko7_benchmark', 'last_frame') 9 | logging_config(base_dir) 10 | batch_size = 1 11 | env = HKOBenchmarkEnv(pd_path=pd_path, 12 | save_dir=base_dir, 13 | mode=mode) 14 | while not env.done: 15 | in_frame_dat, in_datetime_clips, out_datetime_clips,\ 16 | begin_new_episode, need_upload_prediction =\ 17 | env.get_observation(batch_size=batch_size) 18 | if need_upload_prediction: 19 | prediction = np.zeros(shape=(cfg.HKO.BENCHMARK.OUT_LEN,) + in_frame_dat.shape[1:], 20 | dtype=in_frame_dat.dtype) 21 | prediction[:] = in_frame_dat[-1, ...] 22 | env.upload_prediction(prediction=prediction) 23 | env.print_stat_readable() 24 | env.save_eval() 25 | 26 | run(cfg.HKO_PD.RAINY_TEST, "fixed") 27 | run(cfg.HKO_PD.RAINY_VALID, "fixed") 28 | 29 | -------------------------------------------------------------------------------- /experiments/movingmnist/mnist_rnn_factory.py: -------------------------------------------------------------------------------- 1 | from nowcasting.encoder_forecaster import * 2 | from nowcasting.ops import * 3 | from nowcasting.operators import * 4 | 5 | 6 | class MovingMNISTFactory(EncoderForecasterBaseFactory): 7 | def __init__(self, 8 | batch_size, 9 | in_seq_len, 10 | out_seq_len): 11 | super(MovingMNISTFactory, self).__init__(batch_size=batch_size, 12 | in_seq_len=in_seq_len, 13 | out_seq_len=out_seq_len, 14 | height=cfg.MOVINGMNIST.IMG_SIZE, 15 | width=cfg.MOVINGMNIST.IMG_SIZE) 16 | 17 | def loss_sym(self): 18 | self.reset_all() 19 | pred = mx.sym.Variable('pred') # Shape: (out_seq_len, batch_size, 1, H, W) 20 | target = mx.sym.Variable('target') # Shape: (out_seq_len, batch_size, 1, H, W) 21 | avg_mse = mx.sym.mean(mx.sym.square(target - pred)) 22 | avg_mse = mx.sym.MakeLoss(avg_mse, 23 | name="mse") 24 | loss = mx.sym.Group([avg_mse]) 25 | return loss 26 | -------------------------------------------------------------------------------- /hko_data/README.md: -------------------------------------------------------------------------------- 1 | HKO-7 Dataset Files 2 | ------------------- 3 | The HKO-7 dataset used in the benchmark contains radar echo data from 2009 to 2015 collected by 4 | HKO. 5 | 6 | | File Name | Description | 7 | | --------- | :----------- | 8 | | hko7_rainy_train_days.txt | Stores the frame names that belong to the training set of HKO-7 data | 9 | | hko7_rainy_valid_days.txt | Stores the frame names that belong to the validation set of HKO-7 data | 10 | | hko7_rainy_test_days.txt | Stores the frame names that belong to the testing set of HKO-7 data | 11 | | intensity_day.pkl | The daily intensity of the HKO-7 data | 12 | 13 | In the pd directory, we will include all the Pandas DataFrames. The data should be downloaded by `python download_all.py`. To load it, refer to the [documentation of pandas](https://pandas.pydata.org/pandas-docs/stable/generated/pandas.read_pickle.html). 14 | 15 | | File Name | Description | 16 | | --------- | :----------- | 17 | | hko7_all.pkl | Datetimes from 09-15 | 18 | | hko7_all_09_14.pkl | Datetimes from 09-14 | 19 | | hko7_all_15.pkl | Datetimes in year 2015 | 20 | | hko7_rainy_train.pkl | Datetimes of the training set | 21 | | hko7_rainy_valid.pkl | Datetimes of the validation set | 22 | | hko7_rainy_test.pkl | Datetimes of the test set | 23 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | import os 2 | import setuptools 3 | 4 | 5 | setuptools.setup( 6 | name='nowcasting', 7 | version="0.1.dev0", 8 | author="Xingjian Shi, Zhihan Gao, Leonard Lausen, Hao Wang, Dit-Yan Yeung, Wang-chun Woo, Wai-kin Wong", 9 | author_email="xshiab@cse.ust.hk, zhihan.gao@connect.ust.hk, lelausen@connect.ust.hk, hwangaz@cse.ust.hk, dyyeung@cse.ust.hk, wcwoo@hko.gov.hk, wkwong@hko.gov.hk", 10 | packages=setuptools.find_packages(), 11 | description='Deep learning for precipitation nowcasting: A benchmark and a new model', 12 | long_description=open(os.path.join(os.path.dirname( 13 | os.path.abspath(__file__)), 'README.md')).read(), 14 | license='MIT', 15 | url='https://github.com/sxjscience/HKO-7', 16 | install_requires=['numpy', 'scipy', 'matplotlib', 'pandas', 'moviepy', 'numba', 17 | 'pillow', 'six', 'easydict', 'pyyaml'], 18 | classifiers=['Development Status :: 2 - Pre-Alpha', 19 | 'Intended Audience :: Science/Research', 20 | 'License :: OSI Approved :: MIT License', 21 | 'Natural Language :: English', 22 | 'Operating System :: OS Independent', 23 | 'Topic :: Scientific/Engineering :: Artificial Intelligence'], 24 | ) 25 | -------------------------------------------------------------------------------- /nowcasting/mask.py: -------------------------------------------------------------------------------- 1 | # File to deal read and write the .mask extensions 2 | import zlib 3 | import numpy as np 4 | from concurrent.futures import ThreadPoolExecutor, wait 5 | 6 | _executor_pool = ThreadPoolExecutor(max_workers=16) 7 | 8 | def read_mask_file(filepath, out): 9 | """Load mask file to numpy array 10 | 11 | Parameters 12 | ---------- 13 | filepath : str 14 | out : np.ndarray 15 | 16 | Returns 17 | ------- 18 | 19 | """ 20 | f = open(filepath, 'rb') 21 | dat = zlib.decompress(f.read()) 22 | out[:] = np.frombuffer(dat, dtype=bool).reshape((480, 480)) 23 | f.close() 24 | 25 | 26 | def save_mask_file(npy_mask, filepath): 27 | compressed_data = zlib.compress(npy_mask.tobytes(), 2) 28 | f = open(filepath, "wb") 29 | f.write(compressed_data) 30 | f.close() 31 | 32 | 33 | def quick_read_masks(path_list): 34 | num = len(path_list) 35 | read_storage = np.empty((num, 480, 480), dtype=np.bool) 36 | # for i in range(num): 37 | # read_storage[i] = read_mask_file(path_list[i]) 38 | future_objs = [] 39 | for i in range(num): 40 | obj = _executor_pool.submit(read_mask_file, path_list[i], read_storage[i]) 41 | future_objs.append(obj) 42 | wait(future_objs) 43 | ret = read_storage.reshape((num, 1, 480, 480)) 44 | return ret 45 | -------------------------------------------------------------------------------- /hko_data/benchmark_stat/hko7_rainy_valid_out20_stride5_online.json: -------------------------------------------------------------------------------- 1 | { 2 | "pred_seq_num": 2166, 3 | "episode_num": 45, 4 | "episode_start_datetime": [ 5 | "200903010000", 6 | "200905190000", 7 | "200906150000", 8 | "200907250000", 9 | "200909160000", 10 | "200912140000", 11 | "201001050000", 12 | "201002220000", 13 | "201003250000", 14 | "201004180000", 15 | "201006010000", 16 | "201008200000", 17 | "201010050000", 18 | "201106180000", 19 | "201106300000", 20 | "201107170000", 21 | "201108310000", 22 | "201201240000", 23 | "201202080000", 24 | "201202280000", 25 | "201204290000", 26 | "201205190000", 27 | "201207310000", 28 | "201208100000", 29 | "201208160000", 30 | "201209010000", 31 | "201212180000", 32 | "201303300000", 33 | "201304200000", 34 | "201305030000", 35 | "201307100000", 36 | "201309220000", 37 | "201311120000", 38 | "201402080000", 39 | "201402190000", 40 | "201403140000", 41 | "201404140000", 42 | "201405080000", 43 | "201406250000", 44 | "201407220000", 45 | "201408200000", 46 | "201409170000", 47 | "201409300000", 48 | "201410260000", 49 | "201411290000" 50 | ] 51 | } -------------------------------------------------------------------------------- /nowcasting/helpers/ordered_easydict.py: -------------------------------------------------------------------------------- 1 | from collections import OrderedDict 2 | 3 | 4 | class OrderedEasyDict(OrderedDict): 5 | """Using OrderedDict for the `easydict` package 6 | See Also https://pypi.python.org/pypi/easydict/ 7 | """ 8 | def __init__(self, d=None, **kwargs): 9 | super(OrderedEasyDict, self).__init__() 10 | if d is None: 11 | d = OrderedDict() 12 | if kwargs: 13 | d.update(**kwargs) 14 | for k, v in d.items(): 15 | setattr(self, k, v) 16 | # Class attributes 17 | for k in self.__class__.__dict__.keys(): 18 | if not (k.startswith('__') and k.endswith('__')): 19 | setattr(self, k, getattr(self, k)) 20 | 21 | def __setattr__(self, name, value): 22 | # special handling of self.__root and self.__map 23 | if name.startswith('_') and (name.endswith('__root') or name.endswith('__map')): 24 | super(OrderedEasyDict, self).__setattr__(name, value) 25 | else: 26 | if isinstance(value, (list, tuple)): 27 | value = [self.__class__(x) 28 | if isinstance(x, dict) else x for x in value] 29 | else: 30 | value = self.__class__(value) if isinstance(value, dict) else value 31 | super(OrderedEasyDict, self).__setattr__(name, value) 32 | super(OrderedEasyDict, self).__setitem__(name, value) 33 | 34 | __setitem__ = __setattr__ 35 | 36 | if __name__ == "__main__": 37 | import doctest 38 | doctest.testmod() -------------------------------------------------------------------------------- /hko_data/benchmark_stat/hko7_rainy_test_out20_stride5_online.json: -------------------------------------------------------------------------------- 1 | { 2 | "pred_seq_num": 6053, 3 | "episode_num": 56, 4 | "episode_start_datetime": [ 5 | "201501070000", 6 | "201501120000", 7 | "201501240000", 8 | "201502060000", 9 | "201502150000", 10 | "201502210000", 11 | "201502250000", 12 | "201502270000", 13 | "201503030000", 14 | "201503110000", 15 | "201503220000", 16 | "201503260000", 17 | "201503262200", 18 | "201504070000", 19 | "201504200000", 20 | "201505010000", 21 | "201505050000", 22 | "201505090000", 23 | "201505160000", 24 | "201505300000", 25 | "201506060000", 26 | "201506080000", 27 | "201506100000", 28 | "201506151242", 29 | "201506210000", 30 | "201506280000", 31 | "201507090000", 32 | "201507170000", 33 | "201507200000", 34 | "201507290000", 35 | "201508090000", 36 | "201508130000", 37 | "201508151724", 38 | "201508200000", 39 | "201508230006", 40 | "201508260000", 41 | "201508290000", 42 | "201509010000", 43 | "201509070000", 44 | "201509160000", 45 | "201509210000", 46 | "201509260000", 47 | "201510010000", 48 | "201510100000", 49 | "201510250000", 50 | "201510310000", 51 | "201511070000", 52 | "201511100000", 53 | "201511150000", 54 | "201512050000", 55 | "201512080000", 56 | "201512200000", 57 | "201512220000", 58 | "201512250000", 59 | "201512270000", 60 | "201512300000" 61 | ] 62 | } -------------------------------------------------------------------------------- /download_all.py: -------------------------------------------------------------------------------- 1 | from nowcasting.config import cfg 2 | import os 3 | from urllib import request 4 | import argparse 5 | import ssl 6 | ssl._create_default_https_context = ssl._create_unverified_context 7 | 8 | download_jobs =\ 9 | [[cfg.HKO_VALID_DATETIME_PATH, 10 | 'https://www.dropbox.com/s/c372wb6ciygu73i/valid_datetime.pkl?dl=1'], 11 | [cfg.HKO_PD.ALL, 12 | 'https://www.dropbox.com/s/38a2wf9pmkzef8q/hko7_all.pkl?dl=1'], 13 | [cfg.HKO_PD.ALL_09_14, 14 | 'https://www.dropbox.com/s/n7vir3vbkadyrbp/hko7_all_09_14.pkl?dl=1'], 15 | [cfg.HKO_PD.ALL_15, 16 | 'https://www.dropbox.com/s/q7cdp3g00b53fat/hko7_all_15.pkl?dl=1'], 17 | [cfg.HKO_PD.RAINY_TRAIN, 18 | 'https://www.dropbox.com/s/wutmla45tn606cl/hko7_rainy_train.pkl?dl=1'], 19 | [cfg.HKO_PD.RAINY_VALID, 20 | 'https://www.dropbox.com/s/uwumfiw9dbdg5x0/hko7_rainy_valid.pkl?dl=1'], 21 | [cfg.HKO_PD.RAINY_TEST, 22 | 'https://www.dropbox.com/s/2kzt2me55hybfqw/hko7_rainy_test.pkl?dl=1']] 23 | 24 | parser = argparse.ArgumentParser(description='Downloading the necessary data') 25 | parser.add_argument('--overwrite', dest='overwrite', action='store_true', 26 | help='Whether to overwrite the stored data files') 27 | args = parser.parse_args() 28 | 29 | for target_path, src_path in download_jobs: 30 | if not os.path.exists(target_path) or args.overwrite: 31 | print('Downloading from %s to %s...' % (src_path, target_path)) 32 | data_file = request.urlopen(src_path) 33 | with open(target_path, 'wb') as output: 34 | output.write(data_file.read()) 35 | print('Done!') 36 | else: 37 | print('Found %s' % target_path) 38 | -------------------------------------------------------------------------------- /experiments/hko/configurations/convgru_55_55_33_1_64_1_192_1_192_b4.yml: -------------------------------------------------------------------------------- 1 | CNN_ACT_TYPE: leaky 2 | RNN_ACT_TYPE: leaky 3 | FRAME_STACK: 1 4 | FRAME_SKIP: 1 5 | IN_LEN: 5 6 | OUT_LEN: 20 7 | OUT_TYPE: direct 8 | NORMAL_LOSS_GLOBAL_SCALE: 1.0e-05 9 | USE_BALANCED_LOSS: true 10 | TEMPORAL_WEIGHT_TYPE: same 11 | TEMPORAL_WEIGHT_UPPER: 5 12 | L1_LAMBDA: 1.0 13 | L2_LAMBDA: 1.0 14 | GDL_LAMBDA: 0.0 15 | USE_SEASONALITY: false 16 | ENCODER_FORECASTER: 17 | HAS_MASK: true 18 | FEATMAP_SIZE: [96, 32, 16] 19 | FIRST_CONV: [8, 7, 5, 1] 20 | LAST_DECONV: [8, 7, 5, 1] 21 | DOWNSAMPLE: 22 | - [5, 3, 1] 23 | - [3, 2, 1] 24 | UPSAMPLE: 25 | - [5, 3, 1] 26 | - [4, 2, 1] 27 | RNN_BLOCKS: 28 | RES_CONNECTION: true 29 | LAYER_TYPE: [ConvGRU, ConvGRU, ConvGRU] 30 | STACK_NUM: [1, 1, 1] 31 | NUM_FILTER: [64, 192, 192] 32 | H2H_KERNEL: 33 | - [5, 5] 34 | - [5, 5] 35 | - [3, 3] 36 | H2H_DILATE: 37 | - [1, 1] 38 | - [1, 1] 39 | - [1, 1] 40 | I2H_KERNEL: 41 | - [3, 3] 42 | - [3, 3] 43 | - [3, 3] 44 | I2H_PAD: 45 | - [1, 1] 46 | - [1, 1] 47 | - [1, 1] 48 | TRAIN: 49 | BATCH_SIZE: 4 50 | TBPTT: false 51 | OPTIMIZER: adam 52 | LR: 0.0001 53 | GAMMA1: 0.9 54 | BETA1: 0.5 55 | EPS: 1.0e-08 56 | MIN_LR: 1.0e-06 57 | GRAD_CLIP: 50.0 58 | WD: 0 59 | MAX_ITER: 100000 60 | LR_DECAY_ITER: 20000 61 | LR_DECAY_FACTOR: 0.7 62 | VALID_ITER: 10000 63 | SAVE_ITER: 10000 64 | TEST: 65 | FINETUNE: true 66 | MODE: online 67 | ONLINE: 68 | OPTIMIZER: adagrad 69 | LR: 0.001 70 | GAMMA1: 0.9 71 | BETA1: 0.5 72 | EPS: 1.0e-06 73 | GRAD_CLIP: 50.0 74 | WD: 0 75 | -------------------------------------------------------------------------------- /experiments/hko/configurations/convgru_55_55_33_1_64_1_192_1_192_nobal_b4.yml: -------------------------------------------------------------------------------- 1 | CNN_ACT_TYPE: leaky 2 | RNN_ACT_TYPE: leaky 3 | FRAME_STACK: 1 4 | FRAME_SKIP: 1 5 | IN_LEN: 5 6 | OUT_LEN: 20 7 | OUT_TYPE: direct 8 | NORMAL_LOSS_GLOBAL_SCALE: 5.0e-05 9 | USE_BALANCED_LOSS: false 10 | TEMPORAL_WEIGHT_TYPE: same 11 | TEMPORAL_WEIGHT_UPPER: 5 12 | L1_LAMBDA: 1.0 13 | L2_LAMBDA: 1.0 14 | GDL_LAMBDA: 0.0 15 | USE_SEASONALITY: false 16 | ENCODER_FORECASTER: 17 | HAS_MASK: true 18 | FEATMAP_SIZE: [96, 32, 16] 19 | FIRST_CONV: [8, 7, 5, 1] 20 | LAST_DECONV: [8, 7, 5, 1] 21 | DOWNSAMPLE: 22 | - [5, 3, 1] 23 | - [3, 2, 1] 24 | UPSAMPLE: 25 | - [5, 3, 1] 26 | - [4, 2, 1] 27 | RNN_BLOCKS: 28 | RES_CONNECTION: true 29 | LAYER_TYPE: [ConvGRU, ConvGRU, ConvGRU] 30 | STACK_NUM: [1, 1, 1] 31 | NUM_FILTER: [64, 192, 192] 32 | H2H_KERNEL: 33 | - [5, 5] 34 | - [5, 5] 35 | - [3, 3] 36 | H2H_DILATE: 37 | - [1, 1] 38 | - [1, 1] 39 | - [1, 1] 40 | I2H_KERNEL: 41 | - [3, 3] 42 | - [3, 3] 43 | - [3, 3] 44 | I2H_PAD: 45 | - [1, 1] 46 | - [1, 1] 47 | - [1, 1] 48 | TRAIN: 49 | BATCH_SIZE: 4 50 | TBPTT: false 51 | OPTIMIZER: adam 52 | LR: 0.0001 53 | GAMMA1: 0.9 54 | BETA1: 0.5 55 | EPS: 1.0e-08 56 | MIN_LR: 1.0e-06 57 | GRAD_CLIP: 50.0 58 | WD: 0 59 | MAX_ITER: 100000 60 | LR_DECAY_ITER: 20000 61 | LR_DECAY_FACTOR: 0.7 62 | VALID_ITER: 10000 63 | SAVE_ITER: 10000 64 | TEST: 65 | FINETUNE: true 66 | MODE: online 67 | ONLINE: 68 | OPTIMIZER: adagrad 69 | LR: 0.001 70 | GAMMA1: 0.9 71 | BETA1: 0.5 72 | EPS: 1.0e-06 73 | GRAD_CLIP: 50.0 74 | WD: 0 75 | -------------------------------------------------------------------------------- /experiments/hko/configurations/trajgru_55_55_33_1_64_1_192_1_192_13_13_9_b4.yml: -------------------------------------------------------------------------------- 1 | CNN_ACT_TYPE: leaky 2 | RNN_ACT_TYPE: leaky 3 | FRAME_STACK: 1 4 | FRAME_SKIP: 1 5 | IN_LEN: 5 6 | OUT_LEN: 20 7 | OUT_TYPE: direct 8 | NORMAL_LOSS_GLOBAL_SCALE: 1.0e-05 9 | USE_BALANCED_LOSS: true 10 | TEMPORAL_WEIGHT_TYPE: same 11 | TEMPORAL_WEIGHT_UPPER: 5 12 | L1_LAMBDA: 1.0 13 | L2_LAMBDA: 1.0 14 | GDL_LAMBDA: 0.0 15 | USE_SEASONALITY: false 16 | ENCODER_FORECASTER: 17 | HAS_MASK: true 18 | FEATMAP_SIZE: [96, 32, 16] 19 | FIRST_CONV: [8, 7, 5, 1] 20 | LAST_DECONV: [8, 7, 5, 1] 21 | DOWNSAMPLE: 22 | - [5, 3, 1] 23 | - [3, 2, 1] 24 | UPSAMPLE: 25 | - [5, 3, 1] 26 | - [4, 2, 1] 27 | RNN_BLOCKS: 28 | RES_CONNECTION: true 29 | LAYER_TYPE: [TrajGRU, TrajGRU, TrajGRU] 30 | STACK_NUM: [1, 1, 1] 31 | NUM_FILTER: [64, 192, 192] 32 | H2H_KERNEL: 33 | - [5, 5] 34 | - [5, 5] 35 | - [3, 3] 36 | H2H_DILATE: 37 | - [1, 1] 38 | - [1, 1] 39 | - [1, 1] 40 | I2H_KERNEL: 41 | - [3, 3] 42 | - [3, 3] 43 | - [3, 3] 44 | I2H_PAD: 45 | - [1, 1] 46 | - [1, 1] 47 | - [1, 1] 48 | L: [13, 13, 9] 49 | TRAIN: 50 | BATCH_SIZE: 4 51 | TBPTT: false 52 | OPTIMIZER: adam 53 | LR: 0.0001 54 | GAMMA1: 0.9 55 | BETA1: 0.5 56 | EPS: 1.0e-08 57 | MIN_LR: 1.0e-06 58 | GRAD_CLIP: 50.0 59 | WD: 0 60 | MAX_ITER: 100000 61 | LR_DECAY_ITER: 20000 62 | LR_DECAY_FACTOR: 0.7 63 | VALID_ITER: 10000 64 | SAVE_ITER: 10000 65 | TEST: 66 | FINETUNE: true 67 | MODE: online 68 | ONLINE: 69 | OPTIMIZER: adagrad 70 | LR: 0.001 71 | GAMMA1: 0.9 72 | BETA1: 0.5 73 | EPS: 1.0e-06 74 | GRAD_CLIP: 50.0 75 | WD: 0 76 | -------------------------------------------------------------------------------- /experiments/movingmnist/analyze_learned_structure.py: -------------------------------------------------------------------------------- 1 | from nowcasting.config import cfg 2 | import os 3 | import numpy as np 4 | import matplotlib.pyplot as plt 5 | base_dir = "flows" 6 | prefixs = ["ebrnn1_0", "ebrnn2_0", "ebrnn3_0", "fbrnn1_0", "fbrnn2_0", "fbrnn3_0"] 7 | forecaster_block_prefix = [] 8 | flow_num = 13 9 | lengths = [10, 10, 10, 10, 10, 10] 10 | heights = [64, 32, 16, 64, 32, 16] 11 | widths = [64, 32, 16, 64, 32, 16] 12 | flows = [] 13 | for prefix, height, width, length in zip(prefixs, heights, widths, lengths): 14 | flow_maps = np.empty((length, flow_num, 2, height, width), dtype=np.float32) 15 | X, Y = np.meshgrid(np.arange(0, width), np.arange(0, height)) 16 | for frame_id in range(length): 17 | path = os.path.join(base_dir, "{}__t{}_flow.npy".format(prefix, frame_id)) 18 | flow = np.load(path) 19 | flow_maps[frame_id, :, :, :, :] = flow.reshape((flow_num, 2, height, width)) 20 | for i in range(flow_num): 21 | plt.title('Flow') 22 | Q = plt.quiver(X, Y[::-1, :], 23 | flow_maps[frame_id, i, 0, :, :], 24 | - flow_maps[frame_id, i, 1, :, :], 25 | units='width') 26 | save_dir = os.path.join(base_dir, "%s_link%d" % (prefix, i)) 27 | if not os.path.exists(save_dir): 28 | os.makedirs(save_dir) 29 | save_path = os.path.join(save_dir, "frame%d.png" % (frame_id,)) 30 | print("Generating the flow for %s, frame_id=%d, flow_id=%d, saving to %s" 31 | %(prefix, frame_id, i, save_path)) 32 | plt.savefig(save_path, 33 | bbox_inches="tight") 34 | plt.close() 35 | flows.append(flow_maps) 36 | print(flows[3].std(axis=2)) 37 | np.savez('trajgru_flows.npz', **dict([[prefix, flow] 38 | for prefix, flow in zip(prefixs, flows)])) -------------------------------------------------------------------------------- /experiments/movingmnist/configurations/convgru_1_64_1_96_1_96_K5.yml: -------------------------------------------------------------------------------- 1 | MOVINGMNIST: 2 | DISTRACTOR_NUM: 0 3 | VELOCITY_LOWER: 0.0 4 | VELOCITY_UPPER: 3.6 5 | SCALE_VARIATION_LOWER: 0.9090909090909091 6 | SCALE_VARIATION_UPPER: 1.1 7 | ROTATION_LOWER: -30 8 | ROTATION_UPPER: 30 9 | ILLUMINATION_LOWER: 0.6 10 | ILLUMINATION_UPPER: 1.0 11 | DIGIT_NUM: 3 12 | IN_LEN: 10 13 | OUT_LEN: 10 14 | IMG_SIZE: 64 15 | MODEL: 16 | CNN_ACT_TYPE: leaky 17 | RNN_ACT_TYPE: leaky 18 | FRAME_STACK: 1 19 | FRAME_SKIP: 1 20 | IN_LEN: 10 21 | OUT_LEN: 10 22 | OUT_TYPE: direct 23 | ENCODER_FORECASTER: 24 | HAS_MASK: false 25 | FEATMAP_SIZE: [64, 32, 16] 26 | FIRST_CONV: [16, 3, 1, 1] 27 | LAST_DECONV: [16, 3, 1, 1] 28 | DOWNSAMPLE: 29 | - [3, 2, 1] 30 | - [3, 2, 1] 31 | UPSAMPLE: 32 | - [4, 2, 1] 33 | - [4, 2, 1] 34 | RNN_BLOCKS: 35 | RES_CONNECTION: true 36 | LAYER_TYPE: [ConvGRU, ConvGRU, ConvGRU] 37 | STACK_NUM: [1, 1, 1] 38 | NUM_FILTER: [64, 96, 96] 39 | H2H_KERNEL: 40 | - [5, 5] 41 | - [5, 5] 42 | - [5, 5] 43 | H2H_DILATE: 44 | - [1, 1] 45 | - [1, 1] 46 | - [1, 1] 47 | I2H_KERNEL: 48 | - [3, 3] 49 | - [3, 3] 50 | - [3, 3] 51 | I2H_PAD: 52 | - [1, 1] 53 | - [1, 1] 54 | - [1, 1] 55 | L: [5, 5, 5] 56 | TRAIN: 57 | BATCH_SIZE: 4 58 | TBPTT: false 59 | OPTIMIZER: adam 60 | LR: 0.0001 61 | GAMMA1: 0.9 62 | BETA1: 0.5 63 | EPS: 1.0e-08 64 | MIN_LR: 1.0e-06 65 | GRAD_CLIP: 1.0 66 | WD: 0 67 | MAX_ITER: 200000 68 | LR_DECAY_ITER: 20000 69 | LR_DECAY_FACTOR: 0.7 70 | VALID_ITER: 5000 71 | SAVE_ITER: 20000 72 | TEST: 73 | FINETUNE: true 74 | MODE: online 75 | ONLINE: 76 | OPTIMIZER: adam 77 | LR: 1.0e-06 78 | GAMMA1: 0.9 79 | BETA1: 0.5 80 | EPS: 1.0e-08 81 | GRAD_CLIP: 10.0 82 | WD: 0 83 | -------------------------------------------------------------------------------- /experiments/movingmnist/configurations/convgru_1_64_1_96_1_96_K7.yml: -------------------------------------------------------------------------------- 1 | MOVINGMNIST: 2 | DISTRACTOR_NUM: 0 3 | VELOCITY_LOWER: 0.0 4 | VELOCITY_UPPER: 3.6 5 | SCALE_VARIATION_LOWER: 0.9090909090909091 6 | SCALE_VARIATION_UPPER: 1.1 7 | ROTATION_LOWER: -30 8 | ROTATION_UPPER: 30 9 | ILLUMINATION_LOWER: 0.6 10 | ILLUMINATION_UPPER: 1.0 11 | DIGIT_NUM: 3 12 | IN_LEN: 10 13 | OUT_LEN: 10 14 | IMG_SIZE: 64 15 | MODEL: 16 | CNN_ACT_TYPE: leaky 17 | RNN_ACT_TYPE: leaky 18 | FRAME_STACK: 1 19 | FRAME_SKIP: 1 20 | IN_LEN: 10 21 | OUT_LEN: 10 22 | OUT_TYPE: direct 23 | ENCODER_FORECASTER: 24 | HAS_MASK: false 25 | FEATMAP_SIZE: [64, 32, 16] 26 | FIRST_CONV: [16, 3, 1, 1] 27 | LAST_DECONV: [16, 3, 1, 1] 28 | DOWNSAMPLE: 29 | - [3, 2, 1] 30 | - [3, 2, 1] 31 | UPSAMPLE: 32 | - [4, 2, 1] 33 | - [4, 2, 1] 34 | RNN_BLOCKS: 35 | RES_CONNECTION: true 36 | LAYER_TYPE: [ConvGRU, ConvGRU, ConvGRU] 37 | STACK_NUM: [1, 1, 1] 38 | NUM_FILTER: [64, 96, 96] 39 | H2H_KERNEL: 40 | - [7, 7] 41 | - [7, 7] 42 | - [7, 7] 43 | H2H_DILATE: 44 | - [1, 1] 45 | - [1, 1] 46 | - [1, 1] 47 | I2H_KERNEL: 48 | - [3, 3] 49 | - [3, 3] 50 | - [3, 3] 51 | I2H_PAD: 52 | - [1, 1] 53 | - [1, 1] 54 | - [1, 1] 55 | L: [5, 5, 5] 56 | TRAIN: 57 | BATCH_SIZE: 4 58 | TBPTT: false 59 | OPTIMIZER: adam 60 | LR: 0.0001 61 | GAMMA1: 0.9 62 | BETA1: 0.5 63 | EPS: 1.0e-08 64 | MIN_LR: 1.0e-06 65 | GRAD_CLIP: 1.0 66 | WD: 0 67 | MAX_ITER: 200000 68 | LR_DECAY_ITER: 20000 69 | LR_DECAY_FACTOR: 0.7 70 | VALID_ITER: 5000 71 | SAVE_ITER: 20000 72 | TEST: 73 | FINETUNE: true 74 | MODE: online 75 | ONLINE: 76 | OPTIMIZER: adam 77 | LR: 1.0e-06 78 | GAMMA1: 0.9 79 | BETA1: 0.5 80 | EPS: 1.0e-08 81 | GRAD_CLIP: 10.0 82 | WD: 0 83 | -------------------------------------------------------------------------------- /experiments/movingmnist/configurations/convgru_1_64_1_96_1_96_K3D2.yml: -------------------------------------------------------------------------------- 1 | MOVINGMNIST: 2 | DISTRACTOR_NUM: 0 3 | VELOCITY_LOWER: 0.0 4 | VELOCITY_UPPER: 3.6 5 | SCALE_VARIATION_LOWER: 0.9090909090909091 6 | SCALE_VARIATION_UPPER: 1.1 7 | ROTATION_LOWER: -30 8 | ROTATION_UPPER: 30 9 | ILLUMINATION_LOWER: 0.6 10 | ILLUMINATION_UPPER: 1.0 11 | DIGIT_NUM: 3 12 | IN_LEN: 10 13 | OUT_LEN: 10 14 | IMG_SIZE: 64 15 | MODEL: 16 | CNN_ACT_TYPE: leaky 17 | RNN_ACT_TYPE: leaky 18 | FRAME_STACK: 1 19 | FRAME_SKIP: 1 20 | IN_LEN: 10 21 | OUT_LEN: 10 22 | OUT_TYPE: direct 23 | ENCODER_FORECASTER: 24 | HAS_MASK: false 25 | FEATMAP_SIZE: [64, 32, 16] 26 | FIRST_CONV: [16, 3, 1, 1] 27 | LAST_DECONV: [16, 3, 1, 1] 28 | DOWNSAMPLE: 29 | - [3, 2, 1] 30 | - [3, 2, 1] 31 | UPSAMPLE: 32 | - [4, 2, 1] 33 | - [4, 2, 1] 34 | RNN_BLOCKS: 35 | RES_CONNECTION: true 36 | LAYER_TYPE: [ConvGRU, ConvGRU, ConvGRU] 37 | STACK_NUM: [1, 1, 1] 38 | NUM_FILTER: [64, 96, 96] 39 | H2H_KERNEL: 40 | - [3, 3] 41 | - [3, 3] 42 | - [3, 3] 43 | H2H_DILATE: 44 | - [2, 2] 45 | - [2, 2] 46 | - [2, 2] 47 | I2H_KERNEL: 48 | - [3, 3] 49 | - [3, 3] 50 | - [3, 3] 51 | I2H_PAD: 52 | - [1, 1] 53 | - [1, 1] 54 | - [1, 1] 55 | L: [5, 5, 5] 56 | TRAIN: 57 | BATCH_SIZE: 4 58 | TBPTT: false 59 | OPTIMIZER: adam 60 | LR: 0.0001 61 | GAMMA1: 0.9 62 | BETA1: 0.5 63 | EPS: 1.0e-08 64 | MIN_LR: 1.0e-06 65 | GRAD_CLIP: 1.0 66 | WD: 0 67 | MAX_ITER: 200000 68 | LR_DECAY_ITER: 20000 69 | LR_DECAY_FACTOR: 0.7 70 | VALID_ITER: 5000 71 | SAVE_ITER: 20000 72 | TEST: 73 | FINETUNE: true 74 | MODE: online 75 | ONLINE: 76 | OPTIMIZER: adam 77 | LR: 1.0e-06 78 | GAMMA1: 0.9 79 | BETA1: 0.5 80 | EPS: 1.0e-08 81 | GRAD_CLIP: 10.0 82 | WD: 0 83 | -------------------------------------------------------------------------------- /experiments/movingmnist/configurations/convgru_1_64_1_96_1_96_K5_DFN.yml: -------------------------------------------------------------------------------- 1 | MOVINGMNIST: 2 | DISTRACTOR_NUM: 0 3 | VELOCITY_LOWER: 0.0 4 | VELOCITY_UPPER: 3.6 5 | SCALE_VARIATION_LOWER: 0.9090909090909091 6 | SCALE_VARIATION_UPPER: 1.1 7 | ROTATION_LOWER: -30 8 | ROTATION_UPPER: 30 9 | ILLUMINATION_LOWER: 0.6 10 | ILLUMINATION_UPPER: 1.0 11 | DIGIT_NUM: 3 12 | IN_LEN: 10 13 | OUT_LEN: 10 14 | IMG_SIZE: 64 15 | MODEL: 16 | CNN_ACT_TYPE: leaky 17 | RNN_ACT_TYPE: leaky 18 | FRAME_STACK: 1 19 | FRAME_SKIP: 1 20 | IN_LEN: 10 21 | OUT_LEN: 10 22 | OUT_TYPE: DFN 23 | ENCODER_FORECASTER: 24 | HAS_MASK: false 25 | FEATMAP_SIZE: [64, 32, 16] 26 | FIRST_CONV: [16, 3, 1, 1] 27 | LAST_DECONV: [16, 3, 1, 1] 28 | DOWNSAMPLE: 29 | - [3, 2, 1] 30 | - [3, 2, 1] 31 | UPSAMPLE: 32 | - [4, 2, 1] 33 | - [4, 2, 1] 34 | RNN_BLOCKS: 35 | RES_CONNECTION: true 36 | LAYER_TYPE: [ConvGRU, ConvGRU, ConvGRU] 37 | STACK_NUM: [1, 1, 1] 38 | NUM_FILTER: [64, 96, 96] 39 | H2H_KERNEL: 40 | - [5, 5] 41 | - [5, 5] 42 | - [5, 5] 43 | H2H_DILATE: 44 | - [1, 1] 45 | - [1, 1] 46 | - [1, 1] 47 | I2H_KERNEL: 48 | - [3, 3] 49 | - [3, 3] 50 | - [3, 3] 51 | I2H_PAD: 52 | - [1, 1] 53 | - [1, 1] 54 | - [1, 1] 55 | L: [5, 5, 5] 56 | TRAIN: 57 | BATCH_SIZE: 4 58 | TBPTT: false 59 | OPTIMIZER: adam 60 | LR: 0.0001 61 | GAMMA1: 0.9 62 | BETA1: 0.5 63 | EPS: 1.0e-08 64 | MIN_LR: 1.0e-06 65 | GRAD_CLIP: 1.0 66 | WD: 0 67 | MAX_ITER: 200000 68 | LR_DECAY_ITER: 20000 69 | LR_DECAY_FACTOR: 0.7 70 | VALID_ITER: 5000 71 | SAVE_ITER: 20000 72 | TEST: 73 | FINETUNE: true 74 | MODE: online 75 | ONLINE: 76 | OPTIMIZER: adam 77 | LR: 1.0e-06 78 | GAMMA1: 0.9 79 | BETA1: 0.5 80 | EPS: 1.0e-08 81 | GRAD_CLIP: 10.0 82 | WD: 0 83 | -------------------------------------------------------------------------------- /nowcasting/helpers/log_analysis.py: -------------------------------------------------------------------------------- 1 | import os 2 | import pandas as pd 3 | import matplotlib.pyplot as plt 4 | import re 5 | from collections import OrderedDict 6 | import numpy as np 7 | 8 | def remove_duplicates_and_convert_npy(val_list): 9 | tmp_dist = OrderedDict() 10 | for ele in val_list: 11 | tmp_dist[ele[0]] = ele[1:] 12 | val_list = [] 13 | for k, v in tmp_dist.items(): 14 | val_list.append(((k,) + v)) 15 | ret_npy = np.zeros((len(val_list), len(val_list[0])), dtype=np.float32) 16 | for i, ele_tuple in enumerate(val_list): 17 | for j in range(len(ele_tuple)): 18 | ret_npy[i, j] = float(ele_tuple[j]) 19 | return ret_npy 20 | 21 | 22 | def temporal_smoothing(training_statistics, stride=10, window_size=100): 23 | """We always assume the first axis in statistics is the iteration 24 | 25 | Parameters 26 | ---------- 27 | training_statistics 28 | stride 29 | window_size 30 | 31 | Returns 32 | ------- 33 | smoothed_mean: 34 | smoothed_std: 35 | """ 36 | slice_obj = slice(window_size-1, None, stride) 37 | iter_slice = training_statistics[slice_obj, 0:1] 38 | rolling_frame = pd.DataFrame(training_statistics[:, 1:]).rolling(window=window_size, center=False) 39 | smoothed_mean = rolling_frame.mean().as_matrix()[slice_obj, :] 40 | smoothed_std = rolling_frame.std().as_matrix()[slice_obj, :] 41 | smoothed_mean = np.concatenate([iter_slice, smoothed_mean], axis=1) 42 | smoothed_std = np.concatenate([iter_slice, smoothed_std], axis=1) 43 | return smoothed_mean, smoothed_std 44 | 45 | 46 | def parse_log(file_path, regex): 47 | """ 48 | 49 | Parameters 50 | ---------- 51 | file_path 52 | regex 53 | 54 | Returns 55 | ------- 56 | 57 | """ 58 | with open(file_path) as f: 59 | content = f.read() 60 | ret = re.findall(regex, content) 61 | ret = remove_duplicates_and_convert_npy(ret) 62 | return ret 63 | -------------------------------------------------------------------------------- /experiments/movingmnist/configurations/trajgru_1_64_1_96_1_96_L13.yml: -------------------------------------------------------------------------------- 1 | MOVINGMNIST: 2 | DISTRACTOR_NUM: 0 3 | VELOCITY_LOWER: 0.0 4 | VELOCITY_UPPER: 3.6 5 | SCALE_VARIATION_LOWER: 0.9090909090909091 6 | SCALE_VARIATION_UPPER: 1.1 7 | ROTATION_LOWER: -30 8 | ROTATION_UPPER: 30 9 | ILLUMINATION_LOWER: 0.6 10 | ILLUMINATION_UPPER: 1.0 11 | DIGIT_NUM: 3 12 | IN_LEN: 10 13 | OUT_LEN: 10 14 | IMG_SIZE: 64 15 | MODEL: 16 | CNN_ACT_TYPE: leaky 17 | RNN_ACT_TYPE: leaky 18 | FRAME_STACK: 1 19 | FRAME_SKIP: 1 20 | IN_LEN: 10 21 | OUT_LEN: 10 22 | OUT_TYPE: direct 23 | TRAJRNN: 24 | INIT_GRID: true 25 | ENCODER_FORECASTER: 26 | HAS_MASK: false 27 | FEATMAP_SIZE: [64, 32, 16] 28 | FIRST_CONV: [16, 3, 1, 1] 29 | LAST_DECONV: [16, 3, 1, 1] 30 | DOWNSAMPLE: 31 | - [3, 2, 1] 32 | - [3, 2, 1] 33 | UPSAMPLE: 34 | - [4, 2, 1] 35 | - [4, 2, 1] 36 | RNN_BLOCKS: 37 | RES_CONNECTION: true 38 | LAYER_TYPE: [TrajGRU, TrajGRU, TrajGRU] 39 | STACK_NUM: [1, 1, 1] 40 | NUM_FILTER: [64, 96, 96] 41 | H2H_KERNEL: 42 | - [5, 5] 43 | - [5, 5] 44 | - [5, 5] 45 | H2H_DILATE: 46 | - [1, 1] 47 | - [1, 1] 48 | - [1, 1] 49 | I2H_KERNEL: 50 | - [3, 3] 51 | - [3, 3] 52 | - [3, 3] 53 | I2H_PAD: 54 | - [1, 1] 55 | - [1, 1] 56 | - [1, 1] 57 | L: [13, 13, 13] 58 | TRAIN: 59 | BATCH_SIZE: 4 60 | TBPTT: false 61 | OPTIMIZER: adam 62 | LR: 0.0001 63 | GAMMA1: 0.9 64 | BETA1: 0.5 65 | EPS: 1.0e-08 66 | MIN_LR: 1.0e-06 67 | GRAD_CLIP: 1.0 68 | WD: 0 69 | MAX_ITER: 200000 70 | LR_DECAY_ITER: 20000 71 | LR_DECAY_FACTOR: 0.7 72 | VALID_ITER: 5000 73 | SAVE_ITER: 20000 74 | TEST: 75 | FINETUNE: true 76 | MODE: online 77 | ONLINE: 78 | OPTIMIZER: adam 79 | LR: 1.0e-06 80 | GAMMA1: 0.9 81 | BETA1: 0.5 82 | EPS: 1.0e-08 83 | GRAD_CLIP: 10.0 84 | WD: 0 85 | -------------------------------------------------------------------------------- /experiments/movingmnist/configurations/trajgru_1_64_1_96_1_96_L17.yml: -------------------------------------------------------------------------------- 1 | MOVINGMNIST: 2 | DISTRACTOR_NUM: 0 3 | VELOCITY_LOWER: 0.0 4 | VELOCITY_UPPER: 3.6 5 | SCALE_VARIATION_LOWER: 0.9090909090909091 6 | SCALE_VARIATION_UPPER: 1.1 7 | ROTATION_LOWER: -30 8 | ROTATION_UPPER: 30 9 | ILLUMINATION_LOWER: 0.6 10 | ILLUMINATION_UPPER: 1.0 11 | DIGIT_NUM: 3 12 | IN_LEN: 10 13 | OUT_LEN: 10 14 | IMG_SIZE: 64 15 | MODEL: 16 | CNN_ACT_TYPE: leaky 17 | RNN_ACT_TYPE: leaky 18 | FRAME_STACK: 1 19 | FRAME_SKIP: 1 20 | IN_LEN: 10 21 | OUT_LEN: 10 22 | OUT_TYPE: direct 23 | TRAJRNN: 24 | INIT_GRID: true 25 | ENCODER_FORECASTER: 26 | HAS_MASK: false 27 | FEATMAP_SIZE: [64, 32, 16] 28 | FIRST_CONV: [16, 3, 1, 1] 29 | LAST_DECONV: [16, 3, 1, 1] 30 | DOWNSAMPLE: 31 | - [3, 2, 1] 32 | - [3, 2, 1] 33 | UPSAMPLE: 34 | - [4, 2, 1] 35 | - [4, 2, 1] 36 | RNN_BLOCKS: 37 | RES_CONNECTION: true 38 | LAYER_TYPE: [TrajGRU, TrajGRU, TrajGRU] 39 | STACK_NUM: [1, 1, 1] 40 | NUM_FILTER: [64, 96, 96] 41 | H2H_KERNEL: 42 | - [5, 5] 43 | - [5, 5] 44 | - [5, 5] 45 | H2H_DILATE: 46 | - [1, 1] 47 | - [1, 1] 48 | - [1, 1] 49 | I2H_KERNEL: 50 | - [3, 3] 51 | - [3, 3] 52 | - [3, 3] 53 | I2H_PAD: 54 | - [1, 1] 55 | - [1, 1] 56 | - [1, 1] 57 | L: [17, 17, 17] 58 | TRAIN: 59 | BATCH_SIZE: 4 60 | TBPTT: false 61 | OPTIMIZER: adam 62 | LR: 0.0001 63 | GAMMA1: 0.9 64 | BETA1: 0.5 65 | EPS: 1.0e-08 66 | MIN_LR: 1.0e-06 67 | GRAD_CLIP: 1.0 68 | WD: 0 69 | MAX_ITER: 200000 70 | LR_DECAY_ITER: 20000 71 | LR_DECAY_FACTOR: 0.7 72 | VALID_ITER: 5000 73 | SAVE_ITER: 20000 74 | TEST: 75 | FINETUNE: true 76 | MODE: online 77 | ONLINE: 78 | OPTIMIZER: adam 79 | LR: 1.0e-06 80 | GAMMA1: 0.9 81 | BETA1: 0.5 82 | EPS: 1.0e-08 83 | GRAD_CLIP: 10.0 84 | WD: 0 85 | -------------------------------------------------------------------------------- /experiments/movingmnist/configurations/trajgru_1_64_1_96_1_96_L5.yml: -------------------------------------------------------------------------------- 1 | MOVINGMNIST: 2 | DISTRACTOR_NUM: 0 3 | VELOCITY_LOWER: 0.0 4 | VELOCITY_UPPER: 3.6 5 | SCALE_VARIATION_LOWER: 0.9090909090909091 6 | SCALE_VARIATION_UPPER: 1.1 7 | ROTATION_LOWER: -30 8 | ROTATION_UPPER: 30 9 | ILLUMINATION_LOWER: 0.6 10 | ILLUMINATION_UPPER: 1.0 11 | DIGIT_NUM: 3 12 | IN_LEN: 10 13 | OUT_LEN: 10 14 | IMG_SIZE: 64 15 | MODEL: 16 | CNN_ACT_TYPE: leaky 17 | RNN_ACT_TYPE: leaky 18 | FRAME_STACK: 1 19 | FRAME_SKIP: 1 20 | IN_LEN: 10 21 | OUT_LEN: 10 22 | OUT_TYPE: direct 23 | TRAJRNN: 24 | INIT_GRID: true 25 | ENCODER_FORECASTER: 26 | HAS_MASK: false 27 | FEATMAP_SIZE: [64, 32, 16] 28 | FIRST_CONV: [16, 3, 1, 1] 29 | LAST_DECONV: [16, 3, 1, 1] 30 | DOWNSAMPLE: 31 | - [3, 2, 1] 32 | - [3, 2, 1] 33 | UPSAMPLE: 34 | - [4, 2, 1] 35 | - [4, 2, 1] 36 | RNN_BLOCKS: 37 | RES_CONNECTION: true 38 | LAYER_TYPE: [TrajGRU, TrajGRU, TrajGRU] 39 | STACK_NUM: [1, 1, 1] 40 | NUM_FILTER: [64, 96, 96] 41 | H2H_KERNEL: 42 | - [5, 5] 43 | - [5, 5] 44 | - [5, 5] 45 | H2H_DILATE: 46 | - [1, 1] 47 | - [1, 1] 48 | - [1, 1] 49 | I2H_KERNEL: 50 | - [3, 3] 51 | - [3, 3] 52 | - [3, 3] 53 | I2H_PAD: 54 | - [1, 1] 55 | - [1, 1] 56 | - [1, 1] 57 | L: [5, 5, 5] 58 | TRAIN: 59 | BATCH_SIZE: 4 60 | TBPTT: false 61 | OPTIMIZER: adam 62 | LR: 0.0001 63 | GAMMA1: 0.9 64 | BETA1: 0.5 65 | EPS: 1.0e-08 66 | MIN_LR: 1.0e-06 67 | GRAD_CLIP: 1.0 68 | WD: 0 69 | MAX_ITER: 200000 70 | LR_DECAY_ITER: 20000 71 | LR_DECAY_FACTOR: 0.7 72 | VALID_ITER: 5000 73 | SAVE_ITER: 20000 74 | TEST: 75 | FINETUNE: true 76 | MODE: online 77 | ONLINE: 78 | OPTIMIZER: adam 79 | LR: 1.0e-06 80 | GAMMA1: 0.9 81 | BETA1: 0.5 82 | EPS: 1.0e-08 83 | GRAD_CLIP: 10.0 84 | WD: 0 85 | -------------------------------------------------------------------------------- /experiments/movingmnist/configurations/trajgru_1_64_1_96_1_96_L9.yml: -------------------------------------------------------------------------------- 1 | MOVINGMNIST: 2 | DISTRACTOR_NUM: 0 3 | VELOCITY_LOWER: 0.0 4 | VELOCITY_UPPER: 3.6 5 | SCALE_VARIATION_LOWER: 0.9090909090909091 6 | SCALE_VARIATION_UPPER: 1.1 7 | ROTATION_LOWER: -30 8 | ROTATION_UPPER: 30 9 | ILLUMINATION_LOWER: 0.6 10 | ILLUMINATION_UPPER: 1.0 11 | DIGIT_NUM: 3 12 | IN_LEN: 10 13 | OUT_LEN: 10 14 | IMG_SIZE: 64 15 | MODEL: 16 | CNN_ACT_TYPE: leaky 17 | RNN_ACT_TYPE: leaky 18 | FRAME_STACK: 1 19 | FRAME_SKIP: 1 20 | IN_LEN: 10 21 | OUT_LEN: 10 22 | OUT_TYPE: direct 23 | TRAJRNN: 24 | INIT_GRID: true 25 | ENCODER_FORECASTER: 26 | HAS_MASK: false 27 | FEATMAP_SIZE: [64, 32, 16] 28 | FIRST_CONV: [16, 3, 1, 1] 29 | LAST_DECONV: [16, 3, 1, 1] 30 | DOWNSAMPLE: 31 | - [3, 2, 1] 32 | - [3, 2, 1] 33 | UPSAMPLE: 34 | - [4, 2, 1] 35 | - [4, 2, 1] 36 | RNN_BLOCKS: 37 | RES_CONNECTION: true 38 | LAYER_TYPE: [TrajGRU, TrajGRU, TrajGRU] 39 | STACK_NUM: [1, 1, 1] 40 | NUM_FILTER: [64, 96, 96] 41 | H2H_KERNEL: 42 | - [5, 5] 43 | - [5, 5] 44 | - [5, 5] 45 | H2H_DILATE: 46 | - [1, 1] 47 | - [1, 1] 48 | - [1, 1] 49 | I2H_KERNEL: 50 | - [3, 3] 51 | - [3, 3] 52 | - [3, 3] 53 | I2H_PAD: 54 | - [1, 1] 55 | - [1, 1] 56 | - [1, 1] 57 | L: [9, 9, 9] 58 | TRAIN: 59 | BATCH_SIZE: 4 60 | TBPTT: false 61 | OPTIMIZER: adam 62 | LR: 0.0001 63 | GAMMA1: 0.9 64 | BETA1: 0.5 65 | EPS: 1.0e-08 66 | MIN_LR: 1.0e-06 67 | GRAD_CLIP: 1.0 68 | WD: 0 69 | MAX_ITER: 200000 70 | LR_DECAY_ITER: 20000 71 | LR_DECAY_FACTOR: 0.7 72 | VALID_ITER: 5000 73 | SAVE_ITER: 20000 74 | TEST: 75 | FINETUNE: true 76 | MODE: online 77 | ONLINE: 78 | OPTIMIZER: adam 79 | LR: 1.0e-06 80 | GAMMA1: 0.9 81 | BETA1: 0.5 82 | EPS: 1.0e-08 83 | GRAD_CLIP: 10.0 84 | WD: 0 85 | -------------------------------------------------------------------------------- /VarFlow/ProfTimer.h: -------------------------------------------------------------------------------- 1 | //------------------------------------------------------------------------------------------------------ 2 | // ProfTimer class 3 | // 4 | // Performs high resolution timing on Windows and Linux platforms 5 | // 6 | // Windows code source: http://blog.kowalczyk.info/article/High-resolution-timer-for-timing-code-fragments.html 7 | // 8 | // Windows author: Krzysztof J. Kowalczyk (presumed based on website) 9 | // 10 | //------------------------------------------------------------------------------------------------------ 11 | 12 | #ifndef PROFTIMER 13 | #define PROFTIMER 14 | 15 | #ifdef _WIN32 16 | 17 | #include 18 | 19 | class ProfTimer { 20 | 21 | public: 22 | 23 | ProfTimer(){ 24 | }; 25 | 26 | void Start(void) { 27 | QueryPerformanceCounter(&mTimeStart); 28 | }; 29 | 30 | void Stop(void) { 31 | QueryPerformanceCounter(&mTimeStop); 32 | }; 33 | 34 | float GetDurationInSecs(void) 35 | { 36 | LARGE_INTEGER freq; 37 | QueryPerformanceFrequency(&freq); 38 | float duration = (float)(mTimeStop.QuadPart-mTimeStart.QuadPart)/(float)freq.QuadPart; 39 | return duration; 40 | } 41 | 42 | private: 43 | 44 | LARGE_INTEGER mTimeStart; 45 | LARGE_INTEGER mTimeStop; 46 | 47 | }; 48 | 49 | #else 50 | 51 | #include 52 | 53 | class ProfTimer { 54 | 55 | public: 56 | 57 | ProfTimer(){ 58 | }; 59 | 60 | void Start(void) { 61 | gettimeofday(&timeStart, NULL); 62 | }; 63 | 64 | void Stop(void) { 65 | gettimeofday(&timeStop, NULL); 66 | }; 67 | 68 | float GetDurationInSecs(void) 69 | { 70 | float dur = (timeStop.tv_sec - timeStart.tv_sec) + (timeStop.tv_usec - timeStart.tv_usec)/1000000.0; 71 | 72 | return dur; 73 | }; 74 | 75 | private: 76 | 77 | timeval timeStart; 78 | timeval timeStop; 79 | 80 | }; 81 | 82 | #endif // Win32 / Unix selection 83 | 84 | 85 | #endif //If PROFTIMER not defined 86 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Created by .ignore support plugin (hsz.mobi) 2 | ### Python template 3 | # Byte-compiled / optimized / DLL files 4 | __pycache__/ 5 | *.py[cod] 6 | *$py.class 7 | 8 | # C extensions 9 | *.so 10 | 11 | # Distribution / packaging 12 | .Python 13 | env/ 14 | build/ 15 | develop-eggs/ 16 | dist/ 17 | downloads/ 18 | eggs/ 19 | .eggs/ 20 | lib/ 21 | lib64/ 22 | parts/ 23 | sdist/ 24 | var/ 25 | *.egg-info/ 26 | .installed.cfg 27 | *.egg 28 | 29 | # PyInstaller 30 | # Usually these files are written by a python script from a template 31 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 32 | *.manifest 33 | *.spec 34 | 35 | # Installer logs 36 | pip-log.txt 37 | pip-delete-this-directory.txt 38 | 39 | # Unit test / coverage reports 40 | htmlcov/ 41 | .tox/ 42 | .coverage 43 | .coverage.* 44 | .cache 45 | nosetests.xml 46 | coverage.xml 47 | *,cover 48 | .hypothesis/ 49 | 50 | # Translations 51 | *.mo 52 | *.pot 53 | 54 | # Django stuff: 55 | *.log 56 | local_settings.py 57 | 58 | # Flask stuff: 59 | instance/ 60 | .webassets-cache 61 | 62 | # Scrapy stuff: 63 | .scrapy 64 | 65 | # Sphinx documentation 66 | docs/_build/ 67 | 68 | # PyBuilder 69 | target/ 70 | 71 | # Jupyter Notebook 72 | .ipynb_checkpoints 73 | 74 | # pyenv 75 | .python-version 76 | 77 | # celery beat schedule file 78 | celerybeat-schedule 79 | 80 | # dotenv 81 | .env 82 | 83 | # virtualenv 84 | .venv/ 85 | venv/ 86 | ENV/ 87 | 88 | # Spyder project settings 89 | .spyderproject 90 | 91 | # Rope project settings 92 | .ropeproject 93 | 94 | .idea/ 95 | related-paper/ 96 | *.npz 97 | *.gif 98 | *.pkl 99 | *.params 100 | *.yml 101 | *.json 102 | *.mp4 103 | analyze_training_log.ipynb 104 | !cfg_template.yml 105 | !experiments/hko/*/*.yml 106 | !experiments/movingmnist/*/*.yml 107 | !experiments/movingmnist/configurations/conv2d_3d/*.yml 108 | cse_cluster.tmpl 109 | dycpu.tmpl 110 | 111 | hko_data/radarPNG.tar.gz 112 | hko_data/radarPNG_mask.tar.gz 113 | hko_data/radarPNG/ 114 | hko_data/radarPNG_mask/ 115 | 116 | trained_model 117 | -------------------------------------------------------------------------------- /nowcasting/prediction_base_factory.py: -------------------------------------------------------------------------------- 1 | import mxnet as mx 2 | from nowcasting.config import cfg 3 | from nowcasting.ops import reset_regs 4 | from nowcasting.operators.common import grid_generator 5 | 6 | 7 | class PredictionBaseFactory(object): 8 | def __init__(self, batch_size, in_seq_len, out_seq_len, height, width, name="forecaster"): 9 | self._out_typ = cfg.MODEL.OUT_TYPE 10 | self._batch_size = batch_size 11 | self._in_seq_len = in_seq_len 12 | self._out_seq_len = out_seq_len 13 | self._height = height 14 | self._width = width 15 | self._name = name 16 | self._spatial_grid = grid_generator(batch_size=batch_size, height=height, width=width) 17 | self.rnn_list = self._init_rnn() 18 | self._reset_rnn() 19 | 20 | def _pre_encode_frame(self, frame_data, seqlen): 21 | ret = mx.sym.Concat(frame_data, 22 | mx.sym.broadcast_to(mx.sym.expand_dims(self._spatial_grid, axis=0), 23 | shape=(seqlen, self._batch_size, 24 | 2, self._height, self._width)), 25 | mx.sym.ones(shape=(seqlen, self._batch_size, 1, 26 | self._height, self._width)), 27 | num_args=3, dim=2) 28 | return ret 29 | 30 | def _init_rnn(self): 31 | raise NotImplementedError 32 | 33 | def _reset_rnn(self): 34 | for rnn in self.rnn_list: 35 | rnn.reset() 36 | 37 | def reset_all(self): 38 | reset_regs() 39 | self._reset_rnn() 40 | 41 | 42 | class RecursiveOneStepBaseFactory(PredictionBaseFactory): 43 | def __init__(self, batch_size, in_seq_len, out_seq_len, height, width, use_ss=False, 44 | name="forecaster"): 45 | super(RecursiveOneStepBaseFactory, self).__init__(batch_size=batch_size, 46 | in_seq_len=in_seq_len, 47 | out_seq_len=out_seq_len, 48 | height=height, 49 | width=width, 50 | name=name) 51 | self._use_ss = False 52 | 53 | -------------------------------------------------------------------------------- /hko_data/hko7_rainy_test_days.txt: -------------------------------------------------------------------------------- 1 | 20150107, 0.1 2 | 20150112, 14.9 3 | 20150113, 25.8 4 | 20150124, 0.9 5 | 20150206, 0.3 6 | 20150215, 3.3 7 | 20150221, 0.2 8 | 20150222, 15.6 9 | 20150223, 10.2 10 | 20150225, 0.8 11 | 20150227, 1.2 12 | 20150228, 0.4 13 | 20150303, 0.2 14 | 20150304, 0.2 15 | 20150305, 4.8 16 | 20150306, 0.1 17 | 20150307, 0.2 18 | 20150311, 0.3 19 | 20150312, 3.7 20 | 20150322, 0.1 21 | 20150326, 4.2 22 | 20150327, 14.6 23 | 20150407, 0.1 24 | 20150408, 10 25 | 20150409, 1.3 26 | 20150410, 0.7 27 | 20150411, 52 28 | 20150412, 0.2 29 | 20150420, 0.2 30 | 20150501, 0.5 31 | 20150505, 3.9 32 | 20150506, 0.6 33 | 20150507, 0.3 34 | 20150509, 7.3 35 | 20150510, 20.1 36 | 20150511, 51 37 | 20150516, 18.4 38 | 20150517, 5.7 39 | 20150518, 0.9 40 | 20150519, 1.2 41 | 20150520, 107.7 42 | 20150521, 12.6 43 | 20150522, 0.7 44 | 20150523, 169.4 45 | 20150524, 8.2 46 | 20150525, 29.4 47 | 20150526, 64.6 48 | 20150527, 0.2 49 | 20150528, 1.4 50 | 20150530, 7 51 | 20150531, 1.9 52 | 20150601, 10.6 53 | 20150602, 5.4 54 | 20150606, 0.8 55 | 20150608, 1.6 56 | 20150610, 8.1 57 | 20150611, 0.8 58 | 20150612, 96.8 59 | 20150613, 0.4 60 | 20150614, 1.5 61 | 20150615, 5.2 62 | 20150621, 51 63 | 20150622, 18.1 64 | 20150623, 51.3 65 | 20150624, 9.7 66 | 20150625, 28.5 67 | 20150626, 10.4 68 | 20150628, 1.9 69 | 20150709, 2 70 | 20150710, 24.3 71 | 20150717, 1.2 72 | 20150718, 0.2 73 | 20150720, 46.2 74 | 20150721, 51.2 75 | 20150722, 191.3 76 | 20150723, 45 77 | 20150724, 5.7 78 | 20150725, 9.6 79 | 20150726, 24.9 80 | 20150727, 0.3 81 | 20150729, 3.7 82 | 20150730, 0.6 83 | 20150809, 11.6 84 | 20150810, 23.5 85 | 20150811, 16.8 86 | 20150813, 27.5 87 | 20150814, 18.9 88 | 20150815, 24.6 89 | 20150816, 0.1 90 | 20150820, 6.1 91 | 20150823, 3.4 92 | 20150826, 0.2 93 | 20150829, 0.9 94 | 20150830, 9.7 95 | 20150901, 5.3 96 | 20150902, 39 97 | 20150903, 0.2 98 | 20150907, 7.3 99 | 20150916, 1.5 100 | 20150921, 16.9 101 | 20150922, 2.9 102 | 20150926, 10.2 103 | 20150927, 4.6 104 | 20151001, 0.3 105 | 20151002, 7 106 | 20151003, 46.4 107 | 20151004, 38.1 108 | 20151005, 15.6 109 | 20151006, 50.7 110 | 20151007, 5.8 111 | 20151010, 1 112 | 20151011, 2 113 | 20151025, 0.2 114 | 20151026, 0.7 115 | 20151031, 0.5 116 | 20151107, 0.3 117 | 20151110, 0.3 118 | 20151111, 1.1 119 | 20151112, 0.3 120 | 20151113, 10.4 121 | 20151115, 6.5 122 | 20151116, 3.9 123 | 20151205, 15.7 124 | 20151206, 1 125 | 20151208, 0.7 126 | 20151209, 44.6 127 | 20151220, 0.7 128 | 20151222, 0.6 129 | 20151225, 0.2 130 | 20151227, 0.4 131 | 20151230, 0.4 132 | -------------------------------------------------------------------------------- /nowcasting/utils.py: -------------------------------------------------------------------------------- 1 | try: 2 | import cPickle as pickle 3 | except: 4 | import pickle 5 | import ast 6 | import re 7 | import inspect 8 | import os 9 | import logging 10 | import numpy as np 11 | 12 | def cross_entropy_npy(a, b): 13 | return a * np.log(b + 1E-9) + (1 - a) * np.log(1 - b + 1E-9) 14 | 15 | 16 | def safe_eval(expr): 17 | if type(expr) is str: 18 | return ast.literal_eval(expr) 19 | else: 20 | return expr 21 | 22 | 23 | def logging_config(folder=None, name=None, 24 | level=logging.INFO, 25 | console_level=logging.DEBUG): 26 | """ 27 | 28 | Parameters 29 | ---------- 30 | folder : str or None 31 | name : str or None 32 | level : int 33 | console_level 34 | 35 | Returns 36 | ------- 37 | 38 | """ 39 | if name is None: 40 | name = inspect.stack()[1][1].split('.')[0] 41 | if folder is None: 42 | folder = os.path.join(os.getcwd(), name) 43 | if not os.path.exists(folder): 44 | os.makedirs(folder) 45 | # Remove all the current handlers 46 | for handler in logging.root.handlers: 47 | logging.root.removeHandler(handler) 48 | logging.root.handlers = [] 49 | logpath = os.path.join(folder, name + ".log") 50 | print("All Logs will be saved to %s" %logpath) 51 | logging.root.setLevel(level) 52 | formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s') 53 | logfile = logging.FileHandler(logpath) 54 | logfile.setLevel(level) 55 | logfile.setFormatter(formatter) 56 | logging.root.addHandler(logfile) 57 | # Initialze the console logging 58 | logconsole = logging.StreamHandler() 59 | logconsole.setLevel(console_level) 60 | logconsole.setFormatter(formatter) 61 | logging.root.addHandler(logconsole) 62 | return folder 63 | 64 | 65 | def load_params(prefix, epoch): 66 | """ 67 | 68 | Parameters 69 | ---------- 70 | prefix : str 71 | epoch : int 72 | 73 | Returns 74 | ------- 75 | arg_params : dict 76 | aux_params : dict 77 | """ 78 | import mxnet.ndarray as nd 79 | save_dict = nd.load('%s-%04d.params' % (prefix, epoch)) 80 | arg_params = {} 81 | aux_params = {} 82 | for k, v in save_dict.items(): 83 | tp, name = k.split(':', 1) 84 | if tp == 'arg': 85 | arg_params[name] = v 86 | if tp == 'aux': 87 | aux_params[name] = v 88 | return arg_params, aux_params 89 | 90 | 91 | def parse_ctx(ctx_args): 92 | import mxnet as mx 93 | ctx = re.findall('([a-z]+)(\d*)', ctx_args) 94 | ctx = [(device, int(num)) if len(num) > 0 else (device, 0) for device, num in ctx] 95 | ctx = [mx.Context(*ele) for ele in ctx] 96 | return ctx -------------------------------------------------------------------------------- /VarFlow/VarFlow.h: -------------------------------------------------------------------------------- 1 | //------------------------------------------------------------------ 2 | // Released under the BDS License 3 | // 4 | // Located at http://sourceforge.net/projects/varflow 5 | // 6 | //------------------------------------------------------------------ 7 | 8 | #ifndef VARFLOW_H 9 | #define VARFLOW_H 10 | 11 | #include 12 | #include 13 | 14 | /** 15 | * @brief Calculates dense optical flow using a variational method. 16 | * 17 | * The VarFlow class implements the method described in "Real-Time Optic Flow Computation with Variational Methods" by 18 | * Bruhn et al. (Lecture Notes in Computer Science, Volume 2756/2003, pp 222-229). It uses a recursive multigrid algorithm to 19 | * minimize the energy functional, leading to increased performance. This implementation uses an uncoupled Gauss-Seidel algorithm 20 | * and the temporal derivative of an image is calculated using a simple difference instead of a two point stencil as described 21 | * in the original paper. 22 | * 23 | * Date: April 2009 24 | * 25 | * @author Adam Harmat 26 | */ 27 | 28 | 29 | class VarFlow{ 30 | 31 | public: 32 | 33 | VarFlow(int width_in, int height_in, int max_level_in, int start_level_in, int n1_in, int n2_in, 34 | float rho_in, float alpha_in, float sigma_in); 35 | ~VarFlow(); 36 | int CalcFlow(IplImage* imgA, IplImage* imgB, IplImage* imgU, IplImage* imgV, bool saved_data); 37 | 38 | private: 39 | 40 | void gauss_seidel_recursive(int current_level, int max_level, int first_level, float h, IplImage** J13_array, IplImage** J23_array); 41 | void gauss_seidel_iteration(int current_level, float h, int num_iter, IplImage** J13_array, IplImage** J23_array); 42 | float gauss_seidel_step(IplImage* u, int x, int y, float h, float J11, float J12, float J13, float vi); 43 | float residual_part_step(IplImage* u, int x, int y, float h, float J11, float J12, float vi); 44 | void calculate_residual(int current_level, float h, IplImage** J13_array, IplImage** J23_array); 45 | 46 | 47 | float mask_x[5]; 48 | float mask_y[5]; 49 | CvMat fx_mask; 50 | CvMat fy_mask; 51 | 52 | IplImage* imgAsmall; 53 | IplImage* imgBsmall; 54 | 55 | IplImage* imgAfloat; 56 | IplImage* imgBfloat; 57 | 58 | IplImage* imgAfx; 59 | IplImage* imgAfy; 60 | IplImage* imgAft; 61 | 62 | IplImage** imgAfxfx_array; 63 | IplImage** imgAfxfy_array; 64 | IplImage** imgAfxft_array; 65 | IplImage** imgAfyfy_array; 66 | IplImage** imgAfyft_array; 67 | 68 | IplImage** imgU_array; 69 | IplImage** imgV_array; 70 | IplImage** imgU_res_err_array; 71 | IplImage** imgV_res_err_array; 72 | 73 | int initialized; 74 | 75 | int max_level; 76 | int start_level; 77 | 78 | int n1; 79 | int n2; 80 | 81 | float rho; 82 | float alpha; 83 | float sigma; 84 | }; 85 | 86 | 87 | #endif 88 | -------------------------------------------------------------------------------- /experiments/movingmnist/README.md: -------------------------------------------------------------------------------- 1 | # Training Scripts for the MovingMNIST++ experiments 2 | 3 | The models can be trained using the configurations files in the `configurations` folder. We given two examples: 4 | ``` 5 | # Train ConvGRU-K3D2 model 6 | python mnist_rnn_main.py \ 7 | --cfg configurations/convgru_1_64_1_96_1_96_K3D2.yml \ 8 | --save_dir convgru_K3D2 \ 9 | --ctx gpu0 10 | 11 | # Train ConvGRU-K5 model 12 | python mnist_rnn_main.py \ 13 | --cfg configurations/convgru_1_64_1_96_1_96_K5.yml \ 14 | --save_dir convgru_K5 \ 15 | --ctx gpu0 16 | 17 | # Train ConvGRU-K7 model 18 | python mnist_rnn_main.py \ 19 | --cfg configurations/convgru_1_64_1_96_1_96_K7.yml \ 20 | --save_dir convgru_K7 \ 21 | --ctx gpu0 22 | 23 | # Train DFN model 24 | python mnist_rnn_main.py \ 25 | --cfg configurations/convgru_1_64_1_96_1_96_K5_DFN.yml \ 26 | --save_dir convgru_K5_DFN \ 27 | --ctx gpu0 28 | 29 | # Train TrajGRU-L5 model 30 | python mnist_rnn_main.py \ 31 | --cfg configurations/trajgru_1_64_1_96_1_96_L5.yml \ 32 | --save_dir trajgru_L5 \ 33 | --ctx gpu0 34 | 35 | # Train TrajGRU-L9 model 36 | python mnist_rnn_main.py \ 37 | --cfg configurations/trajgru_1_64_1_96_1_96_L9.yml \ 38 | --save_dir trajgru_L9 \ 39 | --ctx gpu0 40 | 41 | # Train TrajGRU-L13 model 42 | python mnist_rnn_main.py \ 43 | --cfg configurations/trajgru_1_64_1_96_1_96_L13.yml \ 44 | --save_dir trajgru_L13 \ 45 | --ctx gpu0 46 | 47 | # Train TrajGRU-L17 model 48 | python mnist_rnn_main.py \ 49 | --cfg configurations/trajgru_1_64_1_96_1_96_L17.yml \ 50 | --save_dir trajgru_L17 \ 51 | --ctx gpu0 52 | ``` 53 | 54 | Also, we have the training scripts for Conv2D and Conv3D models. 55 | ``` 56 | # Train Conv2D model 57 | python deconvolution.py --cfg configurations/conv2d_3d/conv2d.yml --save_dir conv2d --ctx gpu0 58 | # Train Conv3D model 59 | python deconvolution.py --cfg configurations/conv2d_3d/conv3d.yml --save_dir conv3d --ctx gpu0 60 | ``` 61 | 62 | # Test with Pretrained Models 63 | 64 | The pretrained models can be downloaded from [Dropbox](https://www.dropbox.com/sh/n7gxfdd1pdasoio/AAC8uC4yto5Uam_7f3BEl-3La?dl=0) or using the `download_pretrained.sh` in the folder. After the models are downloaded, you can test them using the following commands: 65 | ``` 66 | # ConvGRU with K=3, D=2 67 | python mnist_rnn_test.py --cfg ConvGRU-K3D2/cfg0.yml --load_dir ConvGRU-K3D2 --load_iter 199999 --save_dir ConvGRU-K3D2 --ctx gpu0 68 | # ConvGRU with K=5, D=1 69 | python mnist_rnn_test.py --cfg ConvGRU-K5/cfg0.yml --load_dir ConvGRU-K5 --load_iter 199999 --save_dir ConvGRU-K5 --ctx gpu0 70 | # ConvGRU with K=7, D=1 71 | python mnist_rnn_test.py --cfg ConvGRU-K7/cfg0.yml --load_dir ConvGRU-K7 --load_iter 199999 --save_dir ConvGRU-K7 --ctx gpu0 72 | # DFN 73 | python mnist_rnn_test.py --cfg DFN/cfg0.yml --load_dir DFN --load_iter 199999 --save_dir DFN --ctx gpu0 74 | # TrajGRU with L=5 75 | python mnist_rnn_test.py --cfg TrajGRU-L5/cfg0.yml --load_dir TrajGRU-L5 --load_iter 199999 --save_dir TrajGRU-L5 --ctx gpu0 76 | # TrajGRU with L=9 77 | python mnist_rnn_test.py --cfg TrajGRU-L9/cfg0.yml --load_dir TrajGRU-L9 --load_iter 199999 --save_dir TrajGRU-L9 --ctx gpu0 78 | # TrajGRU with L=13 79 | python mnist_rnn_test.py --cfg TrajGRU-L13/cfg0.yml --load_dir TrajGRU-L13 --load_iter 199999 --save_dir TrajGRU-L13 --ctx gpu0 80 | # TrajGRU with L=17 81 | python mnist_rnn_test.py --cfg TrajGRU-L17/cfg0.yml --load_dir TrajGRU-L17 --load_iter 199999 --save_dir TrajGRU-L17 --ctx gpu0 82 | ``` -------------------------------------------------------------------------------- /experiments/hko/README.md: -------------------------------------------------------------------------------- 1 | # Scripts for the HKO-7 Benchmark 2 | 3 | - Use `last_frame_prediction.py` to evaluate the naive baseline which uses the last frame to predict all the future frames. 4 | It's a good example of how to test your own models using the offline setting of the benchmark. 5 | 6 | - Use `rover.py` to evaluate the rover algorithm in the following paper. Before running the experiments, make sure that the VarFlow package (https://github.com/sxjscience/HKO7_Benchmark/tree/master/VarFlow) is installed. 7 | ``` 8 | @article{woo2017operational, 9 | title={Operational Application of Optical Flow Techniques to Radar-Based Rainfall Nowcasting}, 10 | author={Woo, Wang-chun and Wong, Wai-kin}, 11 | journal={Atmosphere}, 12 | volume={8}, 13 | number={3}, 14 | pages={48}, 15 | year={2017}, 16 | publisher={Multidisciplinary Digital Publishing Institute} 17 | } 18 | ``` 19 | 20 | - Use `hko_main.py` to train the RNN models for precipitation nowcasting and use `hko_rnn_test.py` to test these models. The training scripts support multiple GPUs. 21 | If you find you cannot train the model using a single GPU, try to decrease the batch_size. 22 | 23 | 1. Commands for training the RNN models: 24 | ``` 25 | # Train ConvGRU model with B-MSE + B-MAE 26 | python3 hko_main.py --cfg configurations/convgru_55_55_33_1_64_1_192_1_192_b4.yml --save_dir convgru_55_55_33_1_64_1_192_1_192_b4 --ctx gpu0,gpu1 27 | # Train TrajGRU model with B-MSE + B-MAE 28 | python3 hko_main.py --cfg configurations/trajgru_55_55_33_1_64_1_192_1_192_13_13_9_b4.yml --save_dir trajgru_55_55_33_1_64_1_192_1_192_13_13_9_b4 --ctx gpu0,gpu1 29 | # Train ConvGRU model without B-MSE + B-MAE 30 | python3 hko_main.py --cfg configurations/convgru_55_55_33_1_64_1_192_1_192_nobal_b4.yml --save_dir convgru_55_55_33_1_64_1_192_1_192_nobal_b4 --ctx gpu0,gpu1 31 | ``` 32 | 33 | - Use `deconvolution.py` to run the 2D/3D models. 34 | 35 | 1. Commands for trainnig the CNN models: 36 | ``` 37 | # Train Conv2D model 38 | python deconvolution.py --cfg configurations/conv2d_3d/conv2d.yml --save_dir Conv2D --ctx gpu0 39 | # Train Conv3D model 40 | python deconvolution.py --cfg configurations/conv2d_3d/conv2d.yml --save_dir Conv3D --ctx gpu0 41 | ``` 42 | 43 | # Test with Pretrained Models 44 | 45 | You can download the pretrained ConvGRU/TrajGRU models by manually visit [Download By Dropbox](https://www.dropbox.com/sh/cp8zpi08umfiyha/AAAS6HJSsDQPjpKlxnBHtHvga?dl=0) or using the "download_pretrained.sh" in the folder. 46 | 47 | ``` 48 | # Test ConvGRU model in the offline setting 49 | python hko_rnn_test.py \ 50 | --cfg ConvGRU/cfg0.yml \ 51 | --load_dir ConvGRU \ 52 | --load_iter 49999 \ 53 | --finetune 0 \ 54 | --ctx gpu0 \ 55 | --save_dir ConvGRU \ 56 | --mode fixed \ 57 | --dataset test 58 | 59 | # Test ConvGRU model in the online setting 60 | python hko_rnn_test.py \ 61 | --cfg ConvGRU/cfg0.yml \ 62 | --load_dir ConvGRU \ 63 | --load_iter 49999 \ 64 | --finetune 1 \ 65 | --lr 1E-4 \ 66 | --ctx gpu0 \ 67 | --save_dir ConvGRU \ 68 | --mode online \ 69 | --dataset test 70 | 71 | # Test TrajGRU model in the offline setting 72 | python hko_rnn_test.py \ 73 | --cfg TrajGRU/cfg0.yml \ 74 | --load_dir TrajGRU \ 75 | --load_iter 79999 \ 76 | --finetune 0 \ 77 | --ctx gpu0 \ 78 | --save_dir TrajGRU \ 79 | --mode fixed \ 80 | --dataset test 81 | 82 | # Test TrajGRU model in the online setting 83 | python hko_rnn_test.py \ 84 | --cfg TrajGRU/cfg0.yml \ 85 | --load_dir TrajGRU \ 86 | --load_iter 79999 \ 87 | --finetune 1 \ 88 | --lr 1E-4 \ 89 | --ctx gpu0 \ 90 | --save_dir TrajGRU \ 91 | --mode online \ 92 | --dataset test 93 | 94 | # Test ConvGRU model without balanced loss 95 | python hko_rnn_test.py \ 96 | --cfg ConvGRU-nobal/cfg0.yml \ 97 | --load_dir ConvGRU-nobal \ 98 | --load_iter 59999 \ 99 | --finetune 0 \ 100 | --ctx gpu0 \ 101 | --save_dir ConvGRU-nobal \ 102 | --mode fixed \ 103 | --dataset test 104 | 105 | # Test ConvGRU model without balanced loss 106 | python hko_rnn_test.py \ 107 | --cfg ConvGRU-nobal/cfg0.yml \ 108 | --load_dir ConvGRU-nobal \ 109 | --load_iter 59999 \ 110 | --finetune 1 \ 111 | --lr 1E-4 \ 112 | --ctx gpu0 \ 113 | --save_dir ConvGRU-nobal \ 114 | --mode online \ 115 | --dataset test 116 | ``` -------------------------------------------------------------------------------- /experiments/hko/hko_rnn_test.py: -------------------------------------------------------------------------------- 1 | from nowcasting.config import cfg, save_cfg, cfg_from_file 2 | from nowcasting.utils import logging_config, parse_ctx 3 | from hko_factory import HKONowcastingFactory 4 | from nowcasting.encoder_forecaster import encoder_forecaster_build_networks, load_encoder_forecaster_params 5 | from hko_main import run_benchmark 6 | import os 7 | import argparse 8 | import logging 9 | 10 | 11 | def parse_args(): 12 | parser = argparse.ArgumentParser(description='Test the HKO nowcasting model') 13 | parser.add_argument('--cfg', dest='cfg_file', help='Optional configuration file', type=str) 14 | parser.add_argument('--load_dir', help='The directory to load the model', default=None, type=str) 15 | parser.add_argument('--load_iter', help='The iterator to load', default=-1, type=int) 16 | parser.add_argument('--save_dir', help='The saving directory', required=True, type=str) 17 | parser.add_argument('--ctx', dest='ctx', help='Running Context. E.g `--ctx gpu` or `--ctx gpu0,gpu1` or `--ctx cpu`', 18 | type=str, default='gpu') 19 | parser.add_argument('--finetune', dest='finetune', help='Whether to do online finetuning', 20 | default=None, type=int) 21 | parser.add_argument('--finetune_min_mse', dest='finetune_min_mse', help='Minimum error for finetuning', 22 | default=None, type=float) 23 | parser.add_argument('--mode', dest='mode', help='Whether to used fixed setting or online setting', 24 | required=True, type=str) 25 | parser.add_argument('--dataset', dest='dataset', help='Whether to used the test set or the validation set', 26 | default="test", type=str) 27 | parser.add_argument('--lr', dest='lr', help='learning rate', default=None, type=float) 28 | parser.add_argument('--wd', dest='wd', help='weight decay', default=None, type=float) 29 | parser.add_argument('--grad_clip', dest='grad_clip', help='gradient clipping threshold', 30 | default=None, type=float) 31 | args = parser.parse_args() 32 | args.ctx = parse_ctx(args.ctx) 33 | if args.cfg_file is not None: 34 | cfg_from_file(args.cfg_file, target=cfg.MODEL) 35 | if args.load_dir is not None: 36 | cfg.MODEL.LOAD_DIR = args.load_dir 37 | if args.load_iter != -1: 38 | cfg.MODEL.LOAD_ITER = args.load_iter 39 | if args.lr is not None: 40 | cfg.MODEL.TEST.ONLINE.LR = args.lr 41 | if args.wd is not None: 42 | cfg.MODEL.TEST.ONLINE.WD = args.wd 43 | if args.grad_clip is not None: 44 | cfg.MODEL.TEST.ONLINE.GRAD_CLIP = args.grad_clip 45 | if args.mode is not None: 46 | cfg.MODEL.TEST.MODE = args.mode 47 | if args.finetune is not None: 48 | cfg.MODEL.TEST.FINETUNE = (args.finetune != 0) 49 | if args.finetune_min_mse is not None: 50 | cfg.MODEL.TEST.ONLINE.FINETUNE_MIN_MSE = args.finetune_min_mse 51 | cfg.MODEL.SAVE_DIR = args.save_dir 52 | logging.info(args) 53 | return args 54 | 55 | 56 | def test_hko(args): 57 | assert cfg.MODEL.FRAME_STACK == 1 and cfg.MODEL.FRAME_SKIP == 1 58 | assert len(cfg.MODEL.LOAD_DIR) > 0 59 | base_dir = args.save_dir 60 | logging_config(folder=base_dir, name="testing") 61 | save_cfg(dir_path=base_dir, source=cfg.MODEL) 62 | hko_nowcasting_online = HKONowcastingFactory(batch_size=1, 63 | in_seq_len=cfg.MODEL.IN_LEN, 64 | out_seq_len=cfg.MODEL.OUT_LEN) 65 | t_encoder_net, t_forecaster_net, t_loss_net =\ 66 | encoder_forecaster_build_networks( 67 | factory=hko_nowcasting_online, 68 | context=args.ctx, 69 | for_finetune=True) 70 | t_encoder_net.summary() 71 | t_forecaster_net.summary() 72 | t_loss_net.summary() 73 | load_encoder_forecaster_params(load_dir=cfg.MODEL.LOAD_DIR, 74 | load_iter=cfg.MODEL.LOAD_ITER, 75 | encoder_net=t_encoder_net, 76 | forecaster_net=t_forecaster_net) 77 | if args.dataset == "test": 78 | pd_path = cfg.HKO_PD.RAINY_TEST 79 | elif args.dataset == "valid": 80 | pd_path = cfg.HKO_PD.RAINY_VALID 81 | else: 82 | raise NotImplementedError 83 | run_benchmark(hko_factory=hko_nowcasting_online, 84 | context=args.ctx[0], 85 | encoder_net=t_encoder_net, 86 | forecaster_net=t_forecaster_net, 87 | loss_net=t_loss_net, 88 | save_dir=os.path.join(base_dir, "iter%d_%s_finetune%d" 89 | % (cfg.MODEL.LOAD_ITER + 1, args.dataset, 90 | cfg.MODEL.TEST.FINETUNE)), 91 | finetune=cfg.MODEL.TEST.FINETUNE, 92 | mode=cfg.MODEL.TEST.MODE, 93 | pd_path=pd_path) 94 | 95 | 96 | if __name__ == "__main__": 97 | args = parse_args() 98 | test_hko(args) 99 | -------------------------------------------------------------------------------- /nowcasting/operators/transformations.py: -------------------------------------------------------------------------------- 1 | import mxnet as mx 2 | import numpy as np 3 | from nowcasting.operators.common import constant 4 | 5 | 6 | def CDNA(data, kernels, mask, batch_size, num_filter, kernel_size): 7 | """We assume that the kernels and masks are the output of an identity activation 8 | 9 | Parameters 10 | ---------- 11 | data : mx.sym.symbol 12 | Shape: (batch_size, C, H, W) 13 | kernels : mx.sym.symbol 14 | Shape: (batch_size, M, K, K) 15 | mask : mx.sym.symbol 16 | Shape: (batch_size, M, H, W) 17 | batch_size : int 18 | num_filter : int 19 | M 20 | kernel_size : int 21 | K 22 | Returns 23 | ------- 24 | ret : mx.sym.symbol 25 | Shape: (batch_size, C, H, W) 26 | """ 27 | assert kernel_size % 2 == 1, "Only support odd kernel size" 28 | # Use softmax activation for the kernel and the mask 29 | kernels = mx.sym.SoftmaxActivation(mx.sym.Reshape(kernels, 30 | shape=(-1, kernel_size * kernel_size))) 31 | kernels = mx.sym.Reshape(kernels, shape=(-1, num_filter, kernel_size, kernel_size)) 32 | mask = mx.sym.SoftmaxActivation(mask, mode="channel") 33 | 34 | data_sliced = mx.sym.SliceChannel(mx.sym.expand_dims(data, axis=2), axis=0, 35 | num_outputs=batch_size, squeeze_axis=True) # Each Shape: (C, 1, H, W) 36 | kernels_sliced = mx.sym.SliceChannel(mx.sym.expand_dims(kernels, axis=2), 37 | axis=0, num_outputs=batch_size, 38 | squeeze_axis=True) # Each Shape: (M, 1, K, K) 39 | out = [] 40 | for i in range(batch_size): 41 | ele = mx.sym.Convolution(data=data_sliced[i], 42 | num_filter=num_filter, 43 | kernel=(kernel_size, kernel_size), 44 | pad=(kernel_size/2, kernel_size/2), 45 | weight=kernels_sliced[i], no_bias=True) # Shape: (C, M, H, W) 46 | out.append(mx.sym.expand_dims(ele, axis=0)) 47 | out = mx.sym.Concat(*out, num_args=batch_size, dim=0) # Shape: (batch_size, C, M, H, W) 48 | mask = mx.sym.Reshape(mask, reverse=True, shape=(batch_size, 1, num_filter, 0, 0)) 49 | out = mx.sym.broadcast_mul(out, mask) 50 | out = mx.sym.sum(out, axis=2) 51 | return out 52 | 53 | def STP(data, affine_transform_matrices, mask, num_filter, kernel_size): 54 | """Spatial Transformer Predictor 55 | 56 | Parameters 57 | ---------- 58 | data : mx.sym.symbol 59 | affine_transform_matrices 60 | mask 61 | 62 | Returns 63 | ------- 64 | 65 | """ 66 | raise NotImplementedError() 67 | 68 | 69 | def DFN(data, local_kernels, K, batch_size): 70 | """[NIPS2016] Dynamic Filter Network 71 | 72 | Parameters 73 | ---------- 74 | data : mx.sym.symbol 75 | Shape: (batch_size, C, H, W) 76 | local_kernels : mx.sym.symbol 77 | Shape: (batch_size, K*K, H, W) 78 | K : int 79 | size of the local convolutional kernel 80 | batch_size : int 81 | size of the minibatch 82 | Returns 83 | ------- 84 | 85 | """ 86 | local_kernels = mx.sym.SoftmaxActivation(local_kernels, mode="channel") 87 | #filter_localexpand_npy = np.eye(K*K, K*K).reshape((K*K, 1, K, K)).astype(np.float32) 88 | #filter_localexpand = constant(filter_localexpand_npy, name="CDNA_kernels") 89 | filter_localexpand = mx.sym.one_hot(indices=mx.sym.arange(K * K), depth=K*K) 90 | filter_localexpand = mx.sym.reshape(mx.sym.transpose(filter_localexpand, axes=(1, 0)), 91 | shape=(K * K, 1, K, K)) 92 | data_sliced = mx.sym.SliceChannel(data, num_outputs=batch_size, axis=0, squeeze_axis=True) 93 | vec = [] 94 | for i in range(batch_size): 95 | ele = mx.sym.Convolution(data=mx.sym.expand_dims(data_sliced[i], axis=1), 96 | weight=filter_localexpand, 97 | num_filter=K*K, 98 | kernel=(K, K), 99 | pad=(K // 2, K // 2), no_bias=True) # Shape (C, K*K, H, W) 100 | vec.append(mx.sym.expand_dims(ele, axis=0)) 101 | input_localexpanded = mx.sym.Concat(*vec, num_args=len(vec), dim=0) # Shape (batch_size, C, K*K, H, W) 102 | output = mx.sym.broadcast_mul(input_localexpanded, mx.sym.expand_dims(local_kernels, axis=1)) 103 | output = mx.sym.sum(output, axis=2) 104 | return output 105 | 106 | 107 | 108 | if __name__ == '__main__': 109 | data = mx.sym.Variable('data') 110 | local_kernels = mx.sym.Variable('local_kernels') 111 | K = 11 112 | C = 3 113 | H = 60 114 | W = 60 115 | batch_size = 32 116 | local_kernels_npy = np.random.normal(size=(batch_size, K*K, H, W)) 117 | data_npy = np.random.normal(size=(batch_size, C, H, W)) 118 | out = data 119 | for i in range(10): 120 | out = DFN(data=out, local_kernels=local_kernels, K=K, batch_size=batch_size) 121 | exe = out.simple_bind(ctx=mx.gpu(), data=(batch_size, C, H, W), 122 | local_kernels=(batch_size, K*K, H, W)) 123 | exe.forward(data=data_npy, local_kernels=local_kernels_npy) 124 | print(exe.outputs[0].asnumpy().shape) 125 | -------------------------------------------------------------------------------- /experiments/hko/hko_factory.py: -------------------------------------------------------------------------------- 1 | import mxnet as mx 2 | from nowcasting.config import cfg 3 | from nowcasting.hko_evaluation import rainfall_to_pixel 4 | from nowcasting.encoder_forecaster import EncoderForecasterBaseFactory 5 | from nowcasting.operators import * 6 | from nowcasting.ops import * 7 | 8 | 9 | def get_loss_weight_symbol(data, mask, seq_len): 10 | if cfg.MODEL.USE_BALANCED_LOSS: 11 | balancing_weights = cfg.HKO.EVALUATION.BALANCING_WEIGHTS 12 | weights = mx.sym.ones_like(data) * balancing_weights[0] 13 | thresholds = [rainfall_to_pixel(ele) for ele in cfg.HKO.EVALUATION.THRESHOLDS] 14 | for i, threshold in enumerate(thresholds): 15 | weights = weights + (balancing_weights[i + 1] - balancing_weights[i]) * (data >= threshold) 16 | weights = weights * mask 17 | else: 18 | weights = mask 19 | if cfg.MODEL.TEMPORAL_WEIGHT_TYPE == "same": 20 | return weights 21 | elif cfg.MODEL.TEMPORAL_WEIGHT_TYPE == "linear": 22 | upper = cfg.MODEL.TEMPORAL_WEIGHT_UPPER 23 | assert upper >= 1.0 24 | temporal_mult = 1 + \ 25 | mx.sym.arange(start=0, stop=seq_len) * (upper - 1.0) / (seq_len - 1.0) 26 | temporal_mult = mx.sym.reshape(temporal_mult, shape=(seq_len, 1, 1, 1, 1)) 27 | weights = mx.sym.broadcast_mul(weights, temporal_mult) 28 | return weights 29 | elif cfg.MODEL.TEMPORAL_WEIGHT_TYPE == "exponential": 30 | upper = cfg.MODEL.TEMPORAL_WEIGHT_UPPER 31 | assert upper >= 1.0 32 | base_factor = np.log(upper) / (seq_len - 1.0) 33 | temporal_mult = mx.sym.exp(mx.sym.arange(start=0, stop=seq_len) * base_factor) 34 | temporal_mult = mx.sym.reshape(temporal_mult, shape=(seq_len, 1, 1, 1, 1)) 35 | weights = mx.sym.broadcast_mul(weights, temporal_mult) 36 | return weights 37 | else: 38 | raise NotImplementedError 39 | 40 | 41 | class HKONowcastingFactory(EncoderForecasterBaseFactory): 42 | def __init__(self, 43 | batch_size, 44 | in_seq_len, 45 | out_seq_len, 46 | ctx_num=1, 47 | name="hko_nowcasting"): 48 | super(HKONowcastingFactory, self).__init__(batch_size=batch_size, 49 | in_seq_len=in_seq_len, 50 | out_seq_len=out_seq_len, 51 | ctx_num=ctx_num, 52 | height=cfg.HKO.ITERATOR.HEIGHT, 53 | width=cfg.HKO.ITERATOR.WIDTH, 54 | name=name) 55 | self._central_region = cfg.HKO.EVALUATION.CENTRAL_REGION 56 | 57 | def _slice_central(self, data): 58 | """Slice the central region in the given symbol 59 | 60 | Parameters 61 | ---------- 62 | data : mx.sym.Symbol 63 | 64 | Returns 65 | ------- 66 | ret : mx.sym.Symbol 67 | """ 68 | x_begin, y_begin, x_end, y_end = self._central_region 69 | return mx.sym.slice(data, 70 | begin=(0, 0, 0, y_begin, x_begin), 71 | end=(None, None, None, y_end, x_end)) 72 | 73 | def _concat_month_code(self): 74 | # TODO 75 | raise NotImplementedError 76 | 77 | def loss_sym(self, 78 | pred=mx.sym.Variable('pred'), 79 | mask=mx.sym.Variable('mask'), 80 | target=mx.sym.Variable('target')): 81 | """Construct loss symbol. 82 | 83 | Optional args: 84 | pred: Shape (out_seq_len, batch_size, C, H, W) 85 | mask: Shape (out_seq_len, batch_size, C, H, W) 86 | target: Shape (out_seq_len, batch_size, C, H, W) 87 | """ 88 | self.reset_all() 89 | weights = get_loss_weight_symbol(data=target, mask=mask, seq_len=self._out_seq_len) 90 | mse = weighted_mse(pred=pred, gt=target, weight=weights) 91 | mae = weighted_mae(pred=pred, gt=target, weight=weights) 92 | gdl = masked_gdl_loss(pred=pred, gt=target, mask=mask) 93 | avg_mse = mx.sym.mean(mse) 94 | avg_mae = mx.sym.mean(mae) 95 | avg_gdl = mx.sym.mean(gdl) 96 | global_grad_scale = cfg.MODEL.NORMAL_LOSS_GLOBAL_SCALE 97 | if cfg.MODEL.L2_LAMBDA > 0: 98 | avg_mse = mx.sym.MakeLoss(avg_mse, 99 | grad_scale=global_grad_scale * cfg.MODEL.L2_LAMBDA, 100 | name="mse") 101 | else: 102 | avg_mse = mx.sym.BlockGrad(avg_mse, name="mse") 103 | if cfg.MODEL.L1_LAMBDA > 0: 104 | avg_mae = mx.sym.MakeLoss(avg_mae, 105 | grad_scale=global_grad_scale * cfg.MODEL.L1_LAMBDA, 106 | name="mae") 107 | else: 108 | avg_mae = mx.sym.BlockGrad(avg_mae, name="mae") 109 | if cfg.MODEL.GDL_LAMBDA > 0: 110 | avg_gdl = mx.sym.MakeLoss(avg_gdl, 111 | grad_scale=global_grad_scale * cfg.MODEL.GDL_LAMBDA, 112 | name="gdl") 113 | else: 114 | avg_gdl = mx.sym.BlockGrad(avg_gdl, name="gdl") 115 | loss = mx.sym.Group([avg_mse, avg_mae, avg_gdl]) 116 | return loss 117 | -------------------------------------------------------------------------------- /VarFlow/varflow/varflow.py: -------------------------------------------------------------------------------- 1 | import ctypes 2 | import os 3 | import numpy as np 4 | import cv2 5 | from concurrent.futures import ThreadPoolExecutor, wait 6 | 7 | _BASE_PATH = os.path.abspath(os.path.join(os.path.dirname(__file__))) 8 | _VALID_DLL_PATH = [os.path.join(_BASE_PATH, '..', 'build', 'Release', 'varflow.dll'), 9 | os.path.join(_BASE_PATH, '..', 'build', 'libvarflow.so')] 10 | 11 | _VARFLOW_DLL_PATH = None 12 | for p in _VALID_DLL_PATH: 13 | if os.path.exists(p): 14 | _VARFLOW_DLL_PATH = p 15 | break 16 | if _VARFLOW_DLL_PATH is None: 17 | raise RuntimeError("DLL not found! Valid PATH=%s" %(_VALID_DLL_PATH)) 18 | _CDLL = ctypes.cdll.LoadLibrary(_VARFLOW_DLL_PATH) 19 | 20 | class VarFlowFactory(object): 21 | def __init__(self, max_level, start_level, n1, n2, rho, alpha, sigma): 22 | self._max_level = max_level 23 | self._start_level = start_level 24 | self._n1 = n1 25 | self._n2 = n2 26 | self._rho = rho 27 | self._alpha = alpha 28 | self._sigma = sigma 29 | self._varflow_executor_pool = ThreadPoolExecutor(max_workers=16) 30 | 31 | def calc_flow(self, I1, I2): 32 | """ 33 | 34 | Parameters 35 | ---------- 36 | I1 : np.ndarray 37 | Shape: (H, W) 38 | I2 : np.ndarray 39 | Shape: (H, W) 40 | Returns 41 | ------- 42 | velocity : np.ndarray 43 | Shape: (2, H, W) 44 | The channel dimension will be flow_x, flow_y 45 | """ 46 | if I1.dtype == np.float32: 47 | I1 = (I1 * 255).astype(np.uint8) 48 | else: 49 | I1 = I1.astype(np.uint8) 50 | if I2.dtype == np.float32: 51 | I2 = (I2 * 255).astype(np.uint8) 52 | else: 53 | I2 = I2.astype(np.uint8) 54 | assert I1.ndim == 2 and I2.ndim == 2 55 | assert I1.shape == I2.shape 56 | np.ascontiguousarray(I1) 57 | np.ascontiguousarray(I2) 58 | height, width = I1.shape 59 | velocity = np.zeros((2,) + I1.shape, dtype=np.float32) 60 | self._base_varflow_call(velocity=velocity, I1=I1, I2=I2, width=width, height=height) 61 | return velocity 62 | 63 | def _base_varflow_call(self, velocity, I1, I2, width, height): 64 | _CDLL.varflow(ctypes.c_int32(width), 65 | ctypes.c_int32(height), 66 | ctypes.c_int32(self._max_level), 67 | ctypes.c_int32(self._start_level), 68 | ctypes.c_int32(self._n1), 69 | ctypes.c_int32(self._n2), 70 | ctypes.c_float(self._rho), 71 | ctypes.c_float(self._alpha), 72 | ctypes.c_float(self._sigma), 73 | velocity[0].ctypes.data_as(ctypes.c_void_p), 74 | velocity[1].ctypes.data_as(ctypes.c_void_p), 75 | I1.ctypes.data_as(ctypes.c_void_p), 76 | I2.ctypes.data_as(ctypes.c_void_p)) 77 | 78 | def batch_calc_flow(self, I1, I2): 79 | """Calculate the optical flow from two 80 | 81 | Parameters 82 | ---------- 83 | I1 : np.ndarray 84 | Shape: (batch_size, H, W) 85 | I2 : np.ndarray 86 | Shape: (batch_size, H, W) 87 | Returns 88 | ------- 89 | velocity : np.ndarray 90 | Shape: (batch_size, 2, H, W) 91 | The channel dimension will be flow_x, flow_y 92 | """ 93 | if I1.dtype == np.float32: 94 | I1 = (I1 * 255).astype(np.uint8) 95 | else: 96 | I1 = I1.astype(np.uint8) 97 | if I2.dtype == np.float32: 98 | I2 = (I2 * 255).astype(np.uint8) 99 | else: 100 | I2 = I2.astype(np.uint8) 101 | np.ascontiguousarray(I1) 102 | np.ascontiguousarray(I2) 103 | assert I1.ndim == 3 and I2.ndim == 3 104 | assert I1.shape == I2.shape 105 | batch_size, height, width = I1.shape 106 | velocity = np.zeros((batch_size, 2, height, width), dtype=np.float32) 107 | future_objs = [] 108 | for i in range(batch_size): 109 | obj = self._varflow_executor_pool.submit( 110 | self._base_varflow_call, velocity[i], I1[i], I2[i], width, height) 111 | future_objs.append(obj) 112 | wait(future_objs) 113 | return velocity 114 | 115 | 116 | if __name__ == '__main__': 117 | I1 = cv2.imread('../Data/yos_img_08.jpg', 0) 118 | I2 = cv2.imread('../Data/yos_img_09.jpg', 0) 119 | varflow_factory = VarFlowFactory(max_level=4, start_level=0, n1=2, n2=2, rho=2.8, alpha=1400, 120 | sigma=1.5) 121 | velocity = varflow_factory.calc_flow(I1=I1, I2=I2) 122 | batch_I1 = np.concatenate([I1.reshape((1,) + I1.shape), I2.reshape((1,) + I2.shape)], axis=0) 123 | batch_I2 = np.concatenate([I2.reshape((1,) + I2.shape), I1.reshape((1,) + I1.shape)], axis=0) 124 | batch_velocity = varflow_factory.batch_calc_flow(I1=batch_I1, I2=batch_I2) 125 | velocity = batch_velocity[0] 126 | import matplotlib.pyplot as plt 127 | Q = plt.quiver(velocity[0, ::5, ::5], velocity[1, ::5, ::5]) 128 | qk = plt.quiverkey(Q, 0.5, 0.98, 2, r'$2 \frac{m}{s}$', labelpos='W', 129 | fontproperties={'weight': 'bold'}) 130 | plt.gca().invert_yaxis() 131 | l, r, b, t = plt.axis() 132 | dx, dy = r - l, t - b 133 | plt.axis([l - 0.05*dx, r + 0.05*dx, b - 0.05*dy, t + 0.05*dy]) 134 | plt.title('Minimal arguments, no kwargs') 135 | plt.show() -------------------------------------------------------------------------------- /nowcasting/hko_factory.py: -------------------------------------------------------------------------------- 1 | # TODO this is a copy of experiments/hko_factory.py and should be removed 2 | # after nowcasting/models/deconvolution.py has been refactored to use a factory 3 | # to get the symbols. 4 | # Currently it needs to import the the factory directly to construct the symbol 5 | # based on the cfg.DATASET variable 6 | 7 | import mxnet as mx 8 | from nowcasting.config import cfg 9 | from nowcasting.hko_evaluation import rainfall_to_pixel 10 | from nowcasting.encoder_forecaster import EncoderForecasterBaseFactory 11 | from nowcasting.operators import * 12 | from nowcasting.ops import * 13 | 14 | 15 | def get_loss_weight_symbol(data, mask, seq_len): 16 | if cfg.MODEL.USE_BALANCED_LOSS: 17 | balancing_weights = cfg.HKO.EVALUATION.BALANCING_WEIGHTS 18 | weights = mx.sym.ones_like(data) * balancing_weights[0] 19 | thresholds = [rainfall_to_pixel(ele) for ele in cfg.HKO.EVALUATION.THRESHOLDS] 20 | for i, threshold in enumerate(thresholds): 21 | weights = weights + (balancing_weights[i + 1] - balancing_weights[i]) * (data >= threshold) 22 | weights = weights * mask 23 | else: 24 | weights = mask 25 | if cfg.MODEL.TEMPORAL_WEIGHT_TYPE == "same": 26 | return weights 27 | elif cfg.MODEL.TEMPORAL_WEIGHT_TYPE == "linear": 28 | upper = cfg.MODEL.TEMPORAL_WEIGHT_UPPER 29 | assert upper >= 1.0 30 | temporal_mult = 1 + \ 31 | mx.sym.arange(start=0, stop=seq_len) * (upper - 1.0) / (seq_len - 1.0) 32 | temporal_mult = mx.sym.reshape(temporal_mult, shape=(seq_len, 1, 1, 1, 1)) 33 | weights = mx.sym.broadcast_mul(weights, temporal_mult) 34 | return weights 35 | elif cfg.MODEL.TEMPORAL_WEIGHT_TYPE == "exponential": 36 | upper = cfg.MODEL.TEMPORAL_WEIGHT_UPPER 37 | assert upper >= 1.0 38 | base_factor = np.log(upper) / (seq_len - 1.0) 39 | temporal_mult = mx.sym.exp(mx.sym.arange(start=0, stop=seq_len) * base_factor) 40 | temporal_mult = mx.sym.reshape(temporal_mult, shape=(seq_len, 1, 1, 1, 1)) 41 | weights = mx.sym.broadcast_mul(weights, temporal_mult) 42 | return weights 43 | else: 44 | raise NotImplementedError 45 | 46 | class HKONowcastingFactory(EncoderForecasterBaseFactory): 47 | def __init__(self, 48 | batch_size, 49 | in_seq_len, 50 | out_seq_len, 51 | name="hko_nowcasting"): 52 | super(HKONowcastingFactory, self).__init__(batch_size=batch_size, 53 | in_seq_len=in_seq_len, 54 | out_seq_len=out_seq_len, 55 | height=cfg.HKO.ITERATOR.HEIGHT, 56 | width=cfg.HKO.ITERATOR.WIDTH, 57 | name=name) 58 | self._central_region = cfg.HKO.EVALUATION.CENTRAL_REGION 59 | 60 | def _slice_central(self, data): 61 | """Slice the central region in the given symbol 62 | 63 | Parameters 64 | ---------- 65 | data : mx.sym.Symbol 66 | 67 | Returns 68 | ------- 69 | ret : mx.sym.Symbol 70 | """ 71 | x_begin, y_begin, x_end, y_end = self._central_region 72 | return mx.sym.slice(data, 73 | begin=(0, 0, 0, y_begin, x_begin), 74 | end=(None, None, None, y_end, x_end)) 75 | 76 | def _concat_month_code(self): 77 | #TODO 78 | raise NotImplementedError 79 | 80 | def loss_sym(self, 81 | pred=mx.sym.Variable('pred'), 82 | mask=mx.sym.Variable('mask'), 83 | target=mx.sym.Variable('target')): 84 | """Construct loss symbol. 85 | 86 | Optional args: 87 | pred: Shape (out_seq_len, batch_size, C, H, W) 88 | mask: Shape (out_seq_len, batch_size, C, H, W) 89 | target: Shape (out_seq_len, batch_size, C, H, W) 90 | """ 91 | self.reset_all() 92 | weights = get_loss_weight_symbol(data=target, mask=mask, seq_len=self._out_seq_len) 93 | mse = weighted_mse(pred=pred, gt=target, weight=weights) 94 | mae = weighted_mae(pred=pred, gt=target, weight=weights) 95 | gdl = masked_gdl_loss(pred=pred, gt=target, mask=mask) 96 | avg_mse = mx.sym.mean(mse) 97 | avg_mae = mx.sym.mean(mae) 98 | avg_gdl = mx.sym.mean(gdl) 99 | global_grad_scale = cfg.MODEL.NORMAL_LOSS_GLOBAL_SCALE 100 | if cfg.MODEL.L2_LAMBDA > 0: 101 | avg_mse = mx.sym.MakeLoss(avg_mse, 102 | grad_scale=global_grad_scale * cfg.MODEL.L2_LAMBDA, 103 | name="mse") 104 | else: 105 | avg_mse = mx.sym.BlockGrad(avg_mse, name="mse") 106 | if cfg.MODEL.L1_LAMBDA > 0: 107 | avg_mae = mx.sym.MakeLoss(avg_mae, 108 | grad_scale=global_grad_scale * cfg.MODEL.L1_LAMBDA, 109 | name="mae") 110 | else: 111 | avg_mae = mx.sym.BlockGrad(avg_mae, name="mae") 112 | if cfg.MODEL.GDL_LAMBDA > 0: 113 | avg_gdl = mx.sym.MakeLoss(avg_gdl, 114 | grad_scale=global_grad_scale * cfg.MODEL.GDL_LAMBDA, 115 | name="gdl") 116 | else: 117 | avg_gdl = mx.sym.BlockGrad(avg_gdl, name="gdl") 118 | loss = mx.sym.Group([avg_mse, avg_mae, avg_gdl]) 119 | return loss 120 | -------------------------------------------------------------------------------- /nowcasting/image.py: -------------------------------------------------------------------------------- 1 | # Python plugin that supports loading batch of images in parallel 2 | import cv2 3 | import numpy 4 | import threading 5 | import os 6 | import struct 7 | import concurrent.futures 8 | from concurrent.futures import ThreadPoolExecutor, wait 9 | 10 | _imread_executor_pool = ThreadPoolExecutor(max_workers=16) 11 | 12 | class UnknownImageFormat(Exception): 13 | pass 14 | 15 | 16 | def quick_imsize(file_path): 17 | """Return (width, height) for a given img file content - no external 18 | dependencies except the os and struct modules from core 19 | 20 | Parameters 21 | ---------- 22 | file_path 23 | 24 | Returns 25 | ------- 26 | 27 | """ 28 | size = os.path.getsize(file_path) 29 | with open(file_path, 'rb') as input: 30 | height = -1 31 | width = -1 32 | data = input.read(25) 33 | 34 | if (size >= 10) and data[:6] in ('GIF87a', 'GIF89a'): 35 | # GIFs 36 | w, h = struct.unpack("= 24) and data.startswith('\211PNG\r\n\032\n') 40 | and (data[12:16] == 'IHDR')): 41 | # PNGs 42 | w, h = struct.unpack(">LL", data[16:24]) 43 | width = int(w) 44 | height = int(h) 45 | elif (size >= 16) and data.startswith('\211PNG\r\n\032\n'): 46 | # older PNGs? 47 | w, h = struct.unpack(">LL", data[8:16]) 48 | width = int(w) 49 | height = int(h) 50 | elif (size >= 2) and data.startswith('\377\330'): 51 | # JPEG 52 | msg = " raised while trying to decode as JPEG." 53 | input.seek(0) 54 | input.read(2) 55 | b = input.read(1) 56 | try: 57 | while (b and ord(b) != 0xDA): 58 | while (ord(b) != 0xFF): b = input.read(1) 59 | while (ord(b) == 0xFF): b = input.read(1) 60 | if (ord(b) >= 0xC0 and ord(b) <= 0xC3): 61 | input.read(3) 62 | h, w = struct.unpack(">HH", input.read(4)) 63 | break 64 | else: 65 | input.read(int(struct.unpack(">H", input.read(2))[0]) - 2) 66 | b = input.read(1) 67 | width = int(w) 68 | height = int(h) 69 | except struct.error: 70 | raise UnknownImageFormat("StructError" + msg) 71 | except ValueError: 72 | raise UnknownImageFormat("ValueError" + msg) 73 | except Exception as e: 74 | raise UnknownImageFormat(e.__class__.__name__ + msg) 75 | else: 76 | raise UnknownImageFormat( 77 | "Sorry, don't know how to get information from this file." 78 | ) 79 | 80 | return width, height 81 | 82 | 83 | def cv2_read_img_resize(path, read_storage, resize_storage, frame_size, grayscale): 84 | if grayscale: 85 | read_storage[:] = cv2.imread(path, 0) 86 | else: 87 | read_storage[:] = cv2.imread(path) 88 | resize_storage[:] = cv2.resize(read_storage, frame_size, interpolation=cv2.INTER_LINEAR) 89 | 90 | 91 | def cv2_read_img(path, read_storage, grayscale): 92 | if grayscale: 93 | read_storage[:] = cv2.imread(path, 0) 94 | else: 95 | read_storage[:] = cv2.imread(path) 96 | 97 | 98 | def quick_read_frames(path_list, im_w=None, im_h=None, resize=False, frame_size=None, grayscale=True): 99 | """Multi-thread Frame Loader 100 | 101 | Parameters 102 | ---------- 103 | path_list : list 104 | resize : bool, optional 105 | frame_size : None or tuple 106 | 107 | Returns 108 | ------- 109 | 110 | """ 111 | img_num = len(path_list) 112 | for i in range(img_num): 113 | if not os.path.exists(path_list[i]): 114 | raise IOError 115 | if im_w is None or im_h is None: 116 | im_w, im_h = quick_imsize(path_list[0]) 117 | if grayscale: 118 | read_storage = numpy.empty((img_num, im_h, im_w), dtype=numpy.uint8) 119 | else: 120 | read_storage = numpy.empty((img_num, im_h, im_w, 3), dtype=numpy.uint8) 121 | if resize: 122 | if grayscale: 123 | resize_storage = numpy.empty((img_num, frame_size[0], frame_size[1]), dtype=numpy.uint8) 124 | else: 125 | resize_storage = numpy.empty((img_num, frame_size[0], frame_size[1], 3), dtype=numpy.uint8) 126 | if img_num == 1: 127 | cv2_read_img_resize(path=path_list[0], read_storage=read_storage[0], 128 | resize_storage=resize_storage[0], 129 | frame_size=frame_size, grayscale=grayscale) 130 | else: 131 | future_objs = [] 132 | for i in range(img_num): 133 | obj = _imread_executor_pool.submit(cv2_read_img_resize, 134 | path_list[i], 135 | read_storage[i], 136 | resize_storage[i], frame_size, grayscale) 137 | future_objs.append(obj) 138 | wait(future_objs) 139 | if grayscale: 140 | resize_storage = resize_storage.reshape((img_num, 1, frame_size[0], frame_size[1])) 141 | else: 142 | resize_storage = resize_storage.transpose((0, 3, 1, 2)) 143 | return resize_storage[:, ::-1, ...] 144 | else: 145 | if img_num == 1: 146 | cv2_read_img(path=path_list[0], read_storage=read_storage[0], grayscale=grayscale) 147 | else: 148 | future_objs = [] 149 | for i in range(img_num): 150 | obj = _imread_executor_pool.submit(cv2_read_img, path_list[i], read_storage[i], grayscale) 151 | future_objs.append(obj) 152 | wait(future_objs) 153 | if grayscale: 154 | read_storage = read_storage.reshape((img_num, 1, im_h, im_w)) 155 | else: 156 | read_storage = read_storage.transpose((0, 3, 1, 2)) 157 | return read_storage[:, ::-1, ...] -------------------------------------------------------------------------------- /nowcasting/helpers/visualization.py: -------------------------------------------------------------------------------- 1 | import matplotlib.pyplot as plt 2 | import matplotlib.gridspec as gridspec 3 | from matplotlib.colors import hsv_to_rgb 4 | import cv2 5 | import moviepy.editor as mpy 6 | import numpy as np 7 | from nowcasting.helpers.gifmaker import save_gif 8 | 9 | def flow_to_img(flow_dat, max_displacement=None): 10 | """Convert optical flow data to HSV images 11 | 12 | Parameters 13 | ---------- 14 | flow_dat : np.ndarray 15 | Shape: (seq_len, 2, H, W) 16 | max_displacement : float or None 17 | 18 | Returns 19 | ------- 20 | rgb_dat : np.ndarray 21 | Shape: (seq_len, 3, H, W) 22 | """ 23 | assert flow_dat.ndim == 4 24 | flow_scale = np.square(flow_dat).sum(axis=1, keepdims=True) 25 | flow_x = flow_dat[:, :1, :, :] 26 | flow_y = flow_dat[:, 1:, :, :] 27 | flow_angle = np.arctan2(flow_y, flow_x) 28 | flow_angle[flow_angle < 0] += np.pi * 2 29 | v = np.ones((flow_dat.shape[0], 1, flow_dat.shape[2], flow_dat.shape[3]), 30 | dtype=np.float32) 31 | if max_displacement is None: 32 | flow_scale_max = np.sqrt(flow_scale.max()) 33 | else: 34 | flow_scale_max = max_displacement 35 | h = flow_angle / (2 * np.pi) 36 | s = np.sqrt(flow_scale) / flow_scale_max 37 | 38 | hsv_dat = np.concatenate((h, s, v), axis=1) 39 | rgb_dat = hsv_to_rgb(hsv_dat.transpose((0, 2, 3, 1))).transpose((0, 3, 1, 2)) 40 | return rgb_dat 41 | 42 | 43 | def _ax_imshow(ax, im, **kwargs): 44 | assert im.ndim == 3 or im.ndim == 2 45 | if im.ndim == 2: 46 | ax.imshow(im, **kwargs) 47 | ax.set_axis_off() 48 | else: 49 | if im.shape[0] == 1: 50 | ax.imshow(im[0, :, :], **kwargs) 51 | ax.set_axis_off() 52 | elif im.shape[0] == 3: 53 | ax.imshow(im.transpose((1, 2, 0)), **kwargs) 54 | ax.set_axis_off() 55 | else: 56 | raise NotImplementedError 57 | ax.set_adjustable('box-forced') 58 | ax.autoscale(False) 59 | 60 | 61 | def get_color_flow_legend_image(size=50): 62 | U, V = np.meshgrid(np.arange(-size, size + 1, dtype=np.float32), 63 | np.arange(-size, size + 1, dtype=np.float32)) 64 | flow_scale = np.sqrt(U**2 + V**2) 65 | flow_angle = np.arctan2(V, U) 66 | flow_angle[flow_angle < 0] += np.pi * 2 67 | max_flow_scale = float(size) * np.sqrt(2) 68 | h = flow_angle / (2 * np.pi) 69 | s = flow_scale / max_flow_scale 70 | v = np.ones((size * 2 + 1, size * 2 + 1), 71 | dtype=np.float32) 72 | hsv_dat = np.concatenate((h.reshape((1, size * 2 + 1, size * 2 + 1)), 73 | s.reshape((1, size * 2 + 1, size * 2 + 1)), 74 | v.reshape((1, size * 2 + 1, size * 2 + 1))), axis=0) 75 | rgb_dat = hsv_to_rgb(hsv_dat.transpose((1, 2, 0))).transpose((2, 0, 1)) 76 | a = np.ones((1, size * 2 + 1, size * 2 + 1), dtype=np.float32) 77 | rgb_dat[:, flow_scale > max_flow_scale] = 0 78 | a[:, flow_scale > max_flow_scale] = 0 79 | rgba_dat = np.concatenate((rgb_dat, a), axis=0) 80 | return rgba_dat 81 | 82 | 83 | def save_hko_gif(im_dat, save_path): 84 | """Save the HKO images to gif 85 | 86 | Parameters 87 | ---------- 88 | im_dat : np.ndarray 89 | Shape: (seqlen, H, W) 90 | save_path : str 91 | Returns 92 | ------- 93 | """ 94 | assert im_dat.ndim == 3 95 | save_gif(im_dat, fname=save_path) 96 | return 97 | 98 | 99 | def merge_rgba_cv2(front_img, back_img): 100 | """Merge the front image with the background image using the `Painter's algorithm` 101 | 102 | Parameters 103 | ---------- 104 | front_img : np.ndarray 105 | back_img : np.ndarray 106 | 107 | Returns 108 | ------- 109 | result_img : np.ndarray 110 | """ 111 | assert front_img.shape == back_img.shape 112 | if front_img.dtype == np.uint8: 113 | front_img = front_img.astype(np.float32) / 255.0 114 | if back_img.dtype == np.uint8: 115 | back_img = back_img.astype(np.float32) / 255.0 116 | result_img = np.zeros(front_img.shape, dtype=np.float32) 117 | result_img[:, :, 3] = front_img[:, :, 3] + back_img[:, :, 3] * (1 - front_img[:, :, 3]) 118 | result_img[:, :, :3] = (front_img[:, :, :3] * front_img[:, :, 3:] + 119 | back_img[:, :, :3] * back_img[:, :, 3:] * (1 - front_img[:, :, 3:])) /\ 120 | result_img[:, :, 3:] 121 | result_img = (result_img * 255.0).astype(np.uint8) 122 | return result_img 123 | 124 | 125 | def save_hko_movie(im_dat, datetime_list, mask_dat=None, save_path="hko.mp4", masked=False, 126 | fps=5, prediction_start=None): 127 | """Save the HKO images to a video file 128 | 129 | Parameters 130 | ---------- 131 | im_dat : np.ndarray 132 | Shape : (seq_len, H, W) 133 | datetime_list : list 134 | list of datetimes 135 | mask_dat : np.ndarray or None 136 | Shape : (seq_len, H, W) 137 | save_path : str 138 | masked : bool 139 | whether the mask the inputs when saving the image 140 | fps : float 141 | the fps of the saved movie 142 | prediction_start : int or None 143 | The starting point of the prediction 144 | """ 145 | from nowcasting.config import cfg 146 | central_region = cfg.HKO.EVALUATION.CENTRAL_REGION 147 | seq_len, height, width = im_dat.shape 148 | display_im_dat = [] 149 | mask_color = np.array((0, 170, 160, 150), dtype=np.float32) / 255.0 150 | if im_dat.dtype == np.float32: 151 | im_dat = (im_dat * 255).astype(np.uint8) 152 | for i in range(im_dat.shape[0]): 153 | if not masked: 154 | color_im_dat = cv2.cvtColor(im_dat[i], cv2.COLOR_GRAY2RGBA) 155 | im = color_im_dat 156 | else: 157 | im = im_dat[i] * mask_dat[i] 158 | im = cv2.cvtColor(im, cv2.COLOR_GRAY2RGBA) 159 | # Uncomment the following code to add transparency to the masks 160 | # color_im_dat = cv2.cvtColor(im_dat[i], cv2.COLOR_GRAY2RGBA) 161 | # mask_im_dat = mask_color.reshape((1, 1, 4)) * np.expand_dims(1 - mask_dat[i], axis=2) 162 | # im = merge_rgba_cv2(front_img=mask_im_dat, back_img=color_im_dat) 163 | if prediction_start is not None and i >= prediction_start: 164 | cv2.putText(im, text=datetime_list[i].strftime('%Y/%m/%d %H:%M'), 165 | org=(0, 20), fontFace=cv2.FONT_HERSHEY_SIMPLEX, fontScale=0.4, 166 | color=(255, 0, 0, 0)) 167 | else: 168 | cv2.putText(im, text=datetime_list[i].strftime('%Y/%m/%d %H:%M'), 169 | org=(0, 20), fontFace=cv2.FONT_HERSHEY_SIMPLEX, fontScale=0.4, 170 | color=(255, 255, 255, 0)) 171 | cv2.rectangle(im, 172 | pt1=(central_region[0], central_region[1]), 173 | pt2=(central_region[2], central_region[3]), 174 | color=(0, 255, 0, 0)) 175 | display_im_dat.append(im) 176 | clip = mpy.ImageSequenceClip(display_im_dat, with_mask=False, fps=fps) 177 | clip.write_videofile(save_path, audio=False, verbose=False, threads=4) 178 | -------------------------------------------------------------------------------- /nowcasting/helpers/msssim.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/python 2 | # 3 | # Copyright 2016 The TensorFlow Authors. All Rights Reserved. 4 | # 5 | # Licensed under the Apache License, Version 2.0 (the "License"); 6 | # you may not use this file except in compliance with the License. 7 | # You may obtain a copy of the License at 8 | # 9 | # http://www.apache.org/licenses/LICENSE-2.0 10 | # 11 | # Unless required by applicable law or agreed to in writing, software 12 | # distributed under the License is distributed on an "AS IS" BASIS, 13 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | # See the License for the specific language governing permissions and 15 | # limitations under the License. 16 | # ============================================================================== 17 | 18 | """Python implementation of MS-SSIM. 19 | Usage: 20 | python msssim.py --original_image=original.png --compared_image=distorted.png 21 | """ 22 | import numpy as np 23 | from scipy import signal 24 | from scipy.ndimage.filters import convolve 25 | 26 | def _FSpecialGauss(size, sigma): 27 | """Function to mimic the 'fspecial' gaussian MATLAB function.""" 28 | radius = size // 2 29 | offset = 0.0 30 | start, stop = -radius, radius + 1 31 | if size % 2 == 0: 32 | offset = 0.5 33 | stop -= 1 34 | x, y = np.mgrid[offset + start:stop, offset + start:stop] 35 | assert len(x) == size 36 | g = np.exp(-((x**2 + y**2)/(2.0 * sigma**2))) 37 | return g / g.sum() 38 | 39 | 40 | def _SSIMForMultiScale(img1, img2, max_val=255, filter_size=11, 41 | filter_sigma=1.5, k1=0.01, k2=0.03): 42 | """Return the Structural Similarity Map between `img1` and `img2`. 43 | This function attempts to match the functionality of ssim_index_new.m by 44 | Zhou Wang: http://www.cns.nyu.edu/~lcv/ssim/msssim.zip 45 | Arguments: 46 | img1: Numpy array holding the first RGB image batch. 47 | img2: Numpy array holding the second RGB image batch. 48 | max_val: the dynamic range of the images (i.e., the difference between the 49 | maximum the and minimum allowed values). 50 | filter_size: Size of blur kernel to use (will be reduced for small images). 51 | filter_sigma: Standard deviation for Gaussian blur kernel (will be reduced 52 | for small images). 53 | k1: Constant used to maintain stability in the SSIM calculation (0.01 in 54 | the original paper). 55 | k2: Constant used to maintain stability in the SSIM calculation (0.03 in 56 | the original paper). 57 | Returns: 58 | Pair containing the mean SSIM and contrast sensitivity between `img1` and 59 | `img2`. 60 | Raises: 61 | RuntimeError: If input images don't have the same shape or don't have four 62 | dimensions: [batch_size, height, width, depth]. 63 | """ 64 | if img1.shape != img2.shape: 65 | raise RuntimeError('Input images must have the same shape (%s vs. %s).', 66 | img1.shape, img2.shape) 67 | if img1.ndim != 4: 68 | raise RuntimeError('Input images must have four dimensions, not %d', 69 | img1.ndim) 70 | 71 | img1 = img1.astype(np.float64) 72 | img2 = img2.astype(np.float64) 73 | _, height, width, _ = img1.shape 74 | 75 | # Filter size can't be larger than height or width of images. 76 | size = min(filter_size, height, width) 77 | 78 | # Scale down sigma if a smaller filter size is used. 79 | sigma = size * filter_sigma / filter_size if filter_size else 0 80 | 81 | if filter_size: 82 | window = np.reshape(_FSpecialGauss(size, sigma), (1, size, size, 1)) 83 | mu1 = signal.fftconvolve(img1, window, mode='valid') 84 | mu2 = signal.fftconvolve(img2, window, mode='valid') 85 | sigma11 = signal.fftconvolve(img1 * img1, window, mode='valid') 86 | sigma22 = signal.fftconvolve(img2 * img2, window, mode='valid') 87 | sigma12 = signal.fftconvolve(img1 * img2, window, mode='valid') 88 | else: 89 | # Empty blur kernel so no need to convolve. 90 | mu1, mu2 = img1, img2 91 | sigma11 = img1 * img1 92 | sigma22 = img2 * img2 93 | sigma12 = img1 * img2 94 | 95 | mu11 = mu1 * mu1 96 | mu22 = mu2 * mu2 97 | mu12 = mu1 * mu2 98 | sigma11 -= mu11 99 | sigma22 -= mu22 100 | sigma12 -= mu12 101 | 102 | # Calculate intermediate values used by both ssim and cs_map. 103 | c1 = (k1 * max_val) ** 2 104 | c2 = (k2 * max_val) ** 2 105 | v1 = 2.0 * sigma12 + c2 106 | v2 = sigma11 + sigma22 + c2 107 | ssim = np.mean((((2.0 * mu12 + c1) * v1) / ((mu11 + mu22 + c1) * v2)), axis=(1, 2, 3)) 108 | cs = np.mean(v1 / v2, axis=(1, 2, 3)) 109 | return ssim, cs 110 | 111 | 112 | def MultiScaleSSIM(img1, img2, max_val=255, filter_size=11, filter_sigma=1.5, 113 | k1=0.01, k2=0.03, weights=None): 114 | """Return the MS-SSIM score between `img1` and `img2`. 115 | This function implements Multi-Scale Structural Similarity (MS-SSIM) Image 116 | Quality Assessment according to Zhou Wang's paper, "Multi-scale structural 117 | similarity for image quality assessment" (2003). 118 | Link: https://ece.uwaterloo.ca/~z70wang/publications/msssim.pdf 119 | Author's MATLAB implementation: 120 | http://www.cns.nyu.edu/~lcv/ssim/msssim.zip 121 | Arguments: 122 | img1: Numpy array holding the first RGB image batch. 123 | img2: Numpy array holding the second RGB image batch. 124 | max_val: the dynamic range of the images (i.e., the difference between the 125 | maximum the and minimum allowed values). 126 | filter_size: Size of blur kernel to use (will be reduced for small images). 127 | filter_sigma: Standard deviation for Gaussian blur kernel (will be reduced 128 | for small images). 129 | k1: Constant used to maintain stability in the SSIM calculation (0.01 in 130 | the original paper). 131 | k2: Constant used to maintain stability in the SSIM calculation (0.03 in 132 | the original paper). 133 | weights: List of weights for each level; if none, use five levels and the 134 | weights from the original paper. 135 | Returns: 136 | MS-SSIM score between `img1` and `img2`. 137 | Raises: 138 | RuntimeError: If input images don't have the same shape or don't have four 139 | dimensions: [batch_size, height, width, depth]. 140 | """ 141 | if img1.shape != img2.shape: 142 | raise RuntimeError('Input images must have the same shape (%s vs. %s).', 143 | img1.shape, img2.shape) 144 | if img1.ndim != 4: 145 | raise RuntimeError('Input images must have four dimensions, not %d', 146 | img1.ndim) 147 | 148 | # Note: default weights don't sum to 1.0 but do match the paper / matlab code. 149 | weights = np.array(weights if weights else 150 | [0.0448, 0.2856, 0.3001, 0.2363, 0.1333]) 151 | levels = weights.size 152 | downsample_filter = np.ones((1, 2, 2, 1)) / 4.0 153 | im1, im2 = [x.astype(np.float64) for x in [img1, img2]] 154 | mssim = np.array([]) 155 | mcs = np.array([]) 156 | for _ in range(levels): 157 | ssim, cs = _SSIMForMultiScale( 158 | im1, im2, max_val=max_val, filter_size=filter_size, 159 | filter_sigma=filter_sigma, k1=k1, k2=k2) 160 | mssim = np.append(mssim, ssim.mean()) 161 | mcs = np.append(mcs, cs.mean()) 162 | filtered = [convolve(im, downsample_filter, mode='reflect') 163 | for im in [im1, im2]] 164 | im1, im2 = [x[:, ::2, ::2, :] for x in filtered] 165 | return (np.prod(mcs[0:levels-1] ** weights[0:levels-1]) * 166 | (mssim[levels-1] ** weights[levels-1])) -------------------------------------------------------------------------------- /experiments/hko/rover.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from scipy.interpolate import NearestNDInterpolator 3 | import mxnet as mx 4 | import mxnet.ndarray as nd 5 | import os 6 | from varflow import VarFlowFactory 7 | from nowcasting.config import cfg 8 | from nowcasting.hko_evaluation import HKOEvaluation, pixel_to_dBZ, dBZ_to_pixel 9 | from nowcasting.hko_benchmark import HKOBenchmarkEnv 10 | from nowcasting.hko_iterator import precompute_mask 11 | from nowcasting.helpers.visualization import save_hko_gif 12 | from nowcasting.utils import logging_config 13 | 14 | 15 | class NonLinearRoverTransform(object): 16 | def __init__(self, Zc=33, sharpness=4): 17 | self.Zc = float(Zc) 18 | self.sharpness = float(sharpness) 19 | 20 | def transform(self, img): 21 | dbz_img = pixel_to_dBZ(img) 22 | dbz_lower = pixel_to_dBZ(0.0) 23 | dbz_upper = pixel_to_dBZ(1.0) 24 | transformed_lower = np.arctan((dbz_lower - self.Zc) / self.sharpness) 25 | transformed_upper = np.arctan((dbz_upper - self.Zc) / self.sharpness) 26 | transformed_img = np.arctan((dbz_img - self.Zc) / self.sharpness) 27 | transformed_img = (transformed_img - transformed_lower) / \ 28 | (transformed_upper - transformed_lower) 29 | return transformed_img 30 | 31 | def rev_transform(self, transformed_img): 32 | dbz_lower = pixel_to_dBZ(0.0) 33 | dbz_upper = pixel_to_dBZ(1.0) 34 | transformed_lower = np.arctan((dbz_lower - self.Zc) / self.sharpness) 35 | transformed_upper = np.arctan((dbz_upper - self.Zc) / self.sharpness) 36 | img = transformed_img * (transformed_upper - transformed_lower) + transformed_lower 37 | img = np.tan(img) * self.sharpness + self.Zc 38 | img = dBZ_to_pixel(dBZ_img=img) 39 | return img 40 | 41 | 42 | def nd_advection(im, flow): 43 | """ 44 | 45 | Parameters 46 | ---------- 47 | im : nd.NDArray 48 | Shape: (batch_size, C, H, W) 49 | flow : nd.NDArray 50 | Shape: (batch_size, 2, H, W) 51 | Returns 52 | ------- 53 | new_im : nd.NDArray 54 | """ 55 | grid = nd.GridGenerator(-flow, transform_type="warp") 56 | new_im = nd.BilinearSampler(im, grid) 57 | return new_im 58 | 59 | 60 | def nearest_neighbor_advection(im, flow): 61 | """ 62 | 63 | Parameters 64 | ---------- 65 | im : np.ndarray 66 | Shape: (batch_size, C, H, W) 67 | flow : np.ndarray 68 | Shape: (batch_size, 2, H, W) 69 | Returns 70 | ------- 71 | new_im : nd.NDArray 72 | """ 73 | predict_frame = np.empty(im.shape, dtype=im.dtype) 74 | batch_size, channel_num, height, width = im.shape 75 | assert channel_num == 1 76 | grid_x, grid_y = np.meshgrid(np.arange(width), np.arange(height)) 77 | interp_grid = np.hstack([grid_x.reshape((-1, 1)), grid_y.reshape((-1, 1))]) 78 | for i in range(batch_size): 79 | flow_interpolator = NearestNDInterpolator(interp_grid, im[i].ravel()) 80 | predict_grid = interp_grid + np.hstack([flow[i][0].reshape((-1, 1)), 81 | flow[i][1].reshape((-1, 1))]) 82 | predict_frame[i, 0, ...] = flow_interpolator(predict_grid).reshape((height, width)) 83 | return predict_frame 84 | 85 | 86 | def run(pd_path=cfg.HKO_PD.RAINY_TEST, 87 | mode="fixed", 88 | interp_type="bilinear", 89 | nonlinear_transform=True): 90 | transformer = NonLinearRoverTransform() 91 | flow_factory = VarFlowFactory(max_level=6, start_level=0, 92 | n1=2, n2=2, 93 | rho=1.5, alpha=2000, 94 | sigma=4.5) 95 | assert interp_type == "bilinear", "Nearest interpolation is implemented in CPU and is too slow." \ 96 | " We only support bilinear interpolation for rover." 97 | if nonlinear_transform: 98 | base_dir = os.path.join('hko7_benchmark', 'rover-nonlinear') 99 | else: 100 | base_dir = os.path.join('hko7_benchmark', 'rover-linear') 101 | logging_config(base_dir) 102 | batch_size = 1 103 | env = HKOBenchmarkEnv(pd_path=pd_path, 104 | save_dir=base_dir, 105 | mode=mode) 106 | counter = 0 107 | while not env.done: 108 | in_frame_dat, in_datetime_clips, out_datetime_clips, \ 109 | begin_new_episode, need_upload_prediction = \ 110 | env.get_observation(batch_size=batch_size) 111 | if need_upload_prediction: 112 | counter += 1 113 | prediction = np.zeros(shape=(cfg.HKO.BENCHMARK.OUT_LEN,) + in_frame_dat.shape[1:], 114 | dtype=np.float32) 115 | I1 = in_frame_dat[-2, :, 0, :, :] 116 | I2 = in_frame_dat[-1, :, 0, :, :] 117 | mask_I1 = precompute_mask(I1) 118 | mask_I2 = precompute_mask(I2) 119 | I1 = I1 * mask_I1 120 | I2 = I2 * mask_I2 121 | if nonlinear_transform: 122 | I1 = transformer.transform(I1) 123 | I2 = transformer.transform(I2) 124 | flow = flow_factory.batch_calc_flow(I1=I1, I2=I2) 125 | if interp_type == "bilinear": 126 | init_im = nd.array(I2.reshape((I2.shape[0], 1, I2.shape[1], I2.shape[2])), 127 | ctx=mx.gpu()) 128 | nd_flow = nd.array(np.concatenate((flow[:, :1, :, :], -flow[:, 1:, :, :]), axis=1), 129 | ctx=mx.gpu()) 130 | nd_pred_im = nd.zeros(shape=prediction.shape) 131 | for i in range(cfg.HKO.BENCHMARK.OUT_LEN): 132 | new_im = nd_advection(init_im, flow=nd_flow) 133 | nd_pred_im[i][:] = new_im 134 | init_im[:] = new_im 135 | prediction = nd_pred_im.asnumpy() 136 | elif interp_type == "nearest": 137 | init_im = I2.reshape((I2.shape[0], 1, I2.shape[1], I2.shape[2])) 138 | for i in range(cfg.HKO.BENCHMARK.OUT_LEN): 139 | new_im = nearest_neighbor_advection(init_im, flow) 140 | prediction[i, ...] = new_im 141 | init_im = new_im 142 | if nonlinear_transform: 143 | prediction = transformer.rev_transform(prediction) 144 | env.upload_prediction(prediction=prediction) 145 | if counter % 10 == 0: 146 | save_hko_gif(in_frame_dat[:, 0, 0, :, :], 147 | save_path=os.path.join(base_dir, 'in.gif')) 148 | save_hko_gif(prediction[:, 0, 0, :, :], 149 | save_path=os.path.join(base_dir, 'pred.gif')) 150 | env.print_stat_readable() 151 | # import matplotlib.pyplot as plt 152 | # Q = plt.quiver(flow[1, 0, ::10, ::10], flow[1, 1, ::10, ::10]) 153 | # plt.gca().invert_yaxis() 154 | # plt.show() 155 | # ch = raw_input() 156 | env.save_eval() 157 | 158 | # Running fixed 159 | run(cfg.HKO_PD.RAINY_TEST, mode="fixed", nonlinear_transform=True) 160 | run(cfg.HKO_PD.RAINY_TEST, mode="fixed", nonlinear_transform=False) 161 | run(cfg.HKO_PD.RAINY_VALID, mode="fixed", nonlinear_transform=True) 162 | run(cfg.HKO_PD.RAINY_VALID, mode="fixed", nonlinear_transform=False) 163 | 164 | 165 | # Running online 166 | # run(cfg.HKO_PD.RAINY_TEST, mode="online", nonlinear_transform=True) 167 | # run(cfg.HKO_PD.RAINY_TEST, mode="online", nonlinear_transform=False) 168 | # run(cfg.HKO_PD.RAINY_VALID, mode="online", nonlinear_transform=True) 169 | # run(cfg.HKO_PD.RAINY_VALID, mode="online", nonlinear_transform=False) 170 | -------------------------------------------------------------------------------- /experiments/movingmnist/mnist_rnn_test.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import os 3 | import random 4 | import logging 5 | import numpy as np 6 | import mxnet as mx 7 | import mxnet.ndarray as nd 8 | from mnist_rnn_factory import MovingMNISTFactory 9 | from nowcasting.config import cfg, cfg_from_file, save_cfg 10 | from nowcasting.helpers.gifmaker import save_gif 11 | from nowcasting.my_module import MyModule 12 | from nowcasting.encoder_forecaster import * 13 | from nowcasting.utils import parse_ctx, logging_config, load_params 14 | from nowcasting.movingmnist_iterator import MovingMNISTAdvancedIterator 15 | from nowcasting.helpers.ordered_easydict import OrderedEasyDict as edict 16 | from mnist_rnn_main import save_movingmnist_cfg, mnist_get_prediction 17 | 18 | random.seed(12345) 19 | mx.random.seed(930215) 20 | np.random.seed(921206) 21 | 22 | 23 | def parse_args(): 24 | parser = argparse.ArgumentParser(description='Test the MovingMNIST++ dataset') 25 | parser.add_argument('--batch_size', dest='batch_size', help="batchsize of the testing process", 26 | default=None, type=int) 27 | parser.add_argument('--cfg', dest='cfg_file', help='Testing configuration file', default=None, type=str) 28 | parser.add_argument('--load_dir', help='The loading directory', default=None, type=str) 29 | parser.add_argument('--load_iter', help='The loading iterator', default=None, type=int) 30 | parser.add_argument('--save_dir', help='The saving directory', required=True, type=str) 31 | parser.add_argument('--ctx', dest='ctx', help='Running Context. E.g `--ctx gpu` or `--ctx gpu0,gpu1` or `--ctx cpu`', 32 | type=str, default='gpu') 33 | args = parser.parse_args() 34 | args.ctx = parse_ctx(args.ctx) 35 | if args.cfg_file is not None: 36 | cfg_from_file(args.cfg_file, target=cfg) 37 | if args.load_dir is not None: 38 | cfg.MODEL.LOAD_DIR = args.load_dir 39 | if args.load_iter is not None: 40 | cfg.MODEL.LOAD_ITER = args.load_iter 41 | cfg.MODEL.SAVE_DIR = args.save_dir 42 | logging.info(args) 43 | return args 44 | 45 | def analysis(args): 46 | cfg.MODEL.TRAJRNN.SAVE_MID_RESULTS = True 47 | assert cfg.MODEL.FRAME_STACK == 1 and cfg.MODEL.FRAME_SKIP == 1 48 | base_dir = args.save_dir 49 | logging_config(folder=base_dir, name="testing") 50 | save_movingmnist_cfg(base_dir) 51 | mnist_iter = MovingMNISTAdvancedIterator( 52 | distractor_num=cfg.MOVINGMNIST.DISTRACTOR_NUM, 53 | initial_velocity_range=(cfg.MOVINGMNIST.VELOCITY_LOWER, 54 | cfg.MOVINGMNIST.VELOCITY_UPPER), 55 | rotation_angle_range=(cfg.MOVINGMNIST.ROTATION_LOWER, 56 | cfg.MOVINGMNIST.ROTATION_UPPER), 57 | scale_variation_range=(cfg.MOVINGMNIST.SCALE_VARIATION_LOWER, 58 | cfg.MOVINGMNIST.SCALE_VARIATION_UPPER), 59 | illumination_factor_range=(cfg.MOVINGMNIST.ILLUMINATION_LOWER, 60 | cfg.MOVINGMNIST.ILLUMINATION_UPPER)) 61 | mnist_rnn = MovingMNISTFactory(batch_size=1, 62 | in_seq_len=cfg.MODEL.IN_LEN, 63 | out_seq_len=cfg.MODEL.OUT_LEN) 64 | encoder_net, forecaster_net, loss_net = \ 65 | encoder_forecaster_build_networks( 66 | factory=mnist_rnn, 67 | context=args.ctx) 68 | encoder_net.summary() 69 | forecaster_net.summary() 70 | loss_net.summary() 71 | states = EncoderForecasterStates(factory=mnist_rnn, ctx=args.ctx[0]) 72 | states.reset_all() 73 | # Begin to load the model if load_dir is not empty 74 | assert len(cfg.MODEL.LOAD_DIR) > 0 75 | load_encoder_forecaster_params( 76 | load_dir=cfg.MODEL.LOAD_DIR, load_iter=cfg.MODEL.LOAD_ITER, 77 | encoder_net=encoder_net, forecaster_net=forecaster_net) 78 | for iter_id in range(1): 79 | frame_dat, _ = mnist_iter.sample(batch_size=cfg.MODEL.TRAIN.BATCH_SIZE, 80 | seqlen=cfg.MOVINGMNIST.IN_LEN + cfg.MOVINGMNIST.OUT_LEN) 81 | data_nd = mx.nd.array(frame_dat[0:cfg.MOVINGMNIST.IN_LEN, ...], ctx=args.ctx[0]) / 255.0 82 | target_nd = mx.nd.array( 83 | frame_dat[cfg.MOVINGMNIST.IN_LEN:(cfg.MOVINGMNIST.IN_LEN + cfg.MOVINGMNIST.OUT_LEN), ...], 84 | ctx=args.ctx[0]) / 255.0 85 | pred_nd = mnist_get_prediction(data_nd=data_nd, states=states, 86 | encoder_net=encoder_net, 87 | forecaster_net=forecaster_net) 88 | save_gif(pred_nd.asnumpy()[:, 0, 0, :, :], os.path.join(base_dir, "pred.gif")) 89 | save_gif(data_nd.asnumpy()[:, 0, 0, :, :], os.path.join(base_dir, "in.gif")) 90 | save_gif(target_nd.asnumpy()[:, 0, 0, :, :], os.path.join(base_dir, "gt.gif")) 91 | 92 | def test(args): 93 | cfg.MODEL.TRAJRNN.SAVE_MID_RESULTS = False 94 | assert cfg.MODEL.FRAME_STACK == 1 and cfg.MODEL.FRAME_SKIP == 1 95 | base_dir = args.save_dir 96 | logging_config(folder=base_dir, name="testing") 97 | save_movingmnist_cfg(base_dir) 98 | batch_size = 4 99 | mnist_iter = MovingMNISTAdvancedIterator( 100 | distractor_num=cfg.MOVINGMNIST.DISTRACTOR_NUM, 101 | initial_velocity_range=(cfg.MOVINGMNIST.VELOCITY_LOWER, 102 | cfg.MOVINGMNIST.VELOCITY_UPPER), 103 | rotation_angle_range=(cfg.MOVINGMNIST.ROTATION_LOWER, 104 | cfg.MOVINGMNIST.ROTATION_UPPER), 105 | scale_variation_range=(cfg.MOVINGMNIST.SCALE_VARIATION_LOWER, 106 | cfg.MOVINGMNIST.SCALE_VARIATION_UPPER), 107 | illumination_factor_range=(cfg.MOVINGMNIST.ILLUMINATION_LOWER, 108 | cfg.MOVINGMNIST.ILLUMINATION_UPPER)) 109 | mnist_iter.load(file=cfg.MOVINGMNIST.TEST_FILE) 110 | mnist_rnn = MovingMNISTFactory(batch_size=batch_size, 111 | in_seq_len=cfg.MODEL.IN_LEN, 112 | out_seq_len=cfg.MODEL.OUT_LEN) 113 | encoder_net, forecaster_net, loss_net = \ 114 | encoder_forecaster_build_networks( 115 | factory=mnist_rnn, 116 | context=args.ctx) 117 | encoder_net.summary() 118 | forecaster_net.summary() 119 | loss_net.summary() 120 | states = EncoderForecasterStates(factory=mnist_rnn, ctx=args.ctx[0]) 121 | states.reset_all() 122 | # Begin to load the model if load_dir is not empty 123 | assert len(cfg.MODEL.LOAD_DIR) > 0 124 | load_encoder_forecaster_params( 125 | load_dir=cfg.MODEL.LOAD_DIR, load_iter=cfg.MODEL.LOAD_ITER, 126 | encoder_net=encoder_net, forecaster_net=forecaster_net) 127 | overall_mse = 0 128 | for iter_id in range(10000 // batch_size): 129 | frame_dat, _ = mnist_iter.sample(batch_size=batch_size, 130 | seqlen=cfg.MOVINGMNIST.IN_LEN + cfg.MOVINGMNIST.OUT_LEN, 131 | random=False) 132 | data_nd = mx.nd.array(frame_dat[0:cfg.MOVINGMNIST.IN_LEN, ...], ctx=args.ctx[0]) / 255.0 133 | target_nd = mx.nd.array( 134 | frame_dat[cfg.MOVINGMNIST.IN_LEN:(cfg.MOVINGMNIST.IN_LEN + cfg.MOVINGMNIST.OUT_LEN), ...], 135 | ctx=args.ctx[0]) / 255.0 136 | pred_nd = mnist_get_prediction(data_nd=data_nd, states=states, 137 | encoder_net=encoder_net, forecaster_net=forecaster_net) 138 | overall_mse += mx.nd.mean(mx.nd.square(pred_nd - target_nd)).asscalar() 139 | print(iter_id, overall_mse / (iter_id + 1)) 140 | avg_mse = overall_mse / (10000 // batch_size) 141 | with open(os.path.join(base_dir, 'result.txt'), 'w') as f: 142 | f.write(str(avg_mse)) 143 | print(base_dir, avg_mse) 144 | 145 | if __name__ == "__main__": 146 | backup_test_file_path = cfg.MOVINGMNIST.TEST_FILE 147 | args = parse_args() 148 | cfg.MOVINGMNIST.TEST_FILE = backup_test_file_path 149 | test(args) 150 | -------------------------------------------------------------------------------- /experiments/movingmnist/deconvolution.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import logging 3 | import os 4 | import mxnet as mx 5 | import numpy as np 6 | from nowcasting.config import cfg, save_cfg 7 | from nowcasting.utils import logging_config 8 | from nowcasting.movingmnist_iterator import MovingMNISTAdvancedIterator 9 | from nowcasting.helpers.gifmaker import save_gif, save_gifs 10 | from nowcasting.models.deconvolution import (construct_l2_loss, 11 | construct_modules, train_step, 12 | test_step, get_base_dir, 13 | model_args, training_args, 14 | mode_args, parse_mode_args, 15 | parse_training_args, 16 | parse_model_args) 17 | 18 | 19 | ### Arguments 20 | def argument_parser(): 21 | parser = argparse.ArgumentParser( 22 | description='Deconvolution baseline for MovingMNIST') 23 | 24 | cfg.DATASET = "MOVINGMNIST" 25 | 26 | mode_args(parser) 27 | training_args(parser) 28 | dataset_args(parser) 29 | model_args(parser) 30 | 31 | args = parser.parse_args() 32 | 33 | parse_mode_args(args) 34 | parse_training_args(args) 35 | parse_dataset_args(args) 36 | parse_model_args(args) 37 | 38 | base_dir = get_base_dir(args) 39 | logging_config(folder=base_dir, name="testing") 40 | save_cfg(base_dir, source=cfg.MODEL) 41 | 42 | logging.info(args) 43 | return args 44 | 45 | 46 | def dataset_args(parser): 47 | group = parser.add_argument_group('MovingMNIST', 48 | 'Configure MovingMNIST dataset.') 49 | group.add_argument( 50 | '--num_distractors', 51 | dest='num_distractors', 52 | help="Number of noise distractors for MovingMNIST++", 53 | type=int) 54 | 55 | 56 | def parse_dataset_args(args): 57 | if args.num_distractors: 58 | cfg.MOVINGMNIST.DISTRACTOR_NUM = args.num_distractors 59 | 60 | 61 | ### Training 62 | def train(args): 63 | base_dir = get_base_dir(args) 64 | 65 | ### Get modules 66 | generator_net, loss_net = construct_modules(args) 67 | 68 | ### Prepare data 69 | mnist_iter = MovingMNISTAdvancedIterator( 70 | distractor_num=cfg.MOVINGMNIST.DISTRACTOR_NUM, 71 | initial_velocity_range=(cfg.MOVINGMNIST.VELOCITY_LOWER, 72 | cfg.MOVINGMNIST.VELOCITY_UPPER), 73 | rotation_angle_range=(cfg.MOVINGMNIST.ROTATION_LOWER, 74 | cfg.MOVINGMNIST.ROTATION_UPPER), 75 | scale_variation_range=(cfg.MOVINGMNIST.SCALE_VARIATION_LOWER, 76 | cfg.MOVINGMNIST.SCALE_VARIATION_UPPER), 77 | illumination_factor_range=(cfg.MOVINGMNIST.ILLUMINATION_LOWER, 78 | cfg.MOVINGMNIST.ILLUMINATION_UPPER)) 79 | 80 | for i in range(cfg.MODEL.TRAIN.MAX_ITER): 81 | seq, flow = mnist_iter.sample( 82 | batch_size=cfg.MODEL.TRAIN.BATCH_SIZE, 83 | seqlen=cfg.MOVINGMNIST.IN_LEN + cfg.MOVINGMNIST.OUT_LEN) 84 | in_seq = seq[:cfg.MOVINGMNIST.IN_LEN, ...] 85 | gt_seq = seq[cfg.MOVINGMNIST.IN_LEN:(cfg.MOVINGMNIST.IN_LEN + 86 | cfg.MOVINGMNIST.OUT_LEN), ...] 87 | 88 | # Transform data to NCDHW shape needed for 3D Convolution encoder and normalize 89 | context_nd = mx.nd.array(in_seq) / 255.0 90 | gt_nd = mx.nd.array(gt_seq) / 255.0 91 | context_nd = mx.nd.transpose(context_nd, axes=(1, 2, 0, 3, 4)) 92 | gt_nd = mx.nd.transpose(gt_nd, axes=(1, 2, 0, 3, 4)) 93 | 94 | # Train a step 95 | pred_nd, avg_l2, avg_real_mse, generator_grad_norm =\ 96 | train_step(generator_net, loss_net, context_nd, gt_nd) 97 | 98 | # Logging 99 | logging.info(( 100 | "Iter:{}, L2 Loss:{}, MSE Error:{}, Generator Grad Norm:{}").format( 101 | i, avg_l2, avg_real_mse, generator_grad_norm)) 102 | 103 | logging.info("Iter:%d" % i) 104 | if (i + 1) % 100 == 0: 105 | save_gif(context_nd.asnumpy()[0, 0, :, :, :], 106 | os.path.join(base_dir, "input.gif")) 107 | save_gif(gt_nd.asnumpy()[0, 0, :, :, :], 108 | os.path.join(base_dir, "gt.gif")) 109 | save_gif(pred_nd.asnumpy()[0, 0, :, :, :], 110 | os.path.join(base_dir, "pred.gif")) 111 | if cfg.MODEL.SAVE_ITER > 0 and (i + 1) % cfg.MODEL.SAVE_ITER == 0: 112 | generator_net.save_checkpoint( 113 | prefix=os.path.join(base_dir, "generator"), epoch=i) 114 | 115 | 116 | # Testing 117 | def recursive_generation(args): 118 | assert (cfg.MOVINGMNIST.OUT_LEN == 1) 119 | 120 | base_dir = get_base_dir(args) 121 | 122 | # Get modules 123 | generator_net, = construct_modules(args) 124 | 125 | # Prepare data 126 | mnist_iter = MovingMNISTAdvancedIterator( 127 | distractor_num=cfg.MOVINGMNIST.DISTRACTOR_NUM, 128 | initial_velocity_range=(cfg.MOVINGMNIST.VELOCITY_LOWER, 129 | cfg.MOVINGMNIST.VELOCITY_UPPER), 130 | rotation_angle_range=(cfg.MOVINGMNIST.ROTATION_LOWER, 131 | cfg.MOVINGMNIST.ROTATION_UPPER), 132 | scale_variation_range=(cfg.MOVINGMNIST.SCALE_VARIATION_LOWER, 133 | cfg.MOVINGMNIST.SCALE_VARIATION_UPPER), 134 | illumination_factor_range=(cfg.MOVINGMNIST.ILLUMINATION_LOWER, 135 | cfg.MOVINGMNIST.ILLUMINATION_UPPER)) 136 | 137 | frames = np.empty( 138 | shape=(cfg.MODEL.TEST.MAX_ITER, cfg.MODEL.TRAIN.BATCH_SIZE, 139 | cfg.MOVINGMNIST.TESTING_LEN, 1, cfg.MOVINGMNIST.IMG_SIZE, 140 | cfg.MOVINGMNIST.IMG_SIZE)) 141 | 142 | for i in range(cfg.MODEL.TEST.MAX_ITER): 143 | seq, flow = mnist_iter.sample( 144 | batch_size=cfg.MODEL.TRAIN.BATCH_SIZE, 145 | seqlen=cfg.MOVINGMNIST.IN_LEN + cfg.MOVINGMNIST.TESTING_LEN) 146 | in_seq = seq[:cfg.MOVINGMNIST.IN_LEN, ...] 147 | gt_seq = seq[cfg.MOVINGMNIST.IN_LEN:(cfg.MOVINGMNIST.IN_LEN + 148 | cfg.MOVINGMNIST.TESTING_LEN), ...] 149 | 150 | # Transform data to NCDHW shape needed for 3D Convolution encoder and normalize 151 | context_nd = mx.nd.array(in_seq) / 255.0 152 | gt_nd = mx.nd.array(gt_seq) / 255.0 153 | context_nd = mx.nd.transpose(context_nd, axes=(1, 2, 0, 3, 4)) 154 | gt_nd = mx.nd.transpose(gt_nd, axes=(1, 2, 0, 3, 4)) 155 | 156 | # Train a step 157 | frames[i] = test_step(generator_net, context_nd) 158 | 159 | frames = frames.reshape( 160 | -1, 161 | cfg.MOVINGMNIST.TESTING_LEN, 162 | cfg.MOVINGMNIST.IMG_SIZE, 163 | cfg.MOVINGMNIST.IMG_SIZE, ) 164 | save_gifs(frames, os.path.join(base_dir, "pred")) 165 | 166 | 167 | def test(args): 168 | base_dir = get_base_dir(args) 169 | logging_config(folder=base_dir, name='testing') 170 | 171 | # Get modules 172 | generator_net, loss_net = construct_modules(args) 173 | 174 | # Prepare data 175 | mnist_iter = MovingMNISTAdvancedIterator( 176 | distractor_num=cfg.MOVINGMNIST.DISTRACTOR_NUM, 177 | initial_velocity_range=(cfg.MOVINGMNIST.VELOCITY_LOWER, 178 | cfg.MOVINGMNIST.VELOCITY_UPPER), 179 | rotation_angle_range=(cfg.MOVINGMNIST.ROTATION_LOWER, 180 | cfg.MOVINGMNIST.ROTATION_UPPER), 181 | scale_variation_range=(cfg.MOVINGMNIST.SCALE_VARIATION_LOWER, 182 | cfg.MOVINGMNIST.SCALE_VARIATION_UPPER), 183 | illumination_factor_range=(cfg.MOVINGMNIST.ILLUMINATION_LOWER, 184 | cfg.MOVINGMNIST.ILLUMINATION_UPPER)) 185 | num_samples, seqlen = mnist_iter.load(file=cfg.MOVINGMNIST.TEST_FILE) 186 | 187 | overall_mse = 0 188 | for iter_id in range(num_samples // cfg.MODEL.TRAIN.BATCH_SIZE): 189 | frame_dat, _ = mnist_iter.sample( 190 | batch_size=cfg.MODEL.TRAIN.BATCH_SIZE, seqlen=seqlen, random=False) 191 | 192 | context_nd = mx.nd.array( 193 | frame_dat[:cfg.MOVINGMNIST.IN_LEN], ctx=args.ctx[0]) / 255.0 194 | gt_nd = mx.nd.array( 195 | frame_dat[cfg.MOVINGMNIST.IN_LEN:], ctx=args.ctx[0]) / 255.0 196 | 197 | # Transform data to NCDHW shape needed for 3D Convolution encoder 198 | context_nd = mx.nd.transpose(context_nd, axes=(1, 2, 0, 3, 4)) 199 | gt_nd = mx.nd.transpose(gt_nd, axes=(1, 2, 0, 3, 4)) 200 | 201 | pred_nd = test_step(generator_net, context_nd) 202 | overall_mse += mx.nd.mean(mx.nd.square(pred_nd - gt_nd)).asscalar() 203 | print(iter_id, overall_mse / (iter_id + 1)) 204 | 205 | avg_mse = overall_mse / (num_samples // cfg.MODEL.TRAIN.BATCH_SIZE) 206 | with open(os.path.join(base_dir, 'result.txt'), 'w') as f: 207 | f.write(str(avg_mse)) 208 | print(base_dir, avg_mse) 209 | 210 | 211 | if __name__ == "__main__": 212 | args = argument_parser() 213 | if not cfg.MODEL.TESTING: 214 | train(args) 215 | else: 216 | test(args) 217 | -------------------------------------------------------------------------------- /nowcasting/operators/traj_rnn.py: -------------------------------------------------------------------------------- 1 | import mxnet as mx 2 | import logging 3 | from nowcasting.ops import * 4 | from nowcasting.operators.common import identity, grid_generator, group_add, constant, save_npy 5 | from nowcasting.operators.conv_rnn import BaseConvRNN 6 | import numpy as np 7 | 8 | 9 | def flow_conv(data, num_filter, flows, weight, bias, name): 10 | assert isinstance(flows, list) 11 | warpped_data = [] 12 | for i in range(len(flows)): 13 | flow = flows[i] 14 | grid = mx.sym.GridGenerator(data=-flow, transform_type="warp") 15 | ele_dat = mx.sym.BilinearSampler(data=data, grid=grid) 16 | warpped_data.append(ele_dat) 17 | data = mx.sym.concat(*warpped_data, dim=1) 18 | ret = mx.sym.Convolution(data=data, 19 | num_filter=num_filter, 20 | kernel=(1, 1), 21 | weight=weight, 22 | bias=bias, 23 | name=name) 24 | return ret 25 | 26 | 27 | class TrajGRU(BaseConvRNN): 28 | def __init__(self, b_h_w, num_filter, zoneout=0.0, L=5, 29 | i2h_kernel=(3, 3), i2h_stride=(1, 1), i2h_pad=(1, 1), 30 | h2h_kernel=(5, 5), h2h_dilate=(1, 1), 31 | act_type="leaky", 32 | prefix="TrajGRU", lr_mult=1.0): 33 | super(TrajGRU, self).__init__(num_filter=num_filter, 34 | b_h_w=b_h_w, 35 | h2h_kernel=h2h_kernel, 36 | h2h_dilate=h2h_dilate, 37 | i2h_kernel=i2h_kernel, 38 | i2h_pad=i2h_pad, 39 | i2h_stride=i2h_stride, 40 | act_type=act_type, 41 | prefix=prefix) 42 | self._L = L 43 | self._zoneout = zoneout 44 | self.i2f_conv1_weight = self.params.get("i2f_conv1_weight", lr_mult=lr_mult) 45 | self.i2f_conv1_bias = self.params.get("i2f_conv1_bias", lr_mult=lr_mult) 46 | self.h2f_conv1_weight = self.params.get("h2f_conv1_weight", lr_mult=lr_mult) 47 | self.h2f_conv1_bias = self.params.get("h2f_conv1_bias", lr_mult=lr_mult) 48 | self.f_conv2_weight = self.params.get("f_conv2_weight", lr_mult=lr_mult) 49 | self.f_conv2_bias = self.params.get("f_conv2_bias", lr_mult=lr_mult) 50 | if cfg.MODEL.TRAJRNN.INIT_GRID: 51 | logging.info("TrajGRU: Initialize Grid Using Zeros!") 52 | self.f_out_weight = self.params.get("f_out_weight", 53 | lr_mult=lr_mult * cfg.MODEL.TRAJRNN.FLOW_LR_MULT, 54 | init=mx.init.Zero()) 55 | self.f_out_bias = self.params.get("f_out_bias", 56 | lr_mult=lr_mult * cfg.MODEL.TRAJRNN.FLOW_LR_MULT, 57 | init=mx.init.Zero()) 58 | else: 59 | self.f_out_weight = self.params.get("f_out_weight", lr_mult=lr_mult) 60 | self.f_out_bias = self.params.get("f_out_bias", lr_mult=lr_mult) 61 | self.i2h_weight = self.params.get("i2h_weight", lr_mult=lr_mult) 62 | self.i2h_bias = self.params.get("i2h_bias", lr_mult=lr_mult) 63 | self.h2h_weight = self.params.get("h2h_weight", lr_mult=lr_mult) 64 | self.h2h_bias = self.params.get("h2h_bias", lr_mult=lr_mult) 65 | 66 | @property 67 | def state_postfix(self): 68 | return ['h'] 69 | 70 | @property 71 | def state_info(self): 72 | return [{'shape': (self._batch_size, self._num_filter, 73 | self._state_height, self._state_width), 74 | '__layout__': "NCHW"}] 75 | 76 | def _flow_generator(self, inputs, states, prefix): 77 | if inputs is not None: 78 | i2f_conv1 = mx.sym.Convolution(data=inputs, 79 | weight=self.i2f_conv1_weight, 80 | bias=self.i2f_conv1_bias, 81 | kernel=(5, 5), 82 | dilate=(1, 1), 83 | pad=(2, 2), 84 | num_filter=32, 85 | name="%s_i2f_conv1" % prefix) 86 | else: 87 | i2f_conv1 = None 88 | h2f_conv1 = mx.sym.Convolution(data=states, 89 | weight=self.h2f_conv1_weight, 90 | bias=self.h2f_conv1_bias, 91 | kernel=(5, 5), 92 | dilate=(1, 1), 93 | pad=(2, 2), 94 | num_filter=32, 95 | name="%s_h2f_conv1" % prefix) 96 | f_conv1 = i2f_conv1 + h2f_conv1 if i2f_conv1 is not None else h2f_conv1 97 | f_conv1 = activation(f_conv1, act_type=self._act_type) 98 | # f_conv2 = mx.sym.Convolution(data=f_conv1, 99 | # weight=self.f_conv2_weight, 100 | # bias=self.f_conv2_bias, 101 | # kernel=(5, 5), 102 | # dilate=(1, 1), 103 | # pad=(2, 2), 104 | # num_filter=32, 105 | # name="%s_f_conv2" %prefix) 106 | # f_conv2 = activation(f_conv2, act_type=self._act_type) 107 | flows = mx.sym.Convolution(data=f_conv1, 108 | weight=self.f_out_weight, 109 | bias=self.f_out_bias, 110 | kernel=(5, 5), 111 | pad=(2, 2), 112 | num_filter=self._L * 2) 113 | if cfg.MODEL.TRAJRNN.SAVE_MID_RESULTS: 114 | import os 115 | flows = save_npy(flows, save_name="%s_flow" %prefix, 116 | save_dir=os.path.join(cfg.MODEL.SAVE_DIR, "flows")) 117 | flows = mx.sym.split(flows, num_outputs=self._L, axis=1) 118 | flows = [flows[i] for i in range(self._L)] 119 | return flows 120 | 121 | def __call__(self, inputs, states=None, is_initial=False, ret_mid=False): 122 | self._counter += 1 123 | name = '%s_t%d' % (self._prefix, self._counter) 124 | if is_initial: 125 | states = self.begin_state()[0] 126 | else: 127 | states = states[0] 128 | assert states is not None 129 | if inputs is not None: 130 | i2h = mx.sym.Convolution(data=inputs, 131 | weight=self.i2h_weight, 132 | bias=self.i2h_bias, 133 | kernel=self._i2h_kernel, 134 | stride=self._i2h_stride, 135 | dilate=self._i2h_dilate, 136 | pad=self._i2h_pad, 137 | num_filter=self._num_filter * 3, 138 | name="%s_i2h" % name) 139 | i2h_slice = mx.sym.SliceChannel(i2h, num_outputs=3, axis=1) 140 | else: 141 | i2h_slice = None 142 | prev_h = states 143 | flows = self._flow_generator(inputs=inputs, states=states, prefix=name) 144 | # flows[0] = identity(flows[0], input_debug=True) 145 | h2h = flow_conv(data=prev_h, num_filter=self._num_filter * 3, flows=flows, 146 | weight=self.h2h_weight, bias=self.h2h_bias, name="%s_h2h" % name) 147 | h2h_slice = mx.sym.SliceChannel(h2h, num_outputs=3, axis=1) 148 | if i2h_slice is not None: 149 | reset_gate = mx.sym.Activation(i2h_slice[0] + h2h_slice[0], act_type="sigmoid", 150 | name=name + "_r") 151 | update_gate = mx.sym.Activation(i2h_slice[1] + h2h_slice[1], act_type="sigmoid", 152 | name=name + "_u") 153 | new_mem = activation(i2h_slice[2] + reset_gate * h2h_slice[2], 154 | act_type=self._act_type, 155 | name=name + "_h") 156 | else: 157 | reset_gate = mx.sym.Activation(h2h_slice[0], act_type="sigmoid", 158 | name=name + "_r") 159 | update_gate = mx.sym.Activation(h2h_slice[1], act_type="sigmoid", 160 | name=name + "_u") 161 | new_mem = activation(reset_gate * h2h_slice[2], 162 | act_type=self._act_type, 163 | name=name + "_h") 164 | next_h = update_gate * prev_h + (1 - update_gate) * new_mem 165 | if self._zoneout > 0.0: 166 | mask = mx.sym.Dropout(mx.sym.ones_like(prev_h), p=self._zoneout) 167 | next_h = mx.sym.where(mask, next_h, prev_h) 168 | self._curr_states = [next_h] 169 | if not ret_mid: 170 | return next_h, [next_h] 171 | else: 172 | return next_h, [next_h], [] 173 | -------------------------------------------------------------------------------- /experiments/movingmnist/mnist_rnn_main.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import os 3 | import random 4 | import logging 5 | import numpy as np 6 | import mxnet as mx 7 | import mxnet.ndarray as nd 8 | from mnist_rnn_factory import MovingMNISTFactory 9 | from nowcasting.config import cfg, cfg_from_file, save_cfg 10 | from nowcasting.helpers.gifmaker import save_gif 11 | from nowcasting.encoder_forecaster import encoder_forecaster_build_networks, train_step, EncoderForecasterStates 12 | from nowcasting.utils import parse_ctx, logging_config, load_params 13 | from nowcasting.movingmnist_iterator import MovingMNISTAdvancedIterator 14 | from nowcasting.helpers.ordered_easydict import OrderedEasyDict as edict 15 | 16 | 17 | # random.seed(12345) 18 | # mx.random.seed(930215) 19 | # np.random.seed(921206) 20 | 21 | # random.seed(1234) 22 | # mx.random.seed(93021) 23 | # np.random.seed(92120) 24 | 25 | random.seed(123) 26 | mx.random.seed(9302) 27 | np.random.seed(9212) 28 | 29 | 30 | def parse_args(): 31 | parser = argparse.ArgumentParser(description='Train the MovingMNIST++ dataset') 32 | parser.add_argument('--batch_size', dest='batch_size', help="batchsize of the training process", 33 | default=None, type=int) 34 | parser.add_argument('--cfg', dest='cfg_file', help='Training configuration file', default=None, type=str) 35 | parser.add_argument('--resume', help='Continue to train the previous model', action='store_true', 36 | default=False) 37 | parser.add_argument('--save_dir', help='The saving directory', required=True, type=str) 38 | parser.add_argument('--ctx', dest='ctx', help='Running Context. E.g `--ctx gpu` or `--ctx gpu0,gpu1` or `--ctx cpu`', 39 | type=str, default='gpu') 40 | parser.add_argument('--lr', dest='lr', help='learning rate', default=None, type=float) 41 | parser.add_argument('--wd', dest='wd', help='weight decay', default=None, type=float) 42 | parser.add_argument('--grad_clip', dest='grad_clip', help='gradient clipping threshold', 43 | default=None, type=float) 44 | args = parser.parse_args() 45 | args.ctx = parse_ctx(args.ctx) 46 | if args.cfg_file is not None: 47 | cfg_from_file(args.cfg_file, target=cfg) 48 | if args.batch_size is not None: 49 | cfg.MODEL.TRAIN.BATCH_SIZE = args.batch_size 50 | if args.lr is not None: 51 | cfg.MODEL.TRAIN.LR = args.lr 52 | if args.wd is not None: 53 | cfg.MODEL.TRAIN.WD = args.wd 54 | if args.grad_clip is not None: 55 | cfg.MODEL.TRAIN.GRAD_CLIP = args.grad_clip 56 | if args.wd is not None: 57 | cfg.MODEL.TRAIN.WD = args.wd 58 | cfg.MODEL.SAVE_DIR = args.save_dir 59 | logging.info(args) 60 | return args 61 | 62 | 63 | def save_movingmnist_cfg(dir_path): 64 | tmp_cfg = edict() 65 | tmp_cfg.MOVINGMNIST = cfg.MOVINGMNIST 66 | tmp_cfg.MODEL = cfg.MODEL 67 | save_cfg(dir_path=dir_path, source=tmp_cfg) 68 | 69 | 70 | def load_mnist_params(load_dir, load_iter, encoder_net, forecaster_net): 71 | logging.info("Loading parameters from {}, Iter = {}" 72 | .format(os.path.realpath(load_dir), load_iter)) 73 | encoder_arg_params, encoder_aux_params = load_params(prefix=os.path.join(load_dir, 74 | "encoder_net"), 75 | epoch=load_iter) 76 | encoder_net.init_params(arg_params=encoder_arg_params, aux_params=encoder_aux_params, 77 | allow_missing=False, force_init=True) 78 | forecaster_arg_params, forecaster_aux_params = load_params(prefix=os.path.join(load_dir, 79 | "forecaster_net"), 80 | epoch=load_iter) 81 | forecaster_net.init_params(arg_params=forecaster_arg_params, 82 | aux_params=forecaster_aux_params, 83 | allow_missing=False, force_init=True) 84 | logging.info("Loading Complete!") 85 | 86 | 87 | def mnist_get_prediction(data_nd, states, encoder_net, forecaster_net): 88 | encoder_net.forward(is_train=False, 89 | data_batch=mx.io.DataBatch(data=[data_nd] + states.get_encoder_states())) 90 | states.update(encoder_net.get_outputs()) 91 | # Forward Forecaster 92 | if cfg.MODEL.OUT_TYPE == "direct": 93 | forecaster_net.forward(is_train=False, 94 | data_batch=mx.io.DataBatch(data=states.get_forecaster_state())) 95 | else: 96 | last_frame_nd = data_nd[data_nd.shape[0] - 1] 97 | forecaster_net.forward(is_train=False, 98 | data_batch=mx.io.DataBatch(data=states.get_forecaster_state() + [last_frame_nd])) 99 | forecaster_outputs = forecaster_net.get_outputs() 100 | pred_nd = forecaster_outputs[0] 101 | return pred_nd 102 | 103 | 104 | def train(args): 105 | assert cfg.MODEL.FRAME_STACK == 1 and cfg.MODEL.FRAME_SKIP == 1 106 | base_dir = args.save_dir 107 | logging_config(folder=base_dir, name="training") 108 | save_movingmnist_cfg(base_dir) 109 | mnist_iter = MovingMNISTAdvancedIterator( 110 | distractor_num=cfg.MOVINGMNIST.DISTRACTOR_NUM, 111 | initial_velocity_range=(cfg.MOVINGMNIST.VELOCITY_LOWER, 112 | cfg.MOVINGMNIST.VELOCITY_UPPER), 113 | rotation_angle_range=(cfg.MOVINGMNIST.ROTATION_LOWER, 114 | cfg.MOVINGMNIST.ROTATION_UPPER), 115 | scale_variation_range=(cfg.MOVINGMNIST.SCALE_VARIATION_LOWER, 116 | cfg.MOVINGMNIST.SCALE_VARIATION_UPPER), 117 | illumination_factor_range=(cfg.MOVINGMNIST.ILLUMINATION_LOWER, 118 | cfg.MOVINGMNIST.ILLUMINATION_UPPER)) 119 | 120 | mnist_rnn = MovingMNISTFactory(batch_size=cfg.MODEL.TRAIN.BATCH_SIZE // len(args.ctx), 121 | in_seq_len=cfg.MODEL.IN_LEN, 122 | out_seq_len=cfg.MODEL.OUT_LEN) 123 | 124 | encoder_net, forecaster_net, loss_net = \ 125 | encoder_forecaster_build_networks( 126 | factory=mnist_rnn, 127 | context=args.ctx) 128 | t_encoder_net, t_forecaster_net, t_loss_net = \ 129 | encoder_forecaster_build_networks( 130 | factory=mnist_rnn, 131 | context=args.ctx[0], 132 | shared_encoder_net=encoder_net, 133 | shared_forecaster_net=forecaster_net, 134 | shared_loss_net=loss_net, 135 | for_finetune=True) 136 | encoder_net.summary() 137 | forecaster_net.summary() 138 | loss_net.summary() 139 | # Begin to load the model if load_dir is not empty 140 | if len(cfg.MODEL.LOAD_DIR) > 0: 141 | load_mnist_params(load_dir=cfg.MODEL.LOAD_DIR, load_iter=cfg.MODEL.LOAD_ITER, 142 | encoder_net=encoder_net, forecaster_net=forecaster_net) 143 | states = EncoderForecasterStates(factory=mnist_rnn, ctx=args.ctx[0]) 144 | states.reset_all() 145 | for info in mnist_rnn.init_encoder_state_info: 146 | assert info["__layout__"].find('N') == 0, "Layout=%s is not supported!" %info["__layout__"] 147 | iter_id = 0 148 | while iter_id < cfg.MODEL.TRAIN.MAX_ITER: 149 | frame_dat, _ = mnist_iter.sample(batch_size=cfg.MODEL.TRAIN.BATCH_SIZE, 150 | seqlen=cfg.MOVINGMNIST.IN_LEN + cfg.MOVINGMNIST.OUT_LEN) 151 | data_nd = mx.nd.array(frame_dat[0:cfg.MOVINGMNIST.IN_LEN, ...], ctx=args.ctx[0]) / 255.0 152 | target_nd = mx.nd.array( 153 | frame_dat[cfg.MODEL.IN_LEN:(cfg.MOVINGMNIST.IN_LEN + cfg.MOVINGMNIST.OUT_LEN), ...], 154 | ctx=args.ctx[0]) / 255.0 155 | train_step(batch_size=cfg.MODEL.TRAIN.BATCH_SIZE, 156 | encoder_net=encoder_net, forecaster_net=forecaster_net, 157 | loss_net=loss_net, init_states=states, 158 | data_nd=data_nd, gt_nd=target_nd, mask_nd=None, 159 | iter_id=iter_id) 160 | if (iter_id + 1) % 100 == 0: 161 | new_frame_dat, _ = mnist_iter.sample(batch_size=cfg.MODEL.TRAIN.BATCH_SIZE, 162 | seqlen=cfg.MOVINGMNIST.IN_LEN + cfg.MOVINGMNIST.OUT_LEN) 163 | data_nd = mx.nd.array(frame_dat[0:cfg.MOVINGMNIST.IN_LEN, ...], ctx=args.ctx[0]) / 255.0 164 | target_nd = mx.nd.array( 165 | frame_dat[cfg.MOVINGMNIST.IN_LEN:(cfg.MOVINGMNIST.IN_LEN + cfg.MOVINGMNIST.OUT_LEN), ...], 166 | ctx=args.ctx[0]) / 255.0 167 | pred_nd = mnist_get_prediction(data_nd=data_nd, states=states, 168 | encoder_net=encoder_net, forecaster_net=forecaster_net) 169 | save_gif(pred_nd.asnumpy()[:, 0, 0, :, :], os.path.join(base_dir, "pred.gif")) 170 | save_gif(data_nd.asnumpy()[:, 0, 0, :, :], os.path.join(base_dir, "in.gif")) 171 | save_gif(target_nd.asnumpy()[:, 0, 0, :, :], os.path.join(base_dir, "gt.gif")) 172 | if (iter_id + 1) % cfg.MODEL.SAVE_ITER == 0: 173 | encoder_net.save_checkpoint( 174 | prefix=os.path.join(base_dir, "encoder_net"), 175 | epoch=iter_id) 176 | forecaster_net.save_checkpoint( 177 | prefix=os.path.join(base_dir, "forecaster_net"), 178 | epoch=iter_id) 179 | iter_id += 1 180 | 181 | 182 | if __name__ == "__main__": 183 | args = parse_args() 184 | train(args) 185 | -------------------------------------------------------------------------------- /nowcasting/numba_accelerated.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from numba import jit, float32, boolean, int32, float64 3 | from nowcasting.hko_evaluation import rainfall_to_pixel 4 | from nowcasting.config import cfg 5 | 6 | @jit(float32(float32, float32, boolean)) 7 | def get_GDL_numba(prediction, truth, mask): 8 | """Accelerated version of get_GDL using numba(http://numba.pydata.org/) 9 | 10 | Parameters 11 | ---------- 12 | prediction 13 | truth 14 | mask 15 | 16 | Returns 17 | ------- 18 | gdl 19 | """ 20 | seqlen, batch_size, _, height, width = prediction.shape 21 | gdl = np.zeros(shape=(seqlen, batch_size), dtype=np.float32) 22 | for i in range(seqlen): 23 | for j in range(batch_size): 24 | for m in range(height): 25 | for n in range(width): 26 | if m + 1 < height: 27 | if mask[i][j][0][m+1][n] and mask[i][j][0][m][n]: 28 | pred_diff_h = abs(prediction[i][j][0][m+1][n] - 29 | prediction[i][j][0][m][n]) 30 | gt_diff_h = abs(truth[i][j][0][m+1][n] - truth[i][j][0][m][n]) 31 | gdl[i][j] += abs(pred_diff_h - gt_diff_h) 32 | if n + 1 < width: 33 | if mask[i][j][0][m][n+1] and mask[i][j][0][m][n]: 34 | pred_diff_w = abs(prediction[i][j][0][m][n+1] - 35 | prediction[i][j][0][m][n]) 36 | gt_diff_w = abs(truth[i][j][0][m][n+1] - truth[i][j][0][m][n]) 37 | gdl[i][j] += abs(pred_diff_w - gt_diff_w) 38 | return gdl 39 | 40 | 41 | def get_hit_miss_counts_numba(prediction, truth, mask, thresholds=None): 42 | """This function calculates the overall hits and misses for the prediction, which could be used 43 | to get the skill scores and threat scores: 44 | 45 | 46 | This function assumes the input, i.e, prediction and truth are 3-dim tensors, (timestep, row, col) 47 | and all inputs should be between 0~1 48 | 49 | Parameters 50 | ---------- 51 | prediction : np.ndarray 52 | Shape: (seq_len, batch_size, 1, height, width) 53 | truth : np.ndarray 54 | Shape: (seq_len, batch_size, 1, height, width) 55 | mask : np.ndarray or None 56 | Shape: (seq_len, batch_size, 1, height, width) 57 | 0 --> not use 58 | 1 --> use 59 | thresholds : list or tuple 60 | 61 | Returns 62 | ------- 63 | hits : np.ndarray 64 | (seq_len, batch_size, len(thresholds)) 65 | TP 66 | misses : np.ndarray 67 | (seq_len, batch_size, len(thresholds)) 68 | FN 69 | false_alarms : np.ndarray 70 | (seq_len, batch_size, len(thresholds)) 71 | FP 72 | correct_negatives : np.ndarray 73 | (seq_len, batch_size, len(thresholds)) 74 | TN 75 | """ 76 | if thresholds is None: 77 | thresholds = cfg.HKO.EVALUATION.THRESHOLDS 78 | assert 5 == prediction.ndim 79 | assert 5 == truth.ndim 80 | assert prediction.shape == truth.shape 81 | assert prediction.shape[2] == 1 82 | thresholds = [rainfall_to_pixel(thresholds[i]) for i in range(len(thresholds))] 83 | thresholds = sorted(thresholds) 84 | ret = _get_hit_miss_counts_numba(prediction=prediction, 85 | truth=truth, 86 | mask=mask, 87 | thresholds=thresholds) 88 | return ret[:, :, :, 0], ret[:, :, :, 1], ret[:, :, :, 2], ret[:, :, :, 3] 89 | 90 | 91 | @jit(int32(float32, float32, boolean, float32)) 92 | def _get_hit_miss_counts_numba(prediction, truth, mask, thresholds): 93 | seqlen, batch_size, _, height, width = prediction.shape 94 | threshold_num = len(thresholds) 95 | ret = np.zeros(shape=(seqlen, batch_size, threshold_num, 4), dtype=np.int32) 96 | 97 | for i in range(seqlen): 98 | for j in range(batch_size): 99 | for m in range(height): 100 | for n in range(width): 101 | if mask[i][j][0][m][n]: 102 | for k in range(threshold_num): 103 | bpred = prediction[i][j][0][m][n] >= thresholds[k] 104 | btruth = truth[i][j][0][m][n] >= thresholds[k] 105 | ind = (1 - btruth) * 2 + (1 - bpred) 106 | ret[i][j][k][ind] += 1 107 | # The above code is the same as: 108 | # ret[i][j][k][0] += bpred * btruth 109 | # ret[i][j][k][1] += (1 - bpred) * btruth 110 | # ret[i][j][k][2] += bpred * (1 - btruth) 111 | # ret[i][j][k][3] += (1 - bpred) * (1- btruth) 112 | return ret 113 | 114 | 115 | def get_balancing_weights_numba(data, mask, base_balancing_weights=None, thresholds=None): 116 | """Get the balancing weights 117 | 118 | Parameters 119 | ---------- 120 | data 121 | mask 122 | base_balancing_weights 123 | thresholds 124 | 125 | Returns 126 | ------- 127 | 128 | """ 129 | if thresholds is None: 130 | thresholds = cfg.HKO.EVALUATION.THRESHOLDS 131 | if base_balancing_weights is None: 132 | base_balancing_weights = cfg.HKO.EVALUATION.BALANCING_WEIGHTS 133 | assert data.shape[2] == 1 134 | thresholds = [rainfall_to_pixel(thresholds[i]) for i in range(len(thresholds))] 135 | thresholds = sorted(thresholds) 136 | ret = _get_balancing_weights_numba(data=data, 137 | mask=mask, 138 | base_balancing_weights=base_balancing_weights, 139 | thresholds=thresholds) 140 | return ret 141 | 142 | 143 | @jit(float32(float32, boolean, float32, float32)) 144 | def _get_balancing_weights_numba(data, mask, base_balancing_weights, thresholds): 145 | seqlen, batch_size, _, height, width = data.shape 146 | threshold_num = len(thresholds) 147 | ret = np.zeros(shape=(seqlen, batch_size, 1, height, width), dtype=np.float32) 148 | 149 | for i in range(seqlen): 150 | for j in range(batch_size): 151 | for m in range(height): 152 | for n in range(width): 153 | if mask[i][j][0][m][n]: 154 | ele = data[i][j][0][m][n] 155 | for k in range(threshold_num): 156 | if ele < thresholds[k]: 157 | ret[i][j][0][m][n] = base_balancing_weights[k] 158 | break 159 | if ele >= thresholds[threshold_num - 1]: 160 | ret[i][j][0][m][n] = base_balancing_weights[threshold_num] 161 | return ret 162 | 163 | if __name__ == '__main__': 164 | from nowcasting.hko_evaluation import get_GDL, get_hit_miss_counts, get_balancing_weights 165 | from numpy.testing import assert_allclose, assert_almost_equal 166 | 167 | prediction = np.random.uniform(size=(10, 16, 1, 480, 480)) 168 | truth = np.random.uniform(size=(10, 16, 1, 480, 480)) 169 | mask = np.random.randint(low=0, high=2, size=(10, 16, 1, 480, 480)).astype(np.bool) 170 | import time 171 | 172 | begin = time.time() 173 | gdl = get_GDL(prediction=prediction, truth=truth, mask=mask) 174 | end = time.time() 175 | 176 | print("numpy gdl:", end - begin) 177 | begin = time.time() 178 | gdl_numba = get_GDL_numba(prediction=prediction, truth=truth, mask=mask) 179 | end = time.time() 180 | print("numba gdl:", end - begin) 181 | # gdl_mx = mx_get_GDL(prediction=prediction, truth=truth, mask=mask) 182 | # print gdl_mx 183 | assert_allclose(gdl, gdl_numba, rtol=1E-4, atol=1E-3) 184 | 185 | begin = time.time() 186 | for i in range(5): 187 | hits, misses, false_alarms, correct_negatives = get_hit_miss_counts(prediction=prediction, 188 | truth=truth, 189 | mask=mask) 190 | end = time.time() 191 | print("numpy hits misses:", end - begin) 192 | 193 | begin = time.time() 194 | for i in range(5): 195 | hits_numba, misses_numba, false_alarms_numba, correct_negatives_numba = get_hit_miss_counts_numba( 196 | prediction=prediction, 197 | truth=truth, 198 | mask=mask) 199 | end = time.time() 200 | print("numba hits misses:", end - begin) 201 | print(np.abs(hits - hits_numba).max()) 202 | print(np.abs(misses - misses_numba).max(), np.abs(misses - misses_numba).argmax()) 203 | print(np.abs(false_alarms - false_alarms_numba).max(), 204 | np.abs(false_alarms - false_alarms_numba).argmax()) 205 | print(np.abs(correct_negatives - correct_negatives_numba).max(), 206 | np.abs(correct_negatives - correct_negatives_numba).argmax()) 207 | 208 | begin = time.time() 209 | for i in range(5): 210 | weights_npy = get_balancing_weights(data=truth, mask=mask, 211 | base_balancing_weights=None, thresholds=None) 212 | end = time.time() 213 | print("numpy balancing weights:", end - begin) 214 | 215 | begin = time.time() 216 | for i in range(5): 217 | weights_numba = get_balancing_weights_numba(data=truth, mask=mask, 218 | base_balancing_weights=None, thresholds=None) 219 | end = time.time() 220 | print("numba balancing weights:", end - begin) 221 | print("Inconsistent Number:", (np.abs(weights_npy - weights_numba) > 1E-5).sum()) 222 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | About 2 | ----- 3 | 4 | Source code of the paper [Deep learning for precipitation nowcasting: A benchmark and a new model](http://papers.nips.cc/paper/7145-deep-learning-for-precipitation-nowcasting-a-benchmark-and-a-new-model) 5 | 6 | If you use the code or find it helpful, please cite the following paper: 7 | ``` 8 | @inproceedings{xingjian2017deep, 9 | title={Deep learning for precipitation nowcasting: a benchmark and a new model}, 10 | author={Shi, Xingjian and Gao, Zhihan and Lausen, Leonard and Wang, Hao and Yeung, Dit-Yan and Wong, Wai-kin and Woo, Wang-chun}, 11 | booktitle={Advances in Neural Information Processing Systems}, 12 | year={2017} 13 | } 14 | ``` 15 | 16 | Installation 17 | ------------ 18 | 19 | **Requires Python 3.5 or newer!** 20 | 21 | Both Windows and Linux are supported. 22 | 23 | Install the package 24 | ```bash 25 | python3 setup.py develop 26 | # Use --user if you have no privilege 27 | python3 setup.py develop --user 28 | ``` 29 | 30 | You will also need the python plugin of opencv: 31 | ```bash 32 | pip3 install opencv-contrib-python 33 | ``` 34 | 35 | In addition, you will need to install FFMpeg + X264 (See FAQ). 36 | 37 | For windows users it may be difficult to install some required packages like 38 | numba, ffmpeg or opencv-python. We strongly recommend you to use 39 | [Anaconda](https://www.anaconda.com/download/) and install them by commands like 40 | `conda install numba`. To install opencv-python on windows, you can download the 41 | wheel file from https://www.lfd.uci.edu/~gohlke/pythonlibs/. 42 | 43 | If you want to run the deep models in the paper, e.g., TrajGRU, you will need to install [MXNet](https://github.com/apache/incubator-mxnet). We've tested our code under [MXNet v0.12.0](https://github.com/apache/incubator-mxnet/releases/tag/0.12.0). 44 | Also, in order to run the ROVER algorithm, install the python wrapper of VarFlow by following the guide in [VarFlow](https://github.com/sxjscience/HKO-7/tree/master/VarFlow). 45 | 46 | **IMPORTANT!** You are able to run the HKO-7 benchmark environment without MXNet or VarFlow. You can proceed to use the HKOIterator and HKOBenchmarkEnv after you have installed the python package + Opencv-Python + FFMpeg with X264 encoding enabled and have downloaded the data. (See sections below for more reference). 47 | 48 | MovingMNIST++ 49 | ------------- 50 | Run the following script to draw a sample from the MovingMNIST++ dataset 51 | ```bash 52 | python3 nowcasting/movingmnist_iterator.py 53 | ``` 54 | 55 | Also, you can view samples of the learned connection structure of different layers in the TrajGRU-L13 model: 56 | 57 | - For the encoder, lower-layers will capture lower-level motion features and higher layer will capture some more general motion features. We show one of the learned links for layer1, layer2 and layer3 (from left to right). 58 | 59 | 60 | 61 | 62 | - For the forecaster, higher-layers will generate more global movements and lower layer will generate motions with finer details. We show one of the learned links for layer3, layer2 and layer1 (from left to right). 63 | 64 | 65 | 66 | 67 | 68 | 69 | Download the HKO-7 Dataset and Use the Iterator 70 | ----------------------------------------------- 71 | Please note that our source code does not require HKO-7 Dataset to perform the computation. 72 | 73 | The Hong Kong Observatory (HKO) may provide universities and research institutes the HKO-7 dataset (images + masks) for academic research subject to agreement to the undertaking ([HKO-7_Dataset_Undertaking_fillable.pdf](https://github.com/sxjscience/HKO-7/blob/master/HKO-7_Dataset_Undertaking_fillable.pdf)) by a faculty or formal member of the institute, e.g. professor, lecturer, researcher. Any interested person please review the terms and conditions on the undertaking and, if agreeable, fill in the form, sign it and send an email as follow: 74 | 75 | ``` 76 | Subject: Request for HKO-7 Dataset 77 | ---------------------------------- 78 | Name: YOUR NAME 79 | Institution: YOUR INSTITUTION 80 | Attachment: Completed and Signed Undertaking Form 81 | Other Information: 82 | You can include other information if you want. 83 | ``` 84 | 85 | Preferred email address:(You can also contact anyone in our NIPS2017 paper) 86 | ``` 87 | swirls@hko.gov.hk 88 | ``` 89 | 90 | The email must be sent from an **official email address** ending with the domain name of the institute. As we need to remotely establish the identiy of the data requester, we regret for not being able to process requests sent from general email services, e.g. Gmail, Yahoo, iCloud. If you have difficulty sending official emails, please explain in the email message body. 91 | 92 | Interested undergraduate or post-graduate students please ask their supervisors for advice and instructions. 93 | 94 | Please allow a few weeks for processing the data request. 95 | 96 | After you've downloaded the datasets, extract and put the `radarPNG` or `radarPNG_mask` folders under the `hko_data` folder. To use your own path of `radarPNG` or `radarPNG_mask`, append your paths into the `possible_hko_png_paths` and `possible_hko_mask_paths` in https://github.com/sxjscience/HKO-7/blob/master/nowcasting/config.py. 97 | ```python 98 | possible_hko_png_paths = [os.path.join('E:\\datasets\\HKO-data\\radarPNG\\radarPNG'), 99 | os.path.join(__C.HKO_DATA_BASE_PATH, 'radarPNG'), 100 | YOUR_PNG_PATH] 101 | possible_hko_mask_paths = [os.path.join('E:\\datasets\\HKO-data\\radarPNG\\radarPNG_mask'), 102 | os.path.join(__C.HKO_DATA_BASE_PATH, 'radarPNG_mask'), 103 | YOUR_MASK_PATH] 104 | ``` 105 | 106 | Also, download the necessary files via 107 | ```bash 108 | python3 download_all.py 109 | # You can also force to redownload the dataset: 110 | python3 download_all.py --overwrite 111 | ``` 112 | 113 | You can then try to run the following script to test the FPS of hko iterator. Check whether all the mp4 files are generated successfully. If they are all empty, try to reinstall ffmpeg with x264 encoding enabled (See FAQ below). 114 | ```bash 115 | python3 nowcasting/hko_iterator.py 116 | ``` 117 | 118 | You can use the iterator to sample a random minibatch of radar echo sequence. There is also the `sequent` setting and you can refer to the examples in the `hko_iterator.py` 119 | ```python 120 | from nowcasting.hko_iterator import HKOIterator 121 | from nowcasting.config import cfg 122 | from nowcasting. 123 | 124 | train_hko_iter = HKOIterator(pd_path=cfg.HKO_PD.RAINY_TRAIN, 125 | sample_mode="random", 126 | seq_len=25) 127 | sample_sequence, sample_mask, sample_datetime_clips, new_start =\ 128 | train_hko_iter.sample(batch_size=8) 129 | ``` 130 | 131 | If you have not obtained the HKO-7 dataset and just want to run the TrajGRU model, you can comment out the lines in https://github.com/sxjscience/HKO-7/blob/master/nowcasting/config.py#L39-L41 and https://github.com/sxjscience/HKO-7/blob/master/nowcasting/config.py#L49-L51 and run the MovingMNIST++ experiments. 132 | 133 | Run the HKO-7 Benchmark Environment 134 | ----------------------------------- 135 | The general workflow of the benchmark environment is given in the following: 136 | ```python 137 | from nowcasting.config import cfg 138 | model = INITIALIZE_YOUR_MODEL 139 | mode = "fixed" # Can also be "online" 140 | env = HKOBenchmarkEnv(pd_path=cfg.HKO_PD.RAINTY_TEST, mode=mode) 141 | while not env.done: 142 | # Get the observation 143 | in_frame_dat, in_mask_dat, in_datetime_clips, out_datetime_clips, begin_new_episode, need_upload_prediction =\ 144 | env.get_observation(batch_size=1) 145 | # You can update your model if you are using the online setting 146 | if mode == "online": 147 | # Just an example, need not to be exactly like this 148 | model.update(frames=in_frame_dat, masks=in_mask_dat) 149 | model.store(frames=in_frame_dat, masks=in_mask_dat) 150 | # Running your algorithm to get the prediction 151 | if need_upload_predictoin: 152 | prediction = model.predict(frames=in_frame_dat, masks=in_mask_dat) 153 | # Upload prediction to the environment 154 | env.upload_prediction(prediction) 155 | # Save the evaluation result 156 | env.save_eval() 157 | ``` 158 | 159 | You can refer to the CSI, HSS, B-MSE, B-MAE scores in the saved evaluation file to have an overall understanding of your performance. 160 | 161 | Running Experiments in the Paper 162 | -------------------------------- 163 | Refer to the [MovingMNIST++ Experiment README](https://github.com/sxjscience/HKO-7/tree/master/experiments/movingmnist) and [HKO-7 Experiment README](https://github.com/sxjscience/HKO-7/tree/master/experiments/hko) 164 | 165 | FAQ 166 | --- 167 | 1. Install FFMpeg with X264 encoding 168 | 169 | ```bash 170 | # Install libx264 171 | git clone git://git.videolan.org/x264.git 172 | cd x264 173 | ./configure --enable-static --enable-shared --enable-mp4 --prefix=YOUR_INSTALL_LOCATION --extra-ldflags="-lswresample -lm -lz -llzma" 174 | make -j64 175 | make install 176 | 177 | cd .. 178 | 179 | # Install ffmpeg 180 | git clone http://source.ffmpeg.org/git/ffmpeg.git 181 | cd ffmpeg 182 | ./configure --enable-gpl --enable-libx264 --prefix=YOUR_INSTALL_LOCATION 183 | make -j64 184 | make install 185 | ``` 186 | 187 | Above commands were sufficient to install FFMpeg with X264 encoding on our 188 | servers. Please refer to the official guide 189 | https://trac.ffmpeg.org/wiki/CompilationGuide/Ubuntu in case of any issues. 190 | -------------------------------------------------------------------------------- /experiments/hko/deconvolution.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import os 3 | import logging 4 | import numpy as np 5 | import mxnet as mx 6 | from nowcasting.config import cfg, save_cfg 7 | from nowcasting.utils import logging_config 8 | from nowcasting.hko_iterator import HKOIterator, precompute_mask 9 | from nowcasting.hko_benchmark import HKOBenchmarkEnv 10 | from nowcasting.models.deconvolution import \ 11 | construct_l2_loss, construct_modules, \ 12 | train_step, test_step, get_base_dir, \ 13 | model_args, training_args, mode_args, \ 14 | parse_mode_args, parse_training_args, parse_model_args 15 | from nowcasting.numba_accelerated import get_balancing_weights_numba 16 | 17 | 18 | # Arguments 19 | def argument_parser(): 20 | parser = argparse.ArgumentParser( 21 | description='Deconvolution baseline for HKO') 22 | 23 | cfg.DATASET = "HKO" 24 | 25 | mode_args(parser) 26 | training_args(parser) 27 | dataset_args(parser) 28 | model_args(parser) 29 | 30 | args = parser.parse_args() 31 | 32 | parse_mode_args(args) 33 | parse_training_args(args) 34 | parse_model_args(args) 35 | 36 | base_dir = get_base_dir(args) 37 | logging_config(folder=base_dir, name="training") 38 | save_cfg(base_dir, source=cfg.MODEL) 39 | 40 | logging.info(args) 41 | return args 42 | 43 | 44 | def dataset_args(parser): 45 | group = parser.add_argument_group('HKO', 'Configure HKO dataset.') 46 | group.add_argument( 47 | '--dataset', 48 | default="test", 49 | help="Whether to used the test set or the validation set.", 50 | type=str) 51 | 52 | 53 | # Benchmark 54 | def hko_benchmark(ctx, 55 | generator_net, 56 | loss_net, 57 | sample_num, 58 | finetune=False, 59 | mode="fixed", 60 | save_dir="hko7_rnn", 61 | pd_path=cfg.HKO_PD.RAINY_TEST): 62 | """Run the HKO7 Benchmark given the training sequences 63 | 64 | Args: 65 | ctx 66 | generator_net 67 | sample_num 68 | save_dir 69 | pd_path 70 | """ 71 | logging.info("Begin Evaluation, sample_num=%d," 72 | " results will be saved to %s" % (sample_num, save_dir)) 73 | if finetune: 74 | logging.info(str(cfg.MODEL.TEST.ONLINE)) 75 | env = HKOBenchmarkEnv(pd_path=pd_path, save_dir=save_dir, mode=mode) 76 | 77 | if finetune: 78 | assert (mode == "online") 79 | data_buffer = [] 80 | stored_prediction = [] 81 | finetune_iter = 0 82 | 83 | context_nd = None 84 | 85 | i = 0 86 | while not env.done: 87 | logging.info("Iter {} of evaluation.".format(i)) 88 | i += 1 89 | if finetune: 90 | if len(data_buffer) >= 5: 91 | context_np = data_buffer[0] # HKO.BENCHMARK.IN_LEN frames 92 | gt_np = np.concatenate(data_buffer[1:], axis=0) 93 | gt_np = gt_np[:cfg.HKO.BENCHMARK.OUT_LEN] 94 | 95 | mask_np = precompute_mask(gt_np) 96 | 97 | weights = get_balancing_weights_numba( 98 | data=gt_np, 99 | mask=mask_np, 100 | base_balancing_weights=cfg.HKO.EVALUATION. 101 | BALANCING_WEIGHTS, 102 | thresholds=env._all_eval._thresholds) 103 | weighted_mse = (weights * 104 | np.square(stored_prediction[0] - gt_np)).sum( 105 | axis=(2, 3, 4)) 106 | mean_weighted_mse = weighted_mse.mean() 107 | print("mean_weighted_mse = %g" % mean_weighted_mse) 108 | 109 | if mean_weighted_mse > cfg.MODEL.TEST.ONLINE.FINETUNE_MIN_MSE: 110 | context_nd = mx.nd.array(context_np, ctx=ctx) 111 | context_nd = mx.nd.transpose( 112 | context_nd, axes=(1, 2, 0, 3, 4)) 113 | gt_nd = mx.nd.array(gt_np, ctx=ctx) 114 | gt_nd = mx.nd.transpose(gt_nd, axes=(1, 2, 0, 3, 4)) 115 | mask_nd = mx.nd.array(mask_np, ctx=ctx) 116 | mask_nd = mx.nd.transpose(mask_nd, axes=(1, 2, 0, 3, 4)) 117 | 118 | train_step( 119 | generator_net=generator_net, 120 | loss_net=loss_net, 121 | context_nd=context_nd, 122 | gt_nd=gt_nd, 123 | mask_nd=mask_nd) 124 | 125 | finetune_iter += 1 126 | 127 | del data_buffer[0] 128 | del stored_prediction[0] 129 | 130 | if mode == "online": 131 | context_np, in_datetime_clips, out_datetime_clips,\ 132 | begin_new_episode, need_upload_prediction = env.get_observation( 133 | batch_size=1) 134 | context_np = np.repeat( 135 | context_np, cfg.MODEL.TRAIN.BATCH_SIZE, axis=1) 136 | orig_size = 1 137 | 138 | elif mode == "fixed": 139 | context_np, in_datetime_clips, out_datetime_clips,\ 140 | begin_new_episode, need_upload_prediction = env.get_observation( 141 | batch_size=cfg.MODEL.TRAIN.BATCH_SIZE) 142 | context_nd = mx.nd.array(context_np, ctx=ctx) 143 | context_nd = mx.nd.transpose(context_nd, axes=(1, 2, 0, 3, 4)) 144 | 145 | # Pad context_nd up to batch size if needed 146 | orig_size = context_nd.shape[0] 147 | while context_nd.shape[0] < cfg.MODEL.TRAIN.BATCH_SIZE: 148 | context_nd = mx.nd.concat( 149 | context_nd, context_nd[0:1], num_args=2, dim=0) 150 | else: 151 | raise NotImplementedError 152 | 153 | if finetune: 154 | if begin_new_episode: 155 | data_buffer = [context_np] 156 | prediction_buffer = [] 157 | else: 158 | data_buffer.append(context_np) 159 | 160 | if mode != "fixed": 161 | context_nd = mx.nd.array(context_np, ctx=ctx) 162 | context_nd = mx.nd.transpose(context_nd, axes=(1, 2, 0, 3, 4)) 163 | generator_net.forward( 164 | is_train=False, data_batch=mx.io.DataBatch(data=[context_nd])) 165 | 166 | if need_upload_prediction: 167 | generator_outputs = dict( 168 | zip(generator_net.output_names, generator_net.get_outputs())) 169 | pred_nd = generator_outputs["pred_output"] 170 | 171 | pred_nd = pred_nd[0:orig_size] 172 | 173 | pred_nd = mx.nd.clip(pred_nd, a_min=0, a_max=1) 174 | pred_nd = mx.nd.transpose(pred_nd, axes=(2, 0, 1, 3, 4)) 175 | 176 | env.upload_prediction(prediction=pred_nd.asnumpy()) 177 | 178 | if finetune: 179 | stored_prediction.append(pred_nd.asnumpy()) 180 | 181 | env.save_eval() 182 | 183 | 184 | # Training 185 | def train(args): 186 | base_dir = get_base_dir(args) 187 | 188 | # Get modules 189 | generator_net, loss_net, = construct_modules(args) 190 | 191 | # Prepare data 192 | train_hko_iter = HKOIterator( 193 | pd_path=cfg.HKO_PD.RAINY_TRAIN, 194 | sample_mode="random", 195 | seq_len=cfg.HKO.BENCHMARK.IN_LEN + cfg.HKO.BENCHMARK.OUT_LEN) 196 | 197 | start_iter = 0 if not cfg.MODEL.RESUME else cfg.MODEL.LOAD_ITER 198 | for i in range(start_iter, cfg.MODEL.TRAIN.MAX_ITER): 199 | frame_dat, mask_dat, datetime_clips, _ = train_hko_iter.sample( 200 | batch_size=cfg.MODEL.TRAIN.BATCH_SIZE) 201 | 202 | context_nd = mx.nd.array( 203 | frame_dat[0:cfg.HKO.BENCHMARK.IN_LEN, ...], 204 | ctx=args.ctx[0]) / 255.0 205 | gt_nd = mx.nd.array( 206 | frame_dat[cfg.HKO.BENCHMARK.IN_LEN:(cfg.HKO.BENCHMARK.IN_LEN + 207 | cfg.HKO.BENCHMARK.OUT_LEN)], 208 | ctx=args.ctx[0]) / 255.0 209 | mask_nd = mx.nd.array( 210 | mask_dat[cfg.HKO.BENCHMARK.IN_LEN:(cfg.HKO.BENCHMARK.IN_LEN + 211 | cfg.HKO.BENCHMARK.OUT_LEN)], 212 | ctx=args.ctx[0]) 213 | 214 | # Transform data to NCDHW shape needed for 3D Convolution encoder and normalize 215 | context_nd = mx.nd.transpose(context_nd, axes=(1, 2, 0, 3, 4)) 216 | gt_nd = mx.nd.transpose(gt_nd, axes=(1, 2, 0, 3, 4)) 217 | mask_nd = mx.nd.transpose(mask_nd, axes=(1, 2, 0, 3, 4)) 218 | 219 | # Train a step 220 | pred_nd, avg_l2, avg_real_mse, generator_grad_norm =\ 221 | train_step(generator_net, loss_net, context_nd, gt_nd, mask_nd) 222 | 223 | if (i + 1) % cfg.MODEL.VALID_ITER == 0: 224 | hko_benchmark( 225 | ctx=args.ctx[0], 226 | generator_net=generator_net, 227 | loss_net=loss_net, 228 | sample_num=i, 229 | save_dir=os.path.join(base_dir, "iter{}_valid".format(i + 1)), 230 | pd_path=cfg.HKO_PD.RAINY_VALID) 231 | 232 | # Logging 233 | logging.info(( 234 | "Iter:{}, L2 Loss:{}, MSE Error:{}, Generator Grad Norm:{}").format( 235 | i, avg_l2, avg_real_mse, generator_grad_norm)) 236 | 237 | if cfg.MODEL.SAVE_ITER > 0 and (i + 1) % cfg.MODEL.SAVE_ITER == 0: 238 | generator_net.save_checkpoint( 239 | prefix=os.path.join(base_dir, "generator"), epoch=i) 240 | 241 | hko_benchmark( 242 | ctx=args.ctx[0], 243 | generator_net=generator_net, 244 | loss_net=loss_net, 245 | sample_num=i, 246 | save_dir=os.path.join(base_dir, "iter{}_test".format(i + 1)), 247 | pd_path=cfg.HKO_PD.RAINY_TEST) 248 | 249 | 250 | def test(args): 251 | assert (args.resume is True) 252 | if cfg.MODEL.TEST.FINETUNE: 253 | assert (cfg.MODEL.DECONVBASELINE.BN_GLOBAL_STATS is True) 254 | 255 | base_dir = args.save_dir 256 | logging_config(folder=base_dir, name="testing") 257 | save_cfg(dir_path=base_dir, source=cfg.MODEL) 258 | 259 | generator_net, loss_net, = construct_modules(args) 260 | 261 | if args.dataset == "test": 262 | pd_path = cfg.HKO_PD.RAINY_TEST 263 | elif args.dataset == "valid": 264 | pd_path = cfg.HKO_PD.RAINY_VALID 265 | else: 266 | raise NotImplementedError 267 | 268 | hko_benchmark( 269 | ctx=args.ctx[0], 270 | generator_net=generator_net, 271 | loss_net=loss_net, 272 | sample_num=1, 273 | save_dir=os.path.join(base_dir, "iter{}_{}_finetune{}".format( 274 | cfg.MODEL.LOAD_ITER + 1, args.dataset, cfg.MODEL.TEST.FINETUNE)), 275 | finetune=cfg.MODEL.TEST.FINETUNE, 276 | mode=cfg.MODEL.TEST.MODE, 277 | pd_path=pd_path) 278 | 279 | 280 | if __name__ == "__main__": 281 | args = argument_parser() 282 | if not cfg.MODEL.TESTING: 283 | train(args) 284 | else: 285 | test(args) 286 | -------------------------------------------------------------------------------- /VarFlow/example.cpp: -------------------------------------------------------------------------------- 1 | //------------------------------------------------------------------ 2 | // Released under the BDS License 3 | // 4 | // Located at http://sourceforge.net/projects/varflow 5 | // 6 | //------------------------------------------------------------------ 7 | 8 | //--------------------------------------------------------------------------------------------------------- 9 | // 10 | // Sample program demonstrating usage of the VarFlow class to calculate optical flow 11 | // Two images from the classic Yosemite flyby sequence are used to generate a dense optical flow field. 12 | // The original image is displayed, along with vector and color field representations of the underlying flow field. 13 | // 14 | // Author: Adam Harmat 15 | // Date: April 2009 16 | // 17 | //---------------------------------------------------------------------------------------------------------- 18 | 19 | 20 | #include 21 | #include 22 | #include "VarFlow.h" 23 | //#include "ProfTimer.h" 24 | 25 | using namespace std; 26 | 27 | #define M_PI 3.14159265358979323846 // pi 28 | 29 | // Calculates RGB values from HSV color space 30 | void hsv2rgb(float h, float s, float v, uchar &r, uchar &g, uchar &b){ 31 | 32 | if(h > 360) 33 | { 34 | h = h - 360; 35 | } 36 | 37 | float c = v*s; // chroma 38 | float hp = h / 60; 39 | 40 | float hpmod2 = hp - (float)((int)(hp/2))*2; 41 | 42 | float x = c*(1 - fabs(hpmod2 - 1)); 43 | float m = v - c; 44 | 45 | float r1, g1, b1; 46 | 47 | if(0 <= hp && hp < 1){ 48 | r1 = c; 49 | g1 = x; 50 | b1 = 0; 51 | } 52 | else if(1 <= hp && hp < 2){ 53 | r1 = x; 54 | g1 = c; 55 | b1 = 0; 56 | } 57 | else if(2 <= hp && hp < 3){ 58 | r1 = 0; 59 | g1 = c; 60 | b1 = x; 61 | } 62 | else if(3 <= hp && hp < 4){ 63 | r1 = 0; 64 | g1 = x; 65 | b1 = c; 66 | } 67 | else if(4 <= hp && hp < 5){ 68 | r1 = x; 69 | g1 = 0; 70 | b1 = c; 71 | } 72 | else 73 | { 74 | r1 = c; 75 | g1 = 0; 76 | b1 = x; 77 | } 78 | 79 | r = (uchar)(255*(r1 + m)); 80 | g = (uchar)(255*(g1 + m)); 81 | b = (uchar)(255*(b1 + m)); 82 | 83 | } 84 | 85 | // Draw a vector field based on horizontal and vertical flow fields 86 | void drawMotionField(IplImage* imgU, IplImage* imgV, IplImage* imgMotion, int xSpace, int ySpace, float cutoff, int multiplier, CvScalar color) 87 | { 88 | int x, y; 89 | 90 | CvPoint p0 = cvPoint(0,0); 91 | CvPoint p1 = cvPoint(0,0); 92 | 93 | float deltaX, deltaY, angle, hyp; 94 | 95 | for(y = ySpace; y < imgU->height; y+= ySpace ) { 96 | for(x = xSpace; x < imgU->width; x+= xSpace ){ 97 | 98 | p0.x = x; 99 | p0.y = y; 100 | 101 | deltaX = *((float*)(imgU->imageData + y*imgU->widthStep)+x); 102 | deltaY = -(*((float*)(imgV->imageData + y*imgV->widthStep)+x)); 103 | 104 | angle = atan2(deltaY, deltaX); 105 | hyp = sqrt(deltaX*deltaX + deltaY*deltaY); 106 | 107 | if(hyp > cutoff){ 108 | 109 | p1.x = p0.x + cvRound(multiplier*hyp*cos(angle)); 110 | p1.y = p0.y + cvRound(multiplier*hyp*sin(angle)); 111 | 112 | cvLine( imgMotion, p0, p1, color,1, CV_AA, 0); 113 | 114 | p0.x = p1.x + cvRound(3*cos(angle-M_PI + M_PI/4)); 115 | p0.y = p1.y + cvRound(3*sin(angle-M_PI + M_PI/4)); 116 | cvLine( imgMotion, p0, p1, color,1, CV_AA, 0); 117 | 118 | p0.x = p1.x + cvRound(3*cos(angle-M_PI - M_PI/4)); 119 | p0.y = p1.y + cvRound(3*sin(angle-M_PI - M_PI/4)); 120 | cvLine( imgMotion, p0, p1, color,1, CV_AA, 0); 121 | } 122 | 123 | } 124 | } 125 | 126 | } 127 | 128 | // Draws the circular legend for the color field, indicating direction and magnitude 129 | void drawLegendHSV(IplImage* imgColor, int radius, int cx, int cy) 130 | { 131 | int width = radius*2 + 1; 132 | int height = width; 133 | 134 | IplImage* imgLegend = cvCreateImage( cvSize(width, height), 8, 3 ); 135 | IplImage* imgMask = cvCreateImage( cvSize(width, height), 8, 1 ); 136 | IplImage* sub_img = cvCreateImageHeader(cvSize(width, height),8,3); 137 | 138 | uchar* legend_ptr; 139 | float angle, h, s, v, legend_max_s; 140 | uchar r,g,b; 141 | int deltaX, deltaY; 142 | 143 | legend_max_s = radius*sqrt(2); 144 | 145 | for(int y=0; y < imgLegend->height; y++) 146 | { 147 | legend_ptr = (uchar*)(imgLegend->imageData + y*imgLegend->widthStep); 148 | 149 | for(int x=0; x < imgLegend->width; x++) 150 | { 151 | deltaX = x-radius; 152 | deltaY = -(y-radius); 153 | angle = atan2(deltaY,deltaX); 154 | 155 | if(angle < 0) 156 | angle += 2*M_PI; 157 | 158 | h = angle * 180 / M_PI; 159 | s = sqrt(deltaX*deltaX + deltaY*deltaY) / legend_max_s; 160 | v = 0.9; 161 | 162 | hsv2rgb(h, s, v, r, g, b); 163 | 164 | legend_ptr[3*x] = b; 165 | legend_ptr[3*x+1] = g; 166 | legend_ptr[3*x+2] = r; 167 | 168 | } 169 | } 170 | 171 | cvZero(imgMask); 172 | cvCircle( imgMask, cvPoint(radius,radius) , radius, CV_RGB(255,255,255), -1,8,0 ); 173 | 174 | sub_img->origin = imgColor->origin; 175 | sub_img->widthStep = imgColor->widthStep; 176 | sub_img->imageData = imgColor->imageData + (cy-radius) * imgColor->widthStep + (cx-radius) * imgColor->nChannels; 177 | 178 | cvCopy(imgLegend, sub_img, imgMask); 179 | 180 | cvCircle( imgColor, cvPoint(cx,cy) , radius, CV_RGB(0,0,0), 1,CV_AA,0 ); 181 | 182 | cvReleaseImage(&imgLegend); 183 | cvReleaseImage(&imgMask); 184 | cvReleaseImageHeader(&sub_img); 185 | 186 | } 187 | 188 | 189 | // Draws a color field representation of the flow field 190 | void drawColorField(IplImage* imgU, IplImage* imgV, IplImage* imgColor) 191 | { 192 | IplImage* imgColorHSV = cvCreateImage( cvSize(imgColor->width, imgColor->height), IPL_DEPTH_32F, 3 ); 193 | cvZero(imgColorHSV); 194 | 195 | float max_s = 0; 196 | float *hsv_ptr, *u_ptr, *v_ptr; 197 | uchar *color_ptr; 198 | float angle; 199 | float h,s,v; 200 | uchar r,g,b; 201 | float deltaX, deltaY; 202 | 203 | int x, y; 204 | 205 | // Generate hsv image 206 | for(y = 0; y < imgColorHSV->height; y++ ) { 207 | 208 | hsv_ptr = (float*)(imgColorHSV->imageData + y*imgColorHSV->widthStep); 209 | u_ptr = (float*)(imgU->imageData + y*imgU->widthStep); 210 | v_ptr = (float*)(imgV->imageData + y*imgV->widthStep); 211 | 212 | for(x = 0; x < imgColorHSV->width; x++){ 213 | 214 | deltaX = u_ptr[x]; 215 | deltaY = v_ptr[x]; 216 | 217 | angle = atan2(deltaY,deltaX); 218 | 219 | if(angle < 0) 220 | angle += 2*M_PI; 221 | 222 | hsv_ptr[3*x] = angle * 180 / M_PI; 223 | hsv_ptr[3*x+1] = sqrt(deltaX*deltaX + deltaY*deltaY); 224 | hsv_ptr[3*x+2] = 0.9; 225 | 226 | if(hsv_ptr[3*x+1] > max_s) 227 | max_s = hsv_ptr[3*x+1]; 228 | 229 | } 230 | } 231 | 232 | // Generate color image 233 | for(y = 0; y < imgColor->height; y++ ) { 234 | 235 | hsv_ptr = (float*)(imgColorHSV->imageData + y*imgColorHSV->widthStep); 236 | color_ptr = (uchar*)(imgColor->imageData + y*imgColor->widthStep); 237 | 238 | for(x = 0; x < imgColor->width; x++){ 239 | 240 | h = hsv_ptr[3*x]; 241 | s = hsv_ptr[3*x+1] / max_s; 242 | v = hsv_ptr[3*x+2]; 243 | 244 | hsv2rgb(h, s, v, r, g, b); 245 | 246 | color_ptr[3*x] = b; 247 | color_ptr[3*x+1] = g; 248 | color_ptr[3*x+2] = r; 249 | 250 | } 251 | } 252 | 253 | drawLegendHSV(imgColor, 25, 28, 28); 254 | 255 | cvReleaseImage(&imgColorHSV); 256 | 257 | } 258 | 259 | 260 | int main(int argc, char *argv[]) 261 | { 262 | // Load sample images 263 | IplImage* imgA = cvLoadImage( "Data/yos_img_08.jpg", 0 ); 264 | IplImage* imgB = cvLoadImage( "Data/yos_img_09.jpg", 0 ); 265 | 266 | int width = imgA->width; 267 | int height = imgA->height; 268 | 269 | IplImage* imgU = cvCreateImage(cvSize(width, height), IPL_DEPTH_32F, 1); 270 | IplImage* imgV = cvCreateImage(cvSize(width, height), IPL_DEPTH_32F, 1); 271 | IplImage* imgColor = cvCreateImage( cvSize(width, height), 8, 3 ); 272 | IplImage* imgMotion = cvCreateImage(cvSize(width, height), 8, 3); 273 | 274 | IplImage* swap_img; 275 | 276 | //We will start at level 0 (full size image) and go down to level 4 (coarse image 16 times smaller than original) 277 | //Experiment with these values to see how they affect the flow field as well as calculation time 278 | int max_level = 4; 279 | int start_level = 0; 280 | 281 | //Two pre and post smoothing steps, should be greater than zero 282 | int n1 = 2; 283 | int n2 = 2; 284 | 285 | //Smoothing and regularization parameters, experiment but keep them above zero 286 | float rho = 2.8; 287 | float alpha = 1400; 288 | float sigma = 1.5; 289 | 290 | char c; 291 | 292 | // Set up VarFlow class 293 | VarFlow OpticalFlow(width, height, max_level, start_level, n1, n2, rho, alpha, sigma); 294 | 295 | //ProfTimer t; 296 | 297 | cvZero(imgU); 298 | cvZero(imgV); 299 | cvZero(imgMotion); 300 | cvZero(imgColor); 301 | 302 | // Start timing 303 | //t.Start(); 304 | 305 | // Calculate the flow 306 | OpticalFlow.CalcFlow(imgA, imgB, imgU, imgV, 0); 307 | 308 | // Stop timing 309 | //t.Stop(); 310 | //double dur = t.GetDurationInSecs(); 311 | //cout<<"Executing CalcFlow took "<>> x = mx.nd.ones((2, 3)) 33 | >>> y = mx.nd.ones((5, 6)) 34 | >>> z = mx.nd.ones((4, 2, 3)) 35 | >>> print(nd_global_norm([x, y, z]).asscalar()) 36 | 7.74597 37 | >>> xnone = None 38 | >>> ret = nd_global_norm([x, y, z, xnone]) 39 | >>> print(ret.asscalar()) 40 | 7.74597 41 | """ 42 | ret = None 43 | for arr in t_list: 44 | if arr is not None: 45 | if ret is None: 46 | ret = nd.square(nd.norm(arr)) 47 | else: 48 | ret += nd.square(nd.norm(arr)) 49 | ret = nd.sqrt(ret) 50 | return ret 51 | 52 | 53 | class MyModule(Module): 54 | """Some enhancement to the mx.mod.Module 55 | 56 | """ 57 | 58 | def __init__(self, symbol, data_names=('data',), label_names=('softmax_label',), 59 | logger=logging, context=mx.context.gpu(), work_load_list=None, 60 | fixed_param_names=None, state_names=None, name=None): 61 | self._name = name 62 | super(MyModule, self).__init__(symbol=symbol, 63 | data_names=data_names, 64 | label_names=label_names, 65 | logger=logger, 66 | context=context, 67 | work_load_list=work_load_list, 68 | fixed_param_names=fixed_param_names, 69 | state_names=state_names) 70 | self._tmp_grads = None 71 | 72 | def clip_by_global_norm(self, max_norm=1.0): 73 | """Clips gradient norm. 74 | The norm is computed over all gradients together, as if they were 75 | concatenated into a single vector. Gradients are modified in-place. 76 | The method is first used in 77 | `[ICML2013] On the difficulty of training recurrent neural networks` 78 | Parameters 79 | ---------- 80 | max_norm : float or int 81 | The maximum clipping threshold of the gradient norm. 82 | Returns 83 | ------- 84 | norm_val : float 85 | The computed norm of the gradients. 86 | Examples 87 | -------- 88 | An example of using clip_grad_norm to clip the gradient before updating the parameters:: 89 | >>> #Get the gradient via back-propagation 90 | >>> net.forward_backward(data_batch=data_batch) 91 | >>> norm_val = net.clip_by_global_norm(max_norm=1.0) 92 | >>> net.update() 93 | """ 94 | assert self.binded and self.params_initialized and self.optimizer_initialized 95 | norm_val = self.global_grad_norm() 96 | if norm_val > max_norm: 97 | ratio = max_norm / float(norm_val) 98 | for grads in self._exec_group.grad_arrays: 99 | for grad in grads: 100 | grad *= ratio 101 | return norm_val 102 | 103 | def global_grad_norm(self): 104 | """Calculate global gradient norm. 105 | The L2 norm is computed over all gradients together, as if they were 106 | concatenated into a single vector. 107 | Could be used to debug the optimization process. 108 | See http://videolectures.net/deeplearning2015_goodfellow_network_optimization/ 109 | Returns 110 | ------- 111 | norm_val : float 112 | The computed norm of the gradients. 113 | Examples 114 | -------- 115 | An example of using global_norm to calculate the gradient norm after back-propgation:: 116 | >>> #Get the gradient via back-propagation 117 | >>> net.forward_backward(data_batch=data_batch) 118 | >>> norm_val = net.global_grad_norm() 119 | >>> print(norm_val) 120 | """ 121 | assert self.binded and self.params_initialized and self.optimizer_initialized 122 | # The code in the following will cause the estimated norm to be different for multiple gpus 123 | norm_val = 0.0 124 | for exe in self._exec_group.execs: 125 | norm_val += nd_global_norm(exe.grad_arrays).asscalar() 126 | norm_val /= float(len(self._exec_group.execs)) 127 | norm_val *= self._optimizer.rescale_grad 128 | return norm_val 129 | 130 | def debug_norm_all(self, debug_gnorm=True): 131 | if debug_gnorm: 132 | for k, v, grad_v in zip(self._param_names, self._exec_group.param_arrays, 133 | self._exec_group.grad_arrays): 134 | logging.debug("%s: v-norm: %g, g-norm: %g" 135 | %(k, 136 | nd.norm(v[0]).asnumpy()[0], 137 | nd.norm(grad_v[0]).asnumpy()[0])) 138 | else: 139 | for k, v in zip(self._param_names, self._exec_group.param_arrays): 140 | logging.debug("%s: v-norm: %g" 141 | %(k, 142 | nd.norm(v[0]).asnumpy()[0])) 143 | 144 | def summary(self, level=2): 145 | """Summarize the network parameters. 146 | 147 | Parameters 148 | ---------- 149 | level : int, optional 150 | Level of the summarization logs to print. 151 | The log becomes more verbose with higher summary level. 152 | - Level = 0 153 | Print the total param number + aux param number 154 | - Level = 1 155 | Print the shape of all parameters + The total number of paremter numbers 156 | - Level = 2 157 | Print the shape of the data/state and other available information in Level 1 158 | """ 159 | self.logger.info("Summary of %s" %self._name) 160 | assert self.binded and self.params_initialized 161 | assert 0 <= level <= 2, \ 162 | "Level must be between 0 and 2, level=%d is not supported" % level 163 | 164 | def _log_var(key, value, typ="param"): 165 | if typ == "param": 166 | if k in self._fixed_param_names: 167 | self.logger.info(" %s: %s, %d, req = %s, fixed" 168 | % (key, 169 | str(value.shape), 170 | np.prod(value.shape), 171 | self._exec_group.grad_req[k])) 172 | else: 173 | self.logger.info(" %s: %s, %d, req = %s" 174 | % (key, 175 | str(value.shape), 176 | np.prod(value.shape), 177 | self._exec_group.grad_req[k])) 178 | elif typ == "data" or typ == "aux": 179 | self.logger.info(" %s: %s, %d" 180 | % (key, 181 | str(value.shape), 182 | np.prod(value.shape))) 183 | 184 | total_param_num = 0 185 | total_fixed_param_num = 0 186 | total_aux_param_num = 0 187 | if level >= 2: 188 | if len(self.data_names) == 0: 189 | self.logger.info("Data: None") 190 | else: 191 | self.logger.info("Data:") 192 | for k, v in zip(self.data_names, self.data_shapes): 193 | _log_var(k, v, typ="data") 194 | if len(self._state_names) == 0: 195 | self.logger.info("State: None") 196 | else: 197 | self.logger.info("State:") 198 | for k in self._state_names: 199 | v = self._exec_group.execs[0].arg_dict[k] 200 | _log_var(k, v, typ="data") 201 | if level >= 1: 202 | if len(self._param_names) == 0: 203 | self.logger.info("Param: None") 204 | else: 205 | self.logger.info("Params:") 206 | for k in self._param_names: 207 | v = self._arg_params[k] 208 | _log_var(k, v) 209 | if k in self._fixed_param_names: 210 | total_fixed_param_num += np.prod(v.shape) 211 | else: 212 | total_param_num += np.prod(v.shape) 213 | if len(self._aux_names) == 0: 214 | self.logger.info("Aux States: None") 215 | else: 216 | self.logger.info("Aux States: ") 217 | for k in self._aux_names: 218 | v = self._aux_params[k] 219 | _log_var(k, v, typ="aux") 220 | total_aux_param_num += np.prod(v.shape) 221 | else: 222 | for k in self._param_names: 223 | v = self._arg_params[k] 224 | total_param_num += np.prod(v.shape) 225 | for k in self._aux_names: 226 | v = self._aux_params[k] 227 | total_aux_param_num += np.prod(v.shape) 228 | self.logger.info("Total Param Num (exclude fixed ones): " + str(total_param_num)) 229 | self.logger.info("Total Fixed Param Num: " + str(total_fixed_param_num)) 230 | self.logger.info("Total Aux Param Num: " + str(total_aux_param_num)) 231 | 232 | def get_output_dict(self): 233 | outputs = self.get_outputs() 234 | return OrderedDict([(k, v) for k, v in zip(self._output_names, outputs)]) 235 | 236 | def clear_grad(self): 237 | assert self.binded and self.params_initialized and self.optimizer_initialized 238 | # clear the gradient 239 | for grads in self._exec_group.grad_arrays: 240 | for grad in grads: 241 | grad[:] = 0 242 | 243 | def save_tmp_grad(self): 244 | if self._tmp_grads is None: 245 | self._tmp_grads = [] 246 | for grads in self._exec_group.grad_arrays: 247 | vec = [] 248 | for grad in grads: 249 | vec.append(grad.copyto(grad.context)) 250 | self._tmp_grads.append(vec) 251 | else: 252 | for i, grads in enumerate(self._exec_group.grad_arrays): 253 | for j, grad in enumerate(grads): 254 | self._tmp_grads[i][j][:] = grad 255 | 256 | def acc_grad_with_tmp(self): 257 | assert self._tmp_grads is not None 258 | for i, grads in enumerate(self._exec_group.grad_arrays): 259 | for j, grad in enumerate(grads): 260 | grad += self._tmp_grads[i][j] 261 | 262 | 263 | def load_params_allow_missing(self, fname): 264 | """Loads model parameters from file. 265 | 266 | Parameters 267 | ---------- 268 | fname : str 269 | Path to input param file. 270 | 271 | Examples 272 | -------- 273 | >>> # An example of loading module parameters. 274 | >>> mod.load_params('myfile') 275 | """ 276 | logging.info("Load Param From %s" %fname) 277 | save_dict = mx.nd.load(fname) 278 | arg_params = {} 279 | aux_params = {} 280 | for k, value in save_dict.items(): 281 | arg_type, name = k.split(':', 1) 282 | if arg_type == 'arg': 283 | if name in self._param_names: 284 | logging.info("set %s" %name) 285 | arg_params[name] = value 286 | elif arg_type == 'aux': 287 | if name in self._aux_names: 288 | logging.info("set %s" % name) 289 | aux_params[name] = value 290 | else: 291 | raise ValueError("Invalid param file " + fname) 292 | self.set_params(arg_params, aux_params, allow_missing=True) 293 | --------------------------------------------------------------------------------