add flownet2

models/flownet2_pytorch/LICENSE

@@ -0,0 +1,13 @@
+Copyright 2017 NVIDIA CORPORATION
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.

models/flownet2_pytorch/README.md

@@ -0,0 +1,116 @@
+# flownet2-pytorch 
+
+Pytorch implementation of [FlowNet 2.0: Evolution of Optical Flow Estimation with Deep Networks](https://arxiv.org/abs/1612.01925). 
+
+Multiple GPU training is supported, and the code provides examples for training or inference on [MPI-Sintel](http://sintel.is.tue.mpg.de/) clean and final datasets. The same commands can be used for training or inference with other datasets. See below for more detail.
+
+Inference using fp16 (half-precision) is also supported.
+
+For more help, type <br />
+
+    python main.py --help
+
+## Network architectures
+Below are the different flownet neural network architectures that are provided. <br />
+A batchnorm version for each network is also available.
+
+ - **FlowNet2S**
+ - **FlowNet2C**
+ - **FlowNet2CS**
+ - **FlowNet2CSS**
+ - **FlowNet2SD**
+ - **FlowNet2**
+
+## Custom layers
+
+`FlowNet2` or `FlowNet2C*` achitectures rely on custom layers `Resample2d` or `Correlation`. <br />
+A pytorch implementation of these layers with cuda kernels are available at [./networks](./networks). <br />
+Note : Currently, half precision kernels are not available for these layers.
+
+## Data Loaders
+
+Dataloaders for FlyingChairs, FlyingThings, ChairsSDHom and ImagesFromFolder are available in [datasets.py](./datasets.py). <br />
+
+## Loss Functions
+
+L1 and L2 losses with multi-scale support are available in [losses.py](./losses.py). <br />
+
+## Installation 
+
+    # get flownet2-pytorch source
+    git clone https://github.com/NVIDIA/flownet2-pytorch.git
+    cd flownet2-pytorch
+
+    # install custom layers
+    bash install.sh
+
+### Python requirements 
+Currently, the code supports python 3
+* numpy 
+* PyTorch ( == 0.4.1, for <= 0.4.0 see branch [python36-PyTorch0.4](https://github.com/NVIDIA/flownet2-pytorch/tree/python36-PyTorch0.4))
+* scipy 
+* scikit-image
+* tensorboardX
+* colorama, tqdm, setproctitle 
+
+## Converted Caffe Pre-trained Models
+We've included caffe pre-trained models. Should you use these pre-trained weights, please adhere to the [license agreements](https://drive.google.com/file/d/1TVv0BnNFh3rpHZvD-easMb9jYrPE2Eqd/view?usp=sharing). 
+
+* [FlowNet2](https://drive.google.com/file/d/1hF8vS6YeHkx3j2pfCeQqqZGwA_PJq_Da/view?usp=sharing)[620MB]
+* [FlowNet2-C](https://drive.google.com/file/d/1BFT6b7KgKJC8rA59RmOVAXRM_S7aSfKE/view?usp=sharing)[149MB]
+* [FlowNet2-CS](https://drive.google.com/file/d/1iBJ1_o7PloaINpa8m7u_7TsLCX0Dt_jS/view?usp=sharing)[297MB]
+* [FlowNet2-CSS](https://drive.google.com/file/d/157zuzVf4YMN6ABAQgZc8rRmR5cgWzSu8/view?usp=sharing)[445MB]
+* [FlowNet2-CSS-ft-sd](https://drive.google.com/file/d/1R5xafCIzJCXc8ia4TGfC65irmTNiMg6u/view?usp=sharing)[445MB]
+* [FlowNet2-S](https://drive.google.com/file/d/1V61dZjFomwlynwlYklJHC-TLfdFom3Lg/view?usp=sharing)[148MB]
+* [FlowNet2-SD](https://drive.google.com/file/d/1QW03eyYG_vD-dT-Mx4wopYvtPu_msTKn/view?usp=sharing)[173MB]
+
+## Inference
+    # Example on MPISintel Clean   
+    python main.py --inference --model FlowNet2 --save_flow --inference_dataset MpiSintelClean \
+    --inference_dataset_root /path/to/mpi-sintel/clean/dataset \
+    --resume /path/to/checkpoints 
+
+## Training and validation
+
+    # Example on MPISintel Final and Clean, with L1Loss on FlowNet2 model
+    python main.py --batch_size 8 --model FlowNet2 --loss=L1Loss --optimizer=Adam --optimizer_lr=1e-4 \
+    --training_dataset MpiSintelFinal --training_dataset_root /path/to/mpi-sintel/final/dataset  \
+    --validation_dataset MpiSintelClean --validation_dataset_root /path/to/mpi-sintel/clean/dataset
+
+    # Example on MPISintel Final and Clean, with MultiScale loss on FlowNet2C model 
+    python main.py --batch_size 8 --model FlowNet2C --optimizer=Adam --optimizer_lr=1e-4 --loss=MultiScale --loss_norm=L1 \
+    --loss_numScales=5 --loss_startScale=4 --optimizer_lr=1e-4 --crop_size 384 512 \
+    --training_dataset FlyingChairs --training_dataset_root /path/to/flying-chairs/dataset  \
+    --validation_dataset MpiSintelClean --validation_dataset_root /path/to/mpi-sintel/clean/dataset
+
+## Results on MPI-Sintel
+[![Predicted flows on MPI-Sintel](./image.png)](https://www.youtube.com/watch?v=HtBmabY8aeU "Predicted flows on MPI-Sintel")
+
+## Reference 
+If you find this implementation useful in your work, please acknowledge it appropriately and cite the paper:
+````
+@InProceedings{IMKDB17,
+  author       = "E. Ilg and N. Mayer and T. Saikia and M. Keuper and A. Dosovitskiy and T. Brox",
+  title        = "FlowNet 2.0: Evolution of Optical Flow Estimation with Deep Networks",
+  booktitle    = "IEEE Conference on Computer Vision and Pattern Recognition (CVPR)",
+  month        = "Jul",
+  year         = "2017",
+  url          = "http://lmb.informatik.uni-freiburg.de//Publications/2017/IMKDB17"
+}
+````
+```
+@misc{flownet2-pytorch,
+  author = {Fitsum Reda and Robert Pottorff and Jon Barker and Bryan Catanzaro},
+  title = {flownet2-pytorch: Pytorch implementation of FlowNet 2.0: Evolution of Optical Flow Estimation with Deep Networks},
+  year = {2017},
+  publisher = {GitHub},
+  journal = {GitHub repository},
+  howpublished = {\url{https://github.com/NVIDIA/flownet2-pytorch}}
+}
+```
+## Related Optical Flow Work from Nvidia 
+Code (in Caffe and Pytorch): [PWC-Net](https://github.com/NVlabs/PWC-Net) <br />
+Paper : [PWC-Net: CNNs for Optical Flow Using Pyramid, Warping, and Cost Volume](https://arxiv.org/abs/1709.02371). 
+
+## Acknowledgments
+Parts of this code were derived, as noted in the code, from [ClementPinard/FlowNetPytorch](https://github.com/ClementPinard/FlowNetPytorch).

models/flownet2_pytorch/convert.py

@@ -0,0 +1,137 @@
+#!/usr/bin/env python2.7
+
+import caffe
+from caffe.proto import caffe_pb2
+import sys, os
+
+import torch
+import torch.nn as nn
+
+import argparse, tempfile
+import numpy as np
+
+parser = argparse.ArgumentParser()
+parser.add_argument('caffe_model', help='input model in hdf5 or caffemodel format')
+parser.add_argument('prototxt_template',help='prototxt template')
+parser.add_argument('flownet2_pytorch', help='path to flownet2-pytorch')
+
+args = parser.parse_args()
+
+args.rgb_max = 255
+args.fp16 = False
+args.grads = {}
+
+# load models
+sys.path.append(args.flownet2_pytorch)
+
+import models
+from utils.param_utils import *
+
+width = 256
+height = 256
+keys = {'TARGET_WIDTH': width, 
+        'TARGET_HEIGHT': height,
+        'ADAPTED_WIDTH':width,
+        'ADAPTED_HEIGHT':height,
+        'SCALE_WIDTH':1.,
+        'SCALE_HEIGHT':1.,}
+
+template = '\n'.join(np.loadtxt(args.prototxt_template, dtype=str, delimiter='\n'))
+for k in keys:
+    template = template.replace('$%s$'%(k),str(keys[k]))
+
+prototxt = tempfile.NamedTemporaryFile(mode='w', delete=True)
+prototxt.write(template)
+prototxt.flush()
+
+net = caffe.Net(prototxt.name, args.caffe_model, caffe.TEST)
+
+weights = {}
+biases = {}
+
+for k, v in list(net.params.items()):
+    weights[k] = np.array(v[0].data).reshape(v[0].data.shape)
+    biases[k] = np.array(v[1].data).reshape(v[1].data.shape)
+    print((k, weights[k].shape, biases[k].shape))
+
+if 'FlowNet2/' in args.caffe_model:
+    model = models.FlowNet2(args)
+
+    parse_flownetc(model.flownetc.modules(), weights, biases)
+    parse_flownets(model.flownets_1.modules(), weights, biases, param_prefix='net2_')
+    parse_flownets(model.flownets_2.modules(), weights, biases, param_prefix='net3_')
+    parse_flownetsd(model.flownets_d.modules(), weights, biases, param_prefix='netsd_')
+    parse_flownetfusion(model.flownetfusion.modules(), weights, biases, param_prefix='fuse_')
+
+    state = {'epoch': 0,
+             'state_dict': model.state_dict(),
+             'best_EPE': 1e10}
+    torch.save(state, os.path.join(args.flownet2_pytorch, 'FlowNet2_checkpoint.pth.tar'))
+
+elif 'FlowNet2-C/' in args.caffe_model:
+    model = models.FlowNet2C(args)
+
+    parse_flownetc(model.modules(), weights, biases)
+    state = {'epoch': 0,
+             'state_dict': model.state_dict(),
+             'best_EPE': 1e10}
+    torch.save(state, os.path.join(args.flownet2_pytorch, 'FlowNet2-C_checkpoint.pth.tar'))
+
+elif 'FlowNet2-CS/' in args.caffe_model:
+    model = models.FlowNet2CS(args)
+
+    parse_flownetc(model.flownetc.modules(), weights, biases)
+    parse_flownets(model.flownets_1.modules(), weights, biases, param_prefix='net2_')
+
+    state = {'epoch': 0,
+             'state_dict': model.state_dict(),
+             'best_EPE': 1e10}
+    torch.save(state, os.path.join(args.flownet2_pytorch, 'FlowNet2-CS_checkpoint.pth.tar'))
+
+elif 'FlowNet2-CSS/' in args.caffe_model:
+    model = models.FlowNet2CSS(args)
+
+    parse_flownetc(model.flownetc.modules(), weights, biases)
+    parse_flownets(model.flownets_1.modules(), weights, biases, param_prefix='net2_')
+    parse_flownets(model.flownets_2.modules(), weights, biases, param_prefix='net3_')
+
+    state = {'epoch': 0,
+             'state_dict': model.state_dict(),
+             'best_EPE': 1e10}
+    torch.save(state, os.path.join(args.flownet2_pytorch, 'FlowNet2-CSS_checkpoint.pth.tar'))
+
+elif 'FlowNet2-CSS-ft-sd/' in args.caffe_model:
+    model = models.FlowNet2CSS(args)
+
+    parse_flownetc(model.flownetc.modules(), weights, biases)
+    parse_flownets(model.flownets_1.modules(), weights, biases, param_prefix='net2_')
+    parse_flownets(model.flownets_2.modules(), weights, biases, param_prefix='net3_')
+
+    state = {'epoch': 0,
+             'state_dict': model.state_dict(),
+             'best_EPE': 1e10}
+    torch.save(state, os.path.join(args.flownet2_pytorch, 'FlowNet2-CSS-ft-sd_checkpoint.pth.tar'))
+
+elif 'FlowNet2-S/' in args.caffe_model:
+    model = models.FlowNet2S(args)
+
+    parse_flownetsonly(model.modules(), weights, biases, param_prefix='')
+    state = {'epoch': 0,
+             'state_dict': model.state_dict(),
+             'best_EPE': 1e10}
+    torch.save(state, os.path.join(args.flownet2_pytorch, 'FlowNet2-S_checkpoint.pth.tar'))
+
+elif 'FlowNet2-SD/' in args.caffe_model:
+    model = models.FlowNet2SD(args)
+
+    parse_flownetsd(model.modules(), weights, biases, param_prefix='')
+
+    state = {'epoch': 0,
+             'state_dict': model.state_dict(),
+             'best_EPE': 1e10}
+    torch.save(state, os.path.join(args.flownet2_pytorch, 'FlowNet2-SD_checkpoint.pth.tar'))
+
+else:
+    print(('model type cound not be determined from input caffe model %s'%(args.caffe_model)))
+    quit()
+print(("done converting ", args.caffe_model))