-
Notifications
You must be signed in to change notification settings - Fork 0
RunningCode
- Download the code by running cloning this repository
git clone https://github.com/prgumd/ajna
You'll need the following dependencies to run our code.
- OpenCV 3.3
- TensorFlow 1.15 (GPU or CPU version) for TF1 version of TensorFlow 2.4 (GPU or CPU version) for TF2 version
- Appropriate Cuda and Cudnn version for your Tensorflow and Ubuntu version
- Matplotlib
- tqdm
- Numpy
- Termcolor
We've tested the code on both Ubuntu 20.04 using Tensorflow 1.15 on GPU using Python 3.6 and Ubuntu 22.0f using TensorFlow 2.4 on GPU using Python 3.9.
Now, follow the instructions below for training/testing each network.
The codes for TensorFlow 1 and 2 are in Code/FlowNet/TF1 and Code/FlowNet/TF2 respectively.
To train the Ajna network, you'll need to run Code\FlowNet\TFx\Train.py. Replace x with the 1 or 2 for TensorFlow 1 or 2. The pre-trained models can be downloaded from here. This network was trained with a LinearSoftplus loss from the Science paper on Flying Chairs 2 dataset for 400 epochs and fine tuned on Flying Things 3D dataset for 50 more epochs. They were used for all the experiments without any additional changes in the Science Robotics paper. The following command line flags are implemented for ease of use.
-
--BasePath: Base path from where dataset images are loaded, eg., /home/nitin/Datasets/FlyingChairs. -
--LabelBasePath: Path to labels if present in a different place as in the case of Flying Things 3D. -
--NumEpochs: Number of epochs the training will be done for -
--DivTrain: Factor to reduce Train data by per epoch, used for debugging only or for super large datasets. -
--MiniBatchSize: Size of the MiniBatch to use. -
--LoadCheckPoint: Load Model from latest Checkpoint from CheckPointPath? -
--LogsPath: Path to save logs, eg., /media/nitin/Research/Ajna/Logs/ -
--LossFuncName: Choice of Loss functions, Choice of Loss functions, choose SL1-1 for Supervised L1 loss for forward flow and so on. -
--CheckPointPath: Path to save checkpoints. -
--GPUDevice: What GPU do you want to use? -1 for CPU. -
--LR: Learning Rate. -
--UncType: Choose LinearSoftplus for Science robotics paper.
Other less frequently used flags are not mentioned in the readme.
To run the EVDeblurNet code, you'll need to run code\DeblurNetUnsup\RunEVDeblurNet.py. The following command line flags are implemented for ease of use.
-
--ModelPath: Path to load model from. eg., /media/nitin/Research/EVDodge/CheckpointsDeblurNet/199model.ckpt -
--ReadPath: Path to load images from. eg., /media/nitin/Research/EVDodge/DatasetChethanEvents/processed -
--WritePath: Path to write images to. eg., /media/nitin/Research/EVDodge/DatasetChethanEvents/Deblurred -
--GPUDevice: What GPU do you want to use? -1 for CPU.
To run the network from the paper:
Go to Code\FlowNet\TF2 folder and run:
python3 Train.py --BasePath=PATHTOFLYINGCHAIRS2PATH --NumEpochs=400 --MiniBatchSize=32 --UncType=LinearSoftplus --LossFuncName=SL1-1 --CheckPointPath=PATHTOSAVE --LR=1e-4 --Datset=FC2 --NetworkType=Large --NetworkName=Network.ResNet
Once this has trained, fine-tune on Flying Things 3D. Running the above codes assumes that you have downloaded the Flying Chairs 2 and Flying Things 3D datasets. Make sure to download the Optical flow ground truth separately for Flying Things 3D dataset.
python3 Train.py --BasePath=PATHTOFLYINGTHINGS3DPATH --LabelBasePath=PATHTOFLYINGTHINGS3DLABELS --NumEpochs=50 --MiniBatchSize=32 --UncType=LinearSoftplus --LossFuncName=SL1-1 --CheckPointPath=PATHTOSAVE --LR=5e-5 --Datset=FC2 --LoadCheckPoint=1 --NetworkType=Large --NetworkName=Network.ResNet
To run inference on the Ajna network, you'll need to run Code\FlowNet\TFx\TestSingle.py. Replace x with the 1 or 2 for TensorFlow 1 or 2. The pre-trained models can be downloaded from here. This network was trained with a LinearSoftplus loss from the Science paper on Flying Chairs 2 dataset for 400 epochs and fine tuned on Flying Things 3D dataset for 50 more epochs. They were used for all the experiments without any additional changes in the Science Robotics paper. The following command line flags are implemented for ease of use.
It has the same flags as the Training code with some additional ones explained below.
-
--Img1: Path to first image, eg., ../../../Data/Img1.png -
--Img2: Path to second image, eg., ../../../Data/Img2.png -
--ImgSize: Image size as a list, eg., [240,320,3] -
--GT1: Path to Ground Truth if you have for visualization -
--MiniBatchSize: Size of the MiniBatch to use. -
--CheckPointPath: Path to model, eg., /media/nitin/Education/PRGEyeOmni/Trained/SL1ResNetLR1e-4/199model.ckpt
To run inference of the code from the Science robotics paper run the following command:
python3 TestSingle.py --NetworkType=Large --NetworkName=Network.ResNet --CheckPointPath=PATHTO449.ckptMODEL --Img1=../../../Data/Img1.png --Img2=../../../Data/Img2.png
Optionally, you can also run inference on multiple frames using TestMultiple.py using the following command (Make sure image follows frame%06d formatting):
python3 TestMultiple.py --UncType=Aleatoric --ImgSize=[480,640,3] --ImgPath=PATH_TO_IMG_DIRECTORY WritePath=PATH_TO_SAVE_DIRECTORY --Skip=1