├── .github └── workflows │ ├── build.yml │ ├── lightning.yml │ ├── publish_pypi.yml │ ├── python.yml │ └── pytorch.yml ├── .gitignore ├── .readthedocs.yaml ├── CONTRIBUTING.md ├── LICENSE ├── MANIFEST.in ├── README.md ├── compare_paper_results.py ├── configs └── results │ ├── model_benchmark_all.yaml │ ├── validate_all.yaml │ └── validate_all_things.yaml ├── datasets.yaml ├── docs ├── Makefile ├── make.bat ├── requirements.txt └── source │ ├── _static │ ├── plot-memory(gb)-fp32-sintel-final-occ-val_epe.html │ ├── plot-memory(gb)-fp32-sintel-final-occ-val_flall.html │ ├── plot-params-flops.html │ ├── plot-sintel-final-occ-val_epe-kitti-2015-val_flall.html │ ├── plot-time(ms)-fp32-sintel-final-occ-val_epe.html │ └── plot-time(ms)-fp32-sintel-final-occ-val_flall.html │ ├── api │ ├── ptlflow.data.rst │ ├── ptlflow.models.rst │ ├── ptlflow.utils.rst │ ├── ptlflow │ │ ├── data │ │ │ ├── datasets.rst │ │ │ └── flow_transforms.rst │ │ ├── models │ │ │ ├── base_model.rst │ │ │ └── train_val_dataset.rst │ │ └── utils │ │ │ ├── callbacks │ │ │ └── logger.rst │ │ │ ├── correlation.rst │ │ │ ├── dummy_datasets.rst │ │ │ ├── flow_metrics.rst │ │ │ ├── flow_utils.rst │ │ │ ├── flowpy_torch.rst │ │ │ ├── io_adapter.rst │ │ │ ├── timer.rst │ │ │ └── utils.rst │ ├── scripts.rst │ └── scripts │ │ ├── infer.rst │ │ ├── model_benchmark.rst │ │ ├── summary_metrics.rst │ │ ├── train.rst │ │ └── validate.rst │ ├── conf.py │ ├── custom │ ├── new_dataset.rst │ └── new_model.rst │ ├── datasets │ └── datasets_list.rst │ ├── index.rst │ ├── models │ ├── checkpoint_list.csv │ ├── checkpoint_list.rst │ └── models_list.rst │ ├── results │ ├── accuracy_epe.rst │ ├── accuracy_epe_flall.rst │ ├── metrics_all.csv │ ├── metrics_all_things.csv │ ├── model_benchmark-all.csv │ ├── model_benchmark.rst │ ├── paper_ptlflow.rst │ ├── paper_ptlflow_metrics.csv │ ├── paper_results_things.csv │ ├── result_plots.rst │ ├── summarized_metrics-epe.csv │ ├── summarized_metrics-epe_flall.csv │ └── summarized_metrics-epe_outlier.csv │ └── starting │ ├── benchmark.rst │ ├── config_files.rst │ ├── inference.rst │ ├── installation.rst │ ├── testing.rst │ ├── training.rst │ ├── v04_upgrade_guide.rst │ └── validation.rst ├── infer.py ├── misc ├── create_compare_paper.sh ├── create_metric_summaries.sh └── create_plots.sh ├── model_benchmark.py ├── plot_results.py ├── ptlflow ├── __init__.py ├── data │ ├── AutoFlow_val.txt │ ├── FlyingChairs_val.txt │ ├── Hd1k_val.txt │ ├── Kitti2012_val.txt │ ├── Kitti2015_val.txt │ ├── Sintel_val.txt │ ├── __init__.py │ ├── datasets.py │ ├── flow_datamodule.py │ ├── flow_transforms.py │ └── split_autoflow.py ├── models │ ├── README.md │ ├── __init__.py │ ├── base_model │ │ ├── __init__.py │ │ └── base_model.py │ ├── ccmr │ │ ├── LICENSE │ │ ├── README.md │ │ ├── __init__.py │ │ ├── ccmr.py │ │ ├── corr.py │ │ ├── extractor.py │ │ ├── helpers.py │ │ ├── update.py │ │ ├── utils.py │ │ └── xcit.py │ ├── craft │ │ ├── LICENSE │ │ ├── README.md │ │ ├── __init__.py │ │ ├── corr.py │ │ ├── craft.py │ │ ├── extractor.py │ │ ├── gma.py │ │ ├── setrans.py │ │ ├── setrans_ablation.py │ │ ├── update.py │ │ └── utils.py │ ├── csflow │ │ ├── LICENSE │ │ ├── README.md │ │ ├── __init__.py │ │ └── csflow.py │ ├── dicl │ │ ├── README.md │ │ ├── __init__.py │ │ ├── dicl.py │ │ └── loss_functions.py │ ├── dip │ │ ├── LICENSE │ │ ├── README.md │ │ ├── __init__.py │ │ ├── corr.py │ │ ├── dip.py │ │ ├── extractor.py │ │ ├── path_match.py │ │ ├── update.py │ │ └── utils.py │ ├── dpflow │ │ ├── LICENSE │ │ ├── README.md │ │ ├── __init__.py │ │ ├── assets │ │ │ └── teaser.jpg │ │ ├── cgu.py │ │ ├── cgu_bidir_dual_encoder.py │ │ ├── configs │ │ │ ├── dpflow-test.yaml │ │ │ ├── dpflow-train1-chairs.yaml │ │ │ ├── dpflow-train2-things.yaml │ │ │ ├── dpflow-train3-sintel.yaml │ │ │ ├── dpflow-train4a-kitti.yaml │ │ │ ├── dpflow-train4b-spring.yaml │ │ │ └── dpflow-validate.yaml │ │ ├── conv.py │ │ ├── corr.py │ │ ├── dpflow.py │ │ ├── local_timm │ │ │ ├── activations.py │ │ │ ├── activations_me.py │ │ │ ├── cond_conv2d.py │ │ │ ├── config.py │ │ │ ├── conv2d_same.py │ │ │ ├── create_act.py │ │ │ ├── create_conv2d.py │ │ │ ├── drop.py │ │ │ ├── gelu.py │ │ │ ├── grn.py │ │ │ ├── helpers.py │ │ │ ├── layer_helpers.py │ │ │ ├── mixed_conv2d.py │ │ │ ├── mlp.py │ │ │ ├── norm.py │ │ │ ├── padding.py │ │ │ ├── pool2d_same.py │ │ │ └── weight_init.py │ │ ├── norm.py │ │ ├── pwc_modules.py │ │ ├── requirements.txt │ │ ├── res_stem.py │ │ ├── update.py │ │ └── utils.py │ ├── fastflownet │ │ ├── LICENSE │ │ ├── README.md │ │ ├── __init__.py │ │ └── fastflownet.py │ ├── flow1d │ │ ├── LICENSE │ │ ├── README.md │ │ ├── __init__.py │ │ ├── attention.py │ │ ├── correlation.py │ │ ├── extractor.py │ │ ├── flow1d.py │ │ ├── position.py │ │ ├── update.py │ │ └── utils.py │ ├── flowformer │ │ ├── LICENSE │ │ ├── README.md │ │ ├── __init__.py │ │ ├── attention.py │ │ ├── cnn.py │ │ ├── convnext.py │ │ ├── decoder.py │ │ ├── encoder.py │ │ ├── encoders.py │ │ ├── flowformer.py │ │ ├── gma.py │ │ ├── gru.py │ │ ├── mlpmixer.py │ │ ├── twins.py │ │ └── utils.py │ ├── flowformerplusplus │ │ ├── FlowFormer │ │ │ ├── PerCostFormer3 │ │ │ │ ├── NA.py │ │ │ │ ├── __init__.py │ │ │ │ ├── attention.py │ │ │ │ ├── cnn.py │ │ │ │ ├── crossattentionlayer.py │ │ │ │ ├── decoder.py │ │ │ │ ├── encoder.py │ │ │ │ ├── gma.py │ │ │ │ ├── gru.py │ │ │ │ ├── quater_upsampler.py │ │ │ │ ├── sk.py │ │ │ │ ├── transformer.py │ │ │ │ └── twins.py │ │ │ ├── __init__.py │ │ │ ├── common.py │ │ │ ├── encoders.py │ │ │ └── twins_ft.py │ │ ├── README.md │ │ ├── __init__.py │ │ ├── corr.py │ │ ├── extractor.py │ │ ├── flowformerplusplus.py │ │ ├── loss.py │ │ ├── optimizer │ │ │ └── __init__.py │ │ ├── position_encoding.py │ │ ├── timm0412 │ │ │ ├── data │ │ │ │ ├── __init__.py │ │ │ │ └── constants.py │ │ │ └── models │ │ │ │ ├── features.py │ │ │ │ ├── helpers.py │ │ │ │ ├── hub.py │ │ │ │ ├── layers │ │ │ │ ├── __init__.py │ │ │ │ ├── conv2d_same.py │ │ │ │ ├── drop.py │ │ │ │ ├── helpers.py │ │ │ │ ├── linear.py │ │ │ │ ├── mlp.py │ │ │ │ ├── padding.py │ │ │ │ ├── patch_embed.py │ │ │ │ └── weight_init.py │ │ │ │ ├── registry.py │ │ │ │ └── vision_transformer.py │ │ ├── unsup_loss.py │ │ ├── update.py │ │ └── utils.py │ ├── flownet │ │ ├── LICENSE │ │ ├── README.md │ │ ├── WEIGHTS_LICENSE │ │ ├── __init__.py │ │ ├── flownet2.py │ │ ├── flownet_base.py │ │ ├── flownet_fusion.py │ │ ├── flownetc.py │ │ ├── flownetcs.py │ │ ├── flownetcss.py │ │ ├── flownets.py │ │ ├── flownetsd.py │ │ ├── losses.py │ │ └── submodules.py │ ├── gma │ │ ├── LICENSE │ │ ├── README.md │ │ ├── __init__.py │ │ ├── corr.py │ │ ├── extractor.py │ │ ├── gma.py │ │ ├── gma_utils.py │ │ ├── update.py │ │ └── utils.py │ ├── gmflow │ │ ├── LICENSE │ │ ├── README.md │ │ ├── __init__.py │ │ ├── backbone.py │ │ ├── geometry.py │ │ ├── gmflow.py │ │ ├── matching.py │ │ ├── position.py │ │ ├── transformer.py │ │ ├── trident_conv.py │ │ └── utils.py │ ├── gmflownet │ │ ├── LICENSE │ │ ├── README.md │ │ ├── __init__.py │ │ ├── corr.py │ │ ├── extractor.py │ │ ├── gma.py │ │ ├── gmflownet.py │ │ ├── loss.py │ │ ├── swin_transformer.py │ │ ├── update.py │ │ └── utils │ │ │ ├── __init__.py │ │ │ ├── drop.py │ │ │ ├── helpers.py │ │ │ ├── utils.py │ │ │ └── weight_init.py │ ├── hd3 │ │ ├── LICENSE │ │ ├── README.md │ │ ├── __init__.py │ │ ├── decoder.py │ │ ├── dla.py │ │ ├── dla_up.py │ │ ├── hd3.py │ │ ├── hd3_ops.py │ │ ├── hd3losses.py │ │ └── vgg.py │ ├── irr │ │ ├── LICENSE │ │ ├── README.md │ │ ├── __init__.py │ │ ├── irr_modules.py │ │ ├── irr_pwc.py │ │ ├── losses.py │ │ ├── pwc_modules.py │ │ ├── pwcnet.py │ │ └── pwcnet_irr.py │ ├── lcv │ │ ├── LICENSE_LCV │ │ ├── LICENSE_RAFT │ │ ├── LICENSE_WEIGHTS │ │ ├── README.md │ │ ├── __init__.py │ │ ├── corr_lcv.py │ │ ├── extractor.py │ │ ├── lcv_raft.py │ │ ├── update.py │ │ └── utils.py │ ├── liteflownet │ │ ├── LICENSE_IRR │ │ ├── LICENSE_LFN1 │ │ ├── LICENSE_LFN2 │ │ ├── LICENSE_LFN3 │ │ ├── README.md │ │ ├── __init__.py │ │ ├── liteflownet.py │ │ ├── liteflownet2.py │ │ ├── liteflownet3.py │ │ └── warp.py │ ├── llaflow │ │ ├── LICENSE │ │ ├── README.md │ │ ├── __init__.py │ │ ├── aggregate.py │ │ ├── corr.py │ │ ├── extractor.py │ │ ├── gma.py │ │ ├── llaflow.py │ │ ├── update.py │ │ └── utils.py │ ├── maskflownet │ │ ├── LICENSE │ │ ├── README.md │ │ ├── __init__.py │ │ └── maskflownet.py │ ├── matchflow │ │ ├── LICENSE │ │ ├── QuadtreeAttention │ │ │ ├── QuadtreeAttention │ │ │ │ ├── functions │ │ │ │ │ └── quadtree_attention.py │ │ │ │ ├── modules │ │ │ │ │ └── quadtree_attention.py │ │ │ │ └── src │ │ │ │ │ ├── score_computation.cpp │ │ │ │ │ ├── score_computation.h │ │ │ │ │ ├── score_computation_kernal.cu │ │ │ │ │ ├── utils.h │ │ │ │ │ ├── value_aggregation.cpp │ │ │ │ │ ├── value_aggregation.h │ │ │ │ │ └── value_aggregation_kernel.cu │ │ │ └── setup.py │ │ ├── README.md │ │ ├── __init__.py │ │ ├── common.py │ │ ├── corr.py │ │ ├── extractor.py │ │ ├── gma.py │ │ ├── matchflow.py │ │ ├── matching_encoder.py │ │ ├── quadtree_attention.py │ │ ├── resnet_fpn.py │ │ ├── update.py │ │ └── utils.py │ ├── memflow │ │ ├── LICENSE │ │ ├── MemFlowNet │ │ │ ├── MemFlow.py │ │ │ ├── MemFlow_P.py │ │ │ ├── __init__.py │ │ │ ├── cnn.py │ │ │ ├── corr.py │ │ │ ├── gma.py │ │ │ ├── memory_util.py │ │ │ ├── sk.py │ │ │ ├── sk2.py │ │ │ └── update.py │ │ ├── README.md │ │ ├── __init__.py │ │ ├── encoders.py │ │ ├── kv_memory_store.py │ │ ├── memflow.py │ │ ├── memory_manager_skflow.py │ │ ├── optimizer.py │ │ └── utils.py │ ├── ms_raft_plus │ │ ├── LICENSE │ │ ├── README.md │ │ ├── __init__.py │ │ ├── corr.py │ │ ├── extractor.py │ │ ├── ms_raft_plus.py │ │ ├── update.py │ │ └── utils.py │ ├── neuflow │ │ ├── LICENSE │ │ ├── README.md │ │ ├── __init__.py │ │ ├── backbone.py │ │ ├── matching.py │ │ ├── neuflow.py │ │ ├── refine.py │ │ ├── transformer.py │ │ ├── upsample.py │ │ └── utils.py │ ├── neuflow2 │ │ ├── LICENSE │ │ ├── README.md │ │ ├── __init__.py │ │ ├── backbone_v7.py │ │ ├── corr.py │ │ ├── matching.py │ │ ├── neuflow2.py │ │ ├── refine.py │ │ ├── transformer.py │ │ ├── upsample.py │ │ └── utils.py │ ├── pwcnet │ │ ├── LICENSE │ │ ├── README.md │ │ ├── __init__.py │ │ └── pwcnet.py │ ├── raft │ │ ├── LICENSE │ │ ├── README.md │ │ ├── __init__.py │ │ ├── configs │ │ │ ├── raft-train1-chairs.yaml │ │ │ ├── raft-train2-things.yaml │ │ │ ├── raft-train3-sintel.yaml │ │ │ └── raft-train4-kitti.yaml │ │ ├── corr.py │ │ ├── extractor.py │ │ ├── raft.py │ │ ├── update.py │ │ └── utils.py │ ├── rapidflow │ │ ├── LICENSE │ │ ├── README.md │ │ ├── __init__.py │ │ ├── configs │ │ │ ├── rapidflow-test-kitti.yaml │ │ │ ├── rapidflow-test-sintel.yaml │ │ │ ├── rapidflow-train1-chairs.yaml │ │ │ ├── rapidflow-train2-things.yaml │ │ │ ├── rapidflow-train3-sintel.yaml │ │ │ ├── rapidflow-train4-kitti.yaml │ │ │ └── rapidflow-validate-things.yaml │ │ ├── convert_to_onnx.py │ │ ├── corr.py │ │ ├── image_samples │ │ │ ├── 000000_10.png │ │ │ ├── 000000_11.png │ │ │ └── README.md │ │ ├── local_timm │ │ │ ├── README.md │ │ │ ├── cond_conv2d.py │ │ │ ├── conv2d_same.py │ │ │ ├── create_conv2d.py │ │ │ ├── drop.py │ │ │ ├── gelu.py │ │ │ ├── helpers.py │ │ │ ├── layer_helpers.py │ │ │ ├── mixed_conv2d.py │ │ │ ├── mlp.py │ │ │ ├── norm.py │ │ │ ├── padding.py │ │ │ └── weight_init.py │ │ ├── next1d.py │ │ ├── next1d_encoder.py │ │ ├── onnx_infer.py │ │ ├── pwc_modules.py │ │ ├── rapidflow.py │ │ ├── requirements.txt │ │ ├── tensorrt_test.py │ │ ├── update.py │ │ └── utils.py │ ├── rpknet │ │ ├── LICENSE │ │ ├── README.md │ │ ├── __init__.py │ │ ├── configs │ │ │ ├── rpknet-test-kitti.yaml │ │ │ ├── rpknet-test-sintel.yaml │ │ │ ├── rpknet-test-spring.yaml │ │ │ ├── rpknet-train1-chairs.yaml │ │ │ ├── rpknet-train2-things.yaml │ │ │ ├── rpknet-train3-sintel.yaml │ │ │ ├── rpknet-train4-kitti.yaml │ │ │ └── rpknet-validate-things.yaml │ │ ├── corr.py │ │ ├── local_timm │ │ │ ├── README.md │ │ │ ├── drop.py │ │ │ ├── gelu.py │ │ │ ├── helpers.py │ │ │ ├── layer_helpers.py │ │ │ ├── norm.py │ │ │ ├── padding.py │ │ │ └── weight_init.py │ │ ├── pkconv.py │ │ ├── pkconv_slk.py │ │ ├── pkconv_slk_encoder.py │ │ ├── pklinear.py │ │ ├── pwc_modules.py │ │ ├── requirements.txt │ │ ├── rpknet.py │ │ ├── update_partial.py │ │ └── utils.py │ ├── scopeflow │ │ ├── LICENSE │ │ ├── README.md │ │ ├── __init__.py │ │ ├── irr_modules.py │ │ ├── irr_pwc_v2.py │ │ ├── losses.py │ │ └── pwc_modules.py │ ├── scv │ │ ├── LICENSE │ │ ├── README.md │ │ ├── __init__.py │ │ ├── compute_sparse_correlation.py │ │ ├── extractor.py │ │ ├── knn.py │ │ ├── scv.py │ │ ├── update.py │ │ └── utils.py │ ├── sea_raft │ │ ├── LICENSE │ │ ├── README.md │ │ ├── __init__.py │ │ ├── corr.py │ │ ├── extractor.py │ │ ├── layer.py │ │ ├── sea_raft.py │ │ ├── update.py │ │ └── utils.py │ ├── separableflow │ │ ├── LICENSE │ │ ├── README.md │ │ ├── __init__.py │ │ ├── compile.sh │ │ ├── corr.py │ │ ├── cost_agg.py │ │ ├── extractor.py │ │ ├── libs │ │ │ └── GANet │ │ │ │ ├── functions │ │ │ │ ├── GANet.py │ │ │ │ └── __init__.py │ │ │ │ ├── modules │ │ │ │ ├── GANet.py │ │ │ │ └── __init__.py │ │ │ │ ├── setup.py │ │ │ │ └── src │ │ │ │ ├── GANet_cuda.cpp │ │ │ │ ├── GANet_cuda.h │ │ │ │ ├── GANet_kernel.cu │ │ │ │ ├── GANet_kernel.h │ │ │ │ ├── GANet_kernel_share.cu │ │ │ │ ├── NLF_kernel.cu │ │ │ │ └── costvolume.cu │ │ ├── separableflow.py │ │ ├── update.py │ │ └── utils.py │ ├── skflow │ │ ├── LICENSE │ │ ├── README.md │ │ ├── __init__.py │ │ ├── corr.py │ │ ├── extractor.py │ │ ├── gma.py │ │ ├── skflow.py │ │ ├── update.py │ │ └── utils.py │ ├── splatflow │ │ ├── LICENSE │ │ ├── README.md │ │ ├── __init__.py │ │ ├── attention.py │ │ ├── corr.py │ │ ├── extractor.py │ │ ├── softsplat.py │ │ ├── splatflow.py │ │ ├── update.py │ │ └── util.py │ ├── starflow │ │ ├── LICENSE │ │ ├── README.md │ │ ├── __init__.py │ │ ├── irr_modules.py │ │ ├── pwc_modules.py │ │ └── starflow.py │ ├── streamflow │ │ ├── LICENSE │ │ ├── README.md │ │ ├── __init__.py │ │ ├── corr.py │ │ ├── datasets.py │ │ ├── extractor.py │ │ ├── gma.py │ │ ├── streamflow.py │ │ ├── twins_csc.py │ │ ├── update.py │ │ └── utils.py │ ├── unimatch │ │ ├── LICENSE │ │ ├── README.md │ │ ├── __init__.py │ │ ├── attention.py │ │ ├── backbone.py │ │ ├── geometry.py │ │ ├── matching.py │ │ ├── position.py │ │ ├── reg_refine.py │ │ ├── transformer.py │ │ ├── trident_conv.py │ │ ├── unimatch.py │ │ └── utils.py │ ├── vcn │ │ ├── LICENSE │ │ ├── README.md │ │ ├── __init__.py │ │ ├── conv4d.py │ │ ├── submodule.py │ │ └── vcn.py │ └── videoflow │ │ ├── Networks │ │ ├── BOFNet │ │ │ ├── __init__.py │ │ │ ├── cnn.py │ │ │ ├── corr.py │ │ │ ├── gma.py │ │ │ ├── network.py │ │ │ ├── sk.py │ │ │ ├── sk2.py │ │ │ └── update.py │ │ ├── MOFNetStack │ │ │ ├── __init__.py │ │ │ ├── cnn.py │ │ │ ├── corr.py │ │ │ ├── gma.py │ │ │ ├── network.py │ │ │ ├── resstack.py │ │ │ ├── sk.py │ │ │ ├── sk2.py │ │ │ ├── stack.py │ │ │ ├── stackcat.py │ │ │ └── update.py │ │ ├── __init__.py │ │ ├── common.py │ │ ├── encoders.py │ │ └── twins_ft.py │ │ ├── README.md │ │ ├── __init__.py │ │ ├── utils.py │ │ ├── videoflow_bof.py │ │ └── videoflow_mof.py └── utils │ ├── __init__.py │ ├── callbacks │ ├── RobotoMono-Regular.ttf │ ├── __init__.py │ └── logger.py │ ├── correlation.py │ ├── dummy_datasets.py │ ├── external │ ├── __init__.py │ ├── alt_cuda_corr │ │ ├── LICENSE │ │ ├── README.md │ │ ├── correlation.cpp │ │ ├── correlation_kernel.cu │ │ └── setup.py │ ├── flow_IO.py │ ├── flowpy.py │ ├── raft.py │ └── selflow.py │ ├── flow_metrics.py │ ├── flow_utils.py │ ├── flowpy_torch.py │ ├── io_adapter.py │ ├── lightning │ ├── ptlflow_checkpoint_connector.py │ ├── ptlflow_cli.py │ └── ptlflow_trainer.py │ ├── registry.py │ ├── timer.py │ └── utils.py ├── pyproject.toml ├── requirements.txt ├── summary_metrics.py ├── test.py ├── tests ├── ptlflow │ ├── data │ │ └── test_datasets.py │ ├── models │ │ ├── base_model │ │ │ └── test_get_dataset.py │ │ ├── test_checkpoints.py │ │ └── test_models.py │ ├── test_lightning.py │ └── utils │ │ ├── external │ │ └── test_flowpy.py │ │ ├── test_correlation.py │ │ ├── test_flow_utils.py │ │ └── test_flowpy_torch.py ├── test_infer.py ├── test_model_benchmark.py ├── test_summary_metrics.py ├── test_test.py └── test_validate.py ├── train.py └── validate.py /.github/workflows/build.yml: -------------------------------------------------------------------------------- 1 | # This workflow will install Python dependencies, run tests with a variety of Python versions 2 | # For more information see: https://help.github.com/actions/language-and-framework-guides/using-python-with-github-actions 3 | 4 | # Build the package and install with pip, then run pytest on the pip package. 5 | 6 | name: build 7 | 8 | on: 9 | push: 10 | branches: [ main ] 11 | pull_request: 12 | branches: [ main ] 13 | 14 | jobs: 15 | build: 16 | 17 | runs-on: ubuntu-latest 18 | strategy: 19 | fail-fast: false 20 | matrix: 21 | python-version: ['3.12'] 22 | 23 | steps: 24 | - uses: actions/checkout@v4 25 | - name: Set up Python ${{ matrix.python-version }} 26 | uses: actions/setup-python@v4 27 | with: 28 | python-version: ${{ matrix.python-version }} 29 | - name: Install dependencies 30 | run: | 31 | python -m pip install --upgrade pip 32 | python -m pip install build==1.2.2.post1 33 | python -m pip install --upgrade setuptools==75.6.0 wheel==0.45.1 34 | python -m pip install --upgrade pytest==8.3.3 35 | pip install torch==2.5.1 torchvision==0.20.1 --index-url https://download.pytorch.org/whl/cpu 36 | - name: Install package and remove local dir 37 | run: | 38 | python -m build 39 | python -m pip install dist/*.whl 40 | mv ptlflow ptlflow_tmp 41 | - name: Test with pytest 42 | run: | 43 | python -m pytest tests/ 44 | -------------------------------------------------------------------------------- /.github/workflows/lightning.yml: -------------------------------------------------------------------------------- 1 | # This workflow will install Python dependencies, run tests with a variety of Python versions 2 | # For more information see: https://help.github.com/actions/language-and-framework-guides/using-python-with-github-actions 3 | 4 | name: lightning 5 | 6 | on: 7 | push: 8 | branches: [ main ] 9 | pull_request: 10 | branches: [ main ] 11 | 12 | jobs: 13 | build: 14 | 15 | runs-on: ubuntu-latest 16 | strategy: 17 | fail-fast: false 18 | matrix: 19 | lightning: ["2.1.4", "2.2.5", "2.3.3", "2.4.0", "2.5.1"] 20 | 21 | steps: 22 | - uses: actions/checkout@v4 23 | - name: Replace lightning 24 | uses: jacobtomlinson/gha-find-replace@v3 25 | with: 26 | find: "lightning[pytorch-extra]>=2,<2.5" 27 | replace: "lightning[pytorch-extra]==${{ matrix.lightning }}" 28 | regex: false 29 | include: "requirements.txt" 30 | - name: Install dependencies 31 | run: | 32 | python -m pip install --upgrade pip 33 | pip install torch==2.5.1 torchvision==0.20.1 --index-url https://download.pytorch.org/whl/cpu 34 | pip install -r requirements.txt 35 | - name: Test with pytest 36 | run: | 37 | pip install pytest==8.3.3 38 | python -m pytest tests/ 39 | -------------------------------------------------------------------------------- /.github/workflows/publish_pypi.yml: -------------------------------------------------------------------------------- 1 | # This workflow will upload a Python Package using Twine when a release is created 2 | # For more information see: https://help.github.com/en/actions/language-and-framework-guides/using-python-with-github-actions#publishing-to-package-registries 3 | 4 | # This workflow uses actions that are not certified by GitHub. 5 | # They are provided by a third-party and are governed by 6 | # separate terms of service, privacy policy, and support 7 | # documentation. 8 | 9 | name: Upload Python Package 10 | 11 | on: 12 | release: 13 | types: [published] 14 | 15 | jobs: 16 | deploy: 17 | 18 | runs-on: ubuntu-latest 19 | 20 | steps: 21 | - uses: actions/checkout@v4 22 | - name: Set up Python 23 | uses: actions/setup-python@v4 24 | with: 25 | python-version: '3.10' 26 | - name: Install dependencies 27 | run: | 28 | python -m pip install --upgrade pip 29 | pip install torch==2.5.1 torchvision==0.20.1 --index-url https://download.pytorch.org/whl/cpu 30 | pip install build==1.2.2.post1 31 | - name: Build package 32 | run: python -m build 33 | - name: Publish package 34 | uses: pypa/gh-action-pypi-publish@release/v1 35 | with: 36 | user: __token__ 37 | password: ${{ secrets.PYPI_API_TOKEN }} -------------------------------------------------------------------------------- /.github/workflows/python.yml: -------------------------------------------------------------------------------- 1 | # This workflow will install Python dependencies, run tests with a variety of Python versions 2 | # For more information see: https://help.github.com/actions/language-and-framework-guides/using-python-with-github-actions 3 | 4 | name: python 5 | 6 | on: 7 | push: 8 | branches: [ main ] 9 | pull_request: 10 | branches: [ main ] 11 | 12 | jobs: 13 | build: 14 | 15 | runs-on: ubuntu-latest 16 | strategy: 17 | fail-fast: false 18 | matrix: 19 | python-version: ['3.10', '3.11', '3.12'] 20 | 21 | steps: 22 | - uses: actions/checkout@v4 23 | - name: Set up Python ${{ matrix.python-version }} 24 | uses: actions/setup-python@v4 25 | with: 26 | python-version: ${{ matrix.python-version }} 27 | - name: Install dependencies 28 | run: | 29 | python -m pip install --upgrade pip 30 | pip install torch==2.5.1 torchvision==0.20.1 --index-url https://download.pytorch.org/whl/cpu 31 | pip install -r requirements.txt 32 | - name: Test with pytest 33 | run: | 34 | pip install pytest==8.3.3 35 | python -m pytest tests/ 36 | -------------------------------------------------------------------------------- /.github/workflows/pytorch.yml: -------------------------------------------------------------------------------- 1 | # This workflow will install Python dependencies, run tests with a variety of Python versions 2 | # For more information see: https://help.github.com/actions/language-and-framework-guides/using-python-with-github-actions 3 | 4 | name: pytorch 5 | 6 | on: 7 | push: 8 | branches: [ main ] 9 | pull_request: 10 | branches: [ main ] 11 | 12 | jobs: 13 | build: 14 | 15 | runs-on: ubuntu-latest 16 | strategy: 17 | fail-fast: false 18 | matrix: 19 | pytorch: [ 20 | 'torch==2.6.0 torchvision==0.21.0 --index-url https://download.pytorch.org/whl/cpu', 21 | 'torch==2.5.1 torchvision==0.20.1 --index-url https://download.pytorch.org/whl/cpu', 22 | 'torch==2.4.1 torchvision==0.19.1 --index-url https://download.pytorch.org/whl/cpu', 23 | 'torch==2.3.1 torchvision==0.18.1 --index-url https://download.pytorch.org/whl/cpu', 24 | ] 25 | 26 | steps: 27 | - uses: actions/checkout@v4 28 | - name: Install dependencies 29 | run: | 30 | python -m pip install --upgrade pip 31 | pip install ${{ matrix.pytorch }} 32 | pip install -r requirements.txt 33 | - name: Test with pytest 34 | run: | 35 | pip install pytest==8.3.3 36 | python -m pytest tests/ 37 | -------------------------------------------------------------------------------- /.readthedocs.yaml: -------------------------------------------------------------------------------- 1 | version: "2" 2 | 3 | build: 4 | os: "ubuntu-22.04" 5 | tools: 6 | python: "3.10" 7 | 8 | sphinx: 9 | configuration: docs/source/conf.py 10 | 11 | python: 12 | install: 13 | - requirements: docs/requirements.txt -------------------------------------------------------------------------------- /MANIFEST.in: -------------------------------------------------------------------------------- 1 | global-include *.txt *.md *LICENSE *.ttf -------------------------------------------------------------------------------- /configs/results/model_benchmark_all.yaml: -------------------------------------------------------------------------------- 1 | # lightning.pytorch==2.4.0 2 | # Use this config to benchmark all the models. 3 | # python validate.py --config configs/results/model_benchmark_all.yaml 4 | all: true 5 | select: null 6 | ckpt_path: null 7 | exclude: null 8 | csv_path: null 9 | num_trials: 1 10 | num_samples: 10 11 | sleep_interval: 0.0 12 | input_size: 13 | - 500 14 | - 1000 15 | output_path: outputs/benchmark 16 | final_speed_mode: median 17 | final_memory_mode: first 18 | plot_axes: null 19 | plot_log_x: false 20 | plot_log_y: false 21 | datatypes: 22 | - fp32 23 | batch_size: 1 24 | seed_everything: true 25 | -------------------------------------------------------------------------------- /configs/results/validate_all.yaml: -------------------------------------------------------------------------------- 1 | # lightning.pytorch==2.4.0 2 | # Use this config to generate validation results for all models using all their pretrained ckpts. 3 | # python validate.py --config configs/results/validate_all.yaml 4 | all: true 5 | select: null 6 | exclude: null 7 | ckpt_path: null 8 | output_path: outputs/validate 9 | write_outputs: false 10 | show: false 11 | flow_format: original 12 | max_forward_side: null 13 | scale_factor: null 14 | max_show_side: 1000 15 | max_samples: null 16 | reversed: false 17 | fp16: false 18 | seq_val_mode: all 19 | write_individual_metrics: false 20 | epe_clip: 5.0 21 | metric_exclude: 22 | - wauc 23 | seed_everything: true 24 | data: 25 | predict_dataset: null 26 | test_dataset: null 27 | train_dataset: null 28 | val_dataset: sintel-clean-occ+sintel-final-occ+kitti-2012+kitti-2015 29 | train_batch_size: null 30 | train_num_workers: 4 31 | train_crop_size: null 32 | train_transform_cuda: false 33 | train_transform_fp16: false 34 | autoflow_root_dir: null 35 | flying_chairs_root_dir: null 36 | flying_chairs2_root_dir: null 37 | flying_things3d_root_dir: null 38 | flying_things3d_subset_root_dir: null 39 | mpi_sintel_root_dir: null 40 | kitti_2012_root_dir: null 41 | kitti_2015_root_dir: null 42 | hd1k_root_dir: null 43 | tartanair_root_dir: null 44 | spring_root_dir: null 45 | kubric_root_dir: null 46 | dataset_config_path: ./datasets.yaml 47 | -------------------------------------------------------------------------------- /configs/results/validate_all_things.yaml: -------------------------------------------------------------------------------- 1 | # lightning.pytorch==2.4.0 2 | # Use this config to generate validation results for all models using their "things" pretrained ckpt. 3 | # python validate.py --config configs/results/validate_all_things.yaml 4 | all: true 5 | select: null 6 | exclude: null 7 | ckpt_path: things 8 | output_path: outputs/validate 9 | write_outputs: false 10 | show: false 11 | flow_format: original 12 | max_forward_side: null 13 | scale_factor: null 14 | max_show_side: 1000 15 | max_samples: null 16 | reversed: false 17 | fp16: false 18 | seq_val_mode: all 19 | write_individual_metrics: false 20 | epe_clip: 5.0 21 | metric_exclude: null 22 | seed_everything: true 23 | data: 24 | predict_dataset: null 25 | test_dataset: null 26 | train_dataset: null 27 | val_dataset: sintel-clean+sintel-final+kitti-2015 28 | train_batch_size: null 29 | train_num_workers: 4 30 | train_crop_size: null 31 | train_transform_cuda: false 32 | train_transform_fp16: false 33 | autoflow_root_dir: null 34 | flying_chairs_root_dir: null 35 | flying_chairs2_root_dir: null 36 | flying_things3d_root_dir: null 37 | flying_things3d_subset_root_dir: null 38 | mpi_sintel_root_dir: null 39 | kitti_2012_root_dir: null 40 | kitti_2015_root_dir: null 41 | hd1k_root_dir: null 42 | tartanair_root_dir: null 43 | spring_root_dir: null 44 | kubric_root_dir: null 45 | dataset_config_path: ./datasets.yaml 46 | -------------------------------------------------------------------------------- /datasets.yaml: -------------------------------------------------------------------------------- 1 | autoflow: /path/to/autoflow 2 | flying_chairs: /path/to/FlyingChairs_release 3 | flying_chairs2: /path/to/FlyingChairs2 4 | flying_things3d: /path/to/FlyingThings3D 5 | flying_things3d_subset: /path/to/FlyingThings3D_subset 6 | mpi_sintel: /path/to/MPI-Sintel 7 | kitti_2012: /path/to/KITTI/2012 8 | kitti_2015: /path/to/KITTI/2015 9 | hd1k: /path/to/HD1K 10 | tartanair: /path/to/tartanair 11 | spring: /path/to/spring 12 | kubric: /path/to/kubric 13 | -------------------------------------------------------------------------------- /docs/Makefile: -------------------------------------------------------------------------------- 1 | # Minimal makefile for Sphinx documentation 2 | # 3 | 4 | # You can set these variables from the command line, and also 5 | # from the environment for the first two. 6 | SPHINXOPTS ?= 7 | SPHINXBUILD ?= sphinx-build 8 | SOURCEDIR = source 9 | BUILDDIR = build 10 | 11 | # Put it first so that "make" without argument is like "make help". 12 | help: 13 | @$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) 14 | 15 | .PHONY: help Makefile 16 | 17 | # Catch-all target: route all unknown targets to Sphinx using the new 18 | # "make mode" option. $(O) is meant as a shortcut for $(SPHINXOPTS). 19 | %: Makefile 20 | @$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) 21 | -------------------------------------------------------------------------------- /docs/make.bat: -------------------------------------------------------------------------------- 1 | @ECHO OFF 2 | 3 | pushd %~dp0 4 | 5 | REM Command file for Sphinx documentation 6 | 7 | if "%SPHINXBUILD%" == "" ( 8 | set SPHINXBUILD=sphinx-build 9 | ) 10 | set SOURCEDIR=source 11 | set BUILDDIR=build 12 | 13 | if "%1" == "" goto help 14 | 15 | %SPHINXBUILD% >NUL 2>NUL 16 | if errorlevel 9009 ( 17 | echo. 18 | echo.The 'sphinx-build' command was not found. Make sure you have Sphinx 19 | echo.installed, then set the SPHINXBUILD environment variable to point 20 | echo.to the full path of the 'sphinx-build' executable. Alternatively you 21 | echo.may add the Sphinx directory to PATH. 22 | echo. 23 | echo.If you don't have Sphinx installed, grab it from 24 | echo.http://sphinx-doc.org/ 25 | exit /b 1 26 | ) 27 | 28 | %SPHINXBUILD% -M %1 %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O% 29 | goto end 30 | 31 | :help 32 | %SPHINXBUILD% -M help %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O% 33 | 34 | :end 35 | popd 36 | -------------------------------------------------------------------------------- /docs/requirements.txt: -------------------------------------------------------------------------------- 1 | --extra-index-url https://download.pytorch.org/whl/cpu 2 | 3 | numpydoc==1.8.0 4 | sphinx==8.1.3 5 | sphinx_rtd_theme==3.0.2 6 | 7 | torch==2.5.1 8 | torchvision==0.20.1 9 | 10 | einops<0.9 11 | h5py<3.13 12 | kaleido<0.3 13 | lightning[pytorch-extra]>=2.1,<2.5 14 | loguru<0.8 15 | opencv-python<4.11 16 | pandas<2.3 17 | plotly<5.25 18 | pypng==0.20220715.0 19 | scipy<1.15 20 | tabulate<0.10 21 | tensorboard<2.19 22 | timm<1.1 -------------------------------------------------------------------------------- /docs/source/api/ptlflow.data.rst: -------------------------------------------------------------------------------- 1 | =========== 2 | ptflow.data 3 | =========== 4 | 5 | .. toctree:: 6 | :maxdepth: 1 7 | 8 | ptlflow/data/datasets 9 | ptlflow/data/flow_transforms -------------------------------------------------------------------------------- /docs/source/api/ptlflow.models.rst: -------------------------------------------------------------------------------- 1 | ============== 2 | ptlflow.models 3 | ============== 4 | 5 | .. toctree:: 6 | :maxdepth: 1 7 | 8 | ptlflow/models/base_model -------------------------------------------------------------------------------- /docs/source/api/ptlflow.utils.rst: -------------------------------------------------------------------------------- 1 | ============= 2 | ptlflow.utils 3 | ============= 4 | 5 | .. toctree:: 6 | :maxdepth: 1 7 | 8 | ptlflow/utils/correlation 9 | ptlflow/utils/dummy_datasets 10 | ptlflow/utils/flow_metrics 11 | ptlflow/utils/flow_utils 12 | ptlflow/utils/flowpy_torch 13 | ptlflow/utils/io_adapter 14 | ptlflow/utils/timer 15 | ptlflow/utils/utils 16 | 17 | .. toctree:: 18 | :maxdepth: 1 19 | :caption: Callbacks 20 | 21 | ptlflow/utils/callbacks/logger -------------------------------------------------------------------------------- /docs/source/api/ptlflow/data/datasets.rst: -------------------------------------------------------------------------------- 1 | .. _datasets: 2 | 3 | =========== 4 | datasets.py 5 | =========== 6 | 7 | .. automodule:: ptlflow.data.datasets 8 | :members: 9 | :special-members: __init__ -------------------------------------------------------------------------------- /docs/source/api/ptlflow/data/flow_transforms.rst: -------------------------------------------------------------------------------- 1 | ================== 2 | flow_transforms.py 3 | ================== 4 | 5 | .. automodule:: ptlflow.data.flow_transforms 6 | :members: 7 | :special-members: __init__ -------------------------------------------------------------------------------- /docs/source/api/ptlflow/models/base_model.rst: -------------------------------------------------------------------------------- 1 | .. _base-model: 2 | 3 | ========= 4 | BaseModel 5 | ========= 6 | 7 | .. automodule:: ptlflow.models.base_model.base_model 8 | :members: 9 | :special-members: __init__ -------------------------------------------------------------------------------- /docs/source/api/ptlflow/utils/callbacks/logger.rst: -------------------------------------------------------------------------------- 1 | ========= 2 | logger.py 3 | ========= 4 | 5 | .. automodule:: ptlflow.utils.callbacks.logger 6 | :members: 7 | :special-members: __init__ -------------------------------------------------------------------------------- /docs/source/api/ptlflow/utils/correlation.rst: -------------------------------------------------------------------------------- 1 | ============== 2 | correlation.py 3 | ============== 4 | 5 | .. automodule:: ptlflow.utils.correlation 6 | :members: 7 | :special-members: __init__ -------------------------------------------------------------------------------- /docs/source/api/ptlflow/utils/dummy_datasets.rst: -------------------------------------------------------------------------------- 1 | ================= 2 | dummy_datasets.py 3 | ================= 4 | 5 | .. automodule:: ptlflow.utils.dummy_datasets 6 | :members: -------------------------------------------------------------------------------- /docs/source/api/ptlflow/utils/flow_metrics.rst: -------------------------------------------------------------------------------- 1 | =============== 2 | flow_metrics.py 3 | =============== 4 | 5 | .. automodule:: ptlflow.utils.flow_metrics 6 | :members: 7 | :special-members: __init__ -------------------------------------------------------------------------------- /docs/source/api/ptlflow/utils/flow_utils.rst: -------------------------------------------------------------------------------- 1 | ============= 2 | flow_utils.py 3 | ============= 4 | 5 | .. automodule:: ptlflow.utils.flow_utils 6 | :members: -------------------------------------------------------------------------------- /docs/source/api/ptlflow/utils/flowpy_torch.rst: -------------------------------------------------------------------------------- 1 | =============== 2 | flowpy_torch.py 3 | =============== 4 | 5 | .. automodule:: ptlflow.utils.flowpy_torch 6 | :members: -------------------------------------------------------------------------------- /docs/source/api/ptlflow/utils/io_adapter.rst: -------------------------------------------------------------------------------- 1 | ============= 2 | io_adapter.py 3 | ============= 4 | 5 | .. automodule:: ptlflow.utils.io_adapter 6 | :members: 7 | :special-members: __init__ -------------------------------------------------------------------------------- /docs/source/api/ptlflow/utils/timer.rst: -------------------------------------------------------------------------------- 1 | ======== 2 | timer.py 3 | ======== 4 | 5 | .. automodule:: ptlflow.utils.timer 6 | :members: 7 | :special-members: __init__ -------------------------------------------------------------------------------- /docs/source/api/ptlflow/utils/utils.rst: -------------------------------------------------------------------------------- 1 | ======== 2 | utils.py 3 | ======== 4 | 5 | .. automodule:: ptlflow.utils.utils 6 | :members: -------------------------------------------------------------------------------- /docs/source/api/scripts.rst: -------------------------------------------------------------------------------- 1 | ======= 2 | scripts 3 | ======= 4 | 5 | .. toctree:: 6 | :maxdepth: 1 7 | :caption: Prediction 8 | 9 | scripts/infer 10 | scripts/train 11 | scripts/validate 12 | 13 | .. toctree:: 14 | :maxdepth: 1 15 | :caption: Utils 16 | 17 | scripts/model_benchmark 18 | scripts/summary_metrics -------------------------------------------------------------------------------- /docs/source/api/scripts/infer.rst: -------------------------------------------------------------------------------- 1 | ======== 2 | infer.py 3 | ======== 4 | 5 | .. automodule:: infer 6 | :members: -------------------------------------------------------------------------------- /docs/source/api/scripts/model_benchmark.rst: -------------------------------------------------------------------------------- 1 | ================== 2 | model_benchmark.py 3 | ================== 4 | 5 | .. automodule:: model_benchmark 6 | :members: -------------------------------------------------------------------------------- /docs/source/api/scripts/summary_metrics.rst: -------------------------------------------------------------------------------- 1 | ================== 2 | summary_metrics.py 3 | ================== 4 | 5 | .. automodule:: summary_metrics 6 | :members: -------------------------------------------------------------------------------- /docs/source/api/scripts/train.rst: -------------------------------------------------------------------------------- 1 | ======== 2 | train.py 3 | ======== 4 | 5 | .. automodule:: train 6 | :members: -------------------------------------------------------------------------------- /docs/source/api/scripts/validate.rst: -------------------------------------------------------------------------------- 1 | =========== 2 | validate.py 3 | =========== 4 | 5 | .. automodule:: validate 6 | :members: -------------------------------------------------------------------------------- /docs/source/models/checkpoint_list.rst: -------------------------------------------------------------------------------- 1 | =============== 2 | Checkpoint List 3 | =============== 4 | 5 | The table below lists the available checkpoints and show what are their original counterparts. 6 | 7 | .. csv-table:: 8 | :file: checkpoint_list.csv 9 | :widths: auto 10 | :header-rows: 1 11 | -------------------------------------------------------------------------------- /docs/source/results/accuracy_epe.rst: -------------------------------------------------------------------------------- 1 | ============== 2 | Accuracy - EPE 3 | ============== 4 | 5 | .. csv-table:: 6 | :file: summarized_metrics-epe.csv 7 | :widths: auto 8 | :header-rows: 1 9 | -------------------------------------------------------------------------------- /docs/source/results/accuracy_epe_flall.rst: -------------------------------------------------------------------------------- 1 | ========================= 2 | Accuracy - EPE and Fl-All 3 | ========================= 4 | 5 | .. csv-table:: 6 | :file: summarized_metrics-epe_flall.csv 7 | :widths: auto 8 | :header-rows: 1 9 | -------------------------------------------------------------------------------- /docs/source/results/model_benchmark.rst: -------------------------------------------------------------------------------- 1 | =============== 2 | Model benchmark 3 | =============== 4 | 5 | Test environment 6 | ================ 7 | 8 | - PyTorch 2.2.2 9 | - Lightning 1.9.5 10 | - NVIDIA RTX 3090 11 | - CUDA 12.1 12 | 13 | .. csv-table:: 14 | :file: model_benchmark-all.csv 15 | :widths: auto 16 | :header-rows: 1 -------------------------------------------------------------------------------- /docs/source/results/paper_ptlflow.rst: -------------------------------------------------------------------------------- 1 | .. _paper-ptlflow-metrics: 2 | 3 | ===================== 4 | Paper/PTLFlow metrics 5 | ===================== 6 | 7 | The table below show a comparison between the validation results of PTLFlow and the values in the original papers. 8 | The PTLFlow results correspond to models using the ``things`` pretrained checkpoints (often called C+T in the papers). 9 | 10 | .. csv-table:: 11 | :file: paper_ptlflow_metrics.csv 12 | :widths: auto 13 | :header-rows: 1 14 | -------------------------------------------------------------------------------- /docs/source/results/result_plots.rst: -------------------------------------------------------------------------------- 1 | ============ 2 | Result plots 3 | ============ 4 | 5 | Scatter plots to visualize some metrics more easily. 6 | 7 | Inference time x Sintel Final EPE 8 | ================================= 9 | 10 | .. raw:: html 11 | 12 | 13 | 14 | Memory usage x Sintel Final EPE 15 | =============================== 16 | 17 | .. raw:: html 18 | 19 | 20 | 21 | Sintel Final EPE x KITTI 2015 Fl-All 22 | ==================================== 23 | 24 | .. raw:: html 25 | 26 | 27 | 28 | Model parameters x FLOPs 29 | ======================== 30 | 31 | .. raw:: html 32 | 33 | -------------------------------------------------------------------------------- /docs/source/starting/config_files.rst: -------------------------------------------------------------------------------- 1 | .. _using-config-files: 2 | 3 | ================== 4 | Using config files 5 | ================== 6 | 7 | One of the nicest additions in Lightning 2.0 is the use of YAML files to save and load configurations more easily. 8 | You can create a config file by using the argument ``--print_config`` and then redirecting the output to a YAML file. 9 | For example, you can save the config of a training: 10 | 11 | .. code-block:: bash 12 | 13 | python train.py --model raft --data.train_dataset chairs2 --data.val_dataset kitti-2015-val --print_config > raft_train_config.yaml 14 | 15 | This also provides an easier way to know which arguments are available and to tune them according to your needs. 16 | This config can then be used to run the training by loading it with: 17 | 18 | .. code-block:: bash 19 | 20 | python train.py --config raft_train_config.yaml 21 | 22 | All the basic scripts (train, test, validate, infer, model_benchmark) should be able to save and load config files in this way. 23 | 24 | If you want to learn more about config files, check the `LightningCLI documentation `__, especially the content about YAML files. -------------------------------------------------------------------------------- /misc/create_compare_paper.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | python ../compare_paper_results.py --paper_results_path ../docs/source/results/paper_results_things.csv --validate_results_path ../docs/source/results/metrics_all_things.csv --add_delta -------------------------------------------------------------------------------- /misc/create_metric_summaries.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | python ../summary_metrics.py --metrics_path ../docs/source/results/metrics_all.csv --chosen_metrics epe 3 | 4 | python ../summary_metrics.py --metrics_path ../docs/source/results/metrics_all.csv --chosen_metrics epe outlier -------------------------------------------------------------------------------- /misc/create_plots.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | python ../plot_results.py \ 3 | --metrics_csv_path ../docs/source/results/metrics_all.csv \ 4 | --benchmark_csv_path ../docs/source/results/model_benchmark-all.csv \ 5 | --checkpoint_names things \ 6 | --plot_axes "memory(gb)-fp32" "sintel-final-occ-val/epe" 7 | 8 | python ../plot_results.py \ 9 | --metrics_csv_path ../docs/source/results/metrics_all.csv \ 10 | --benchmark_csv_path ../docs/source/results/model_benchmark-all.csv \ 11 | --checkpoint_names things \ 12 | --plot_axes "memory(gb)-fp32" "sintel-final-occ-val/flall" 13 | 14 | python ../plot_results.py \ 15 | --metrics_csv_path ../docs/source/results/metrics_all.csv \ 16 | --benchmark_csv_path ../docs/source/results/model_benchmark-all.csv \ 17 | --checkpoint_names things \ 18 | --plot_axes "time(ms)-fp32" "sintel-final-occ-val/epe" 19 | 20 | python ../plot_results.py \ 21 | --metrics_csv_path ../docs/source/results/metrics_all.csv \ 22 | --benchmark_csv_path ../docs/source/results/model_benchmark-all.csv \ 23 | --checkpoint_names things \ 24 | --plot_axes "time(ms)-fp32" "sintel-final-occ-val/flall" 25 | 26 | python ../plot_results.py \ 27 | --metrics_csv_path ../docs/source/results/metrics_all.csv \ 28 | --benchmark_csv_path ../docs/source/results/model_benchmark-all.csv \ 29 | --checkpoint_names things \ 30 | --plot_axes "params" "flops" \ 31 | --log_x \ 32 | --log_y 33 | 34 | python ../plot_results.py \ 35 | --metrics_csv_path ../docs/source/results/metrics_all.csv \ 36 | --benchmark_csv_path ../docs/source/results/model_benchmark-all.csv \ 37 | --checkpoint_names things \ 38 | --plot_axes "sintel-final-occ-val/epe" "kitti-2015-val/flall" -------------------------------------------------------------------------------- /ptlflow/data/Kitti2012_val.txt: -------------------------------------------------------------------------------- 1 | 000000_10 2 | 000012_10 3 | 000015_10 4 | 000016_10 5 | 000017_10 6 | 000018_10 7 | 000024_10 8 | 000030_10 9 | 000038_10 10 | 000039_10 11 | 000042_10 12 | 000050_10 13 | 000054_10 14 | 000059_10 15 | 000060_10 16 | 000061_10 17 | 000077_10 18 | 000078_10 19 | 000081_10 20 | 000089_10 21 | 000097_10 22 | 000101_10 23 | 000107_10 24 | 000121_10 25 | 000124_10 26 | 000142_10 27 | 000145_10 28 | 000146_10 29 | 000152_10 30 | 000154_10 31 | 000155_10 32 | 000158_10 33 | 000159_10 34 | 000160_10 35 | 000164_10 36 | 000182_10 37 | 000183_10 38 | 000184_10 39 | 000190_10 -------------------------------------------------------------------------------- /ptlflow/data/Kitti2015_val.txt: -------------------------------------------------------------------------------- 1 | 000010_10 2 | 000011_10 3 | 000012_10 4 | 000025_10 5 | 000026_10 6 | 000030_10 7 | 000031_10 8 | 000040_10 9 | 000041_10 10 | 000042_10 11 | 000046_10 12 | 000052_10 13 | 000053_10 14 | 000072_10 15 | 000073_10 16 | 000074_10 17 | 000075_10 18 | 000076_10 19 | 000080_10 20 | 000081_10 21 | 000085_10 22 | 000086_10 23 | 000095_10 24 | 000096_10 25 | 000097_10 26 | 000098_10 27 | 000104_10 28 | 000116_10 29 | 000117_10 30 | 000120_10 31 | 000121_10 32 | 000126_10 33 | 000127_10 34 | 000153_10 35 | 000172_10 36 | 000175_10 37 | 000183_10 38 | 000184_10 39 | 000190_10 40 | 000199_10 -------------------------------------------------------------------------------- /ptlflow/data/Sintel_val.txt: -------------------------------------------------------------------------------- 1 | ambush_2 2 | bamboo_2 3 | cave_2 4 | market_2 5 | shaman_2 6 | temple_2 -------------------------------------------------------------------------------- /ptlflow/data/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hmorimitsu/ptlflow/fb6d753c3b16e74481258891b45668dcba63e54e/ptlflow/data/__init__.py -------------------------------------------------------------------------------- /ptlflow/models/README.md: -------------------------------------------------------------------------------- 1 | # Models 2 | 3 | Except for `base_model`, each folder contains a model, or a group of variants of a model. `base_model` contains an abstract class that should be inherited by the other concrete models. 4 | 5 | Most models are adapted from other sources and each of them are subjected to different licenses. Inside each folder you can find links to the original sources, as well as their licenses. 6 | 7 | Please be sure to check if the license of the model you want to use is compatible with your application. Also, make sure you cite the respective sources when using these models. 8 | 9 | More information about the models can be found [in this page of the documentation](https://ptlflow.readthedocs.io/en/latest/models/models_list.html). -------------------------------------------------------------------------------- /ptlflow/models/__init__.py: -------------------------------------------------------------------------------- 1 | from .ccmr import * 2 | from .craft import * 3 | from .csflow import * 4 | from .dicl import * 5 | from .dip import * 6 | from .dpflow import * 7 | from .fastflownet import * 8 | from .flow1d import * 9 | from .flowformer import * 10 | from .flowformerplusplus import * 11 | from .flownet import * 12 | from .gma import * 13 | from .gmflow import * 14 | from .gmflownet import * 15 | from .hd3 import * 16 | from .irr.pwcnet import * 17 | from .lcv import * 18 | from .liteflownet import * 19 | from .llaflow import * 20 | from .maskflownet import * 21 | from .matchflow import * 22 | from .memflow import * 23 | from .ms_raft_plus import * 24 | from .neuflow import * 25 | from .neuflow2 import * 26 | from .pwcnet import * 27 | from .raft import * 28 | from .rapidflow import * 29 | from .rpknet import * 30 | from .sea_raft import * 31 | from .scopeflow import * 32 | from .scv.scv import * 33 | from .separableflow import * 34 | from .skflow import * 35 | from .splatflow import * 36 | from .starflow import * 37 | from .streamflow import * 38 | from .unimatch import * 39 | from .vcn import * 40 | from .videoflow import * 41 | -------------------------------------------------------------------------------- /ptlflow/models/base_model/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hmorimitsu/ptlflow/fb6d753c3b16e74481258891b45668dcba63e54e/ptlflow/models/base_model/__init__.py -------------------------------------------------------------------------------- /ptlflow/models/ccmr/LICENSE: -------------------------------------------------------------------------------- 1 | Copyright (c) 2023 Azin Jahedi. All Rights Reserved. 2 | 3 | Redistribution and use in source and binary forms, with or without modification, 4 | are permitted provided that the following conditions are met: 5 | 6 | 1. Redistribution of source code must retain the above copyright notice, this 7 | list of conditions and the following disclaimer. 8 | 9 | 2. Redistribution in binary form must reproduce the above copyright notice, 10 | this list of conditions and the following disclaimer in the documentation 11 | and/or other materials provided with the distribution. 12 | 13 | 3. Neither the name of the copyright holder nor the names of its contributors 14 | may be used to endorse or promote products derived from this software without 15 | specific prior written permission. 16 | 17 | THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND 18 | ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED 19 | WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. 20 | IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, 21 | INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, 22 | BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, 23 | DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF 24 | LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE 25 | OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED 26 | OF THE POSSIBILITY OF SUCH DAMAGE. 27 | 28 | You acknowledge that this software is not designed, licensed or intended for use 29 | in the design, construction, operation or maintenance of any military facility. -------------------------------------------------------------------------------- /ptlflow/models/ccmr/README.md: -------------------------------------------------------------------------------- 1 | # CCMR 2 | 3 | ## Original code 4 | 5 | [https://github.com/cv-stuttgart/CCMR](https://github.com/cv-stuttgart/CCMR) 6 | 7 | ## Code license 8 | 9 | See [LICENSE](LICENSE). 10 | 11 | ## Pretrained weights license 12 | 13 | Not specified. 14 | 15 | ## Citation 16 | 17 | ``` 18 | @inproceedings{jahedi2024ccmr, 19 | title={CCMR: High Resolution Optical Flow Estimation via Coarse-to-Fine Context-Guided Motion Reasoning}, 20 | author={Jahedi, Azin and Luz, Maximilian and Rivinius, Marc and Bruhn, Andr{\'e}s}, 21 | booktitle={Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision}, 22 | pages={6899--6908}, 23 | year={2024} 24 | } 25 | ``` -------------------------------------------------------------------------------- /ptlflow/models/ccmr/__init__.py: -------------------------------------------------------------------------------- 1 | from .ccmr import * 2 | -------------------------------------------------------------------------------- /ptlflow/models/ccmr/helpers.py: -------------------------------------------------------------------------------- 1 | """ Layer/Module Helpers 2 | Hacked together by / Copyright 2020 Ross Wightman 3 | """ 4 | 5 | from itertools import repeat 6 | import collections.abc 7 | 8 | 9 | # From PyTorch internals 10 | def _ntuple(n): 11 | def parse(x): 12 | if isinstance(x, collections.abc.Iterable): 13 | return x 14 | return tuple(repeat(x, n)) 15 | 16 | return parse 17 | 18 | 19 | to_1tuple = _ntuple(1) 20 | to_2tuple = _ntuple(2) 21 | to_3tuple = _ntuple(3) 22 | to_4tuple = _ntuple(4) 23 | to_ntuple = _ntuple 24 | 25 | 26 | def make_divisible(v, divisor=8, min_value=None, round_limit=0.9): 27 | min_value = min_value or divisor 28 | new_v = max(min_value, int(v + divisor / 2) // divisor * divisor) 29 | # Make sure that round down does not go down by more than 10%. 30 | if new_v < round_limit * v: 31 | new_v += divisor 32 | return new_v 33 | -------------------------------------------------------------------------------- /ptlflow/models/craft/LICENSE: -------------------------------------------------------------------------------- 1 | 2 | DO WHAT THE FUCK YOU WANT TO PUBLIC LICENSE 3 | Version 2, December 2004 4 | 5 | Copyright (C) 2004 Sam Hocevar 6 | 7 | Everyone is permitted to copy and distribute verbatim or modified 8 | copies of this license document, and changing it is allowed as long 9 | as the name is changed. 10 | 11 | DO WHAT THE FUCK YOU WANT TO PUBLIC LICENSE 12 | TERMS AND CONDITIONS FOR COPYING, DISTRIBUTION AND MODIFICATION 13 | 14 | 0. You just DO WHAT THE FUCK YOU WANT TO. -------------------------------------------------------------------------------- /ptlflow/models/craft/README.md: -------------------------------------------------------------------------------- 1 | # CRAFT 2 | 3 | ## Original code 4 | 5 | [https://github.com/askerlee/craft](https://github.com/askerlee/craft) 6 | 7 | ## Code license 8 | 9 | See [LICENSE](LICENSE). 10 | 11 | ## Pretrained weights license 12 | 13 | Not specified. 14 | 15 | ## Citation 16 | 17 | ``` 18 | @InProceedings{craft, 19 | author="Sui, Xiuchao and Li, Shaohua and Geng, Xue and Wu, Yan and Xu, Xinxing and Liu, Yong and Goh, Rick Siow Mong and Zhu, Hongyuan", 20 | title="CRAFT: Cross-Attentional Flow Transformers for Robust Optical Flow", 21 | booktitle="CVPR", 22 | year="2022" 23 | } 24 | ``` -------------------------------------------------------------------------------- /ptlflow/models/craft/__init__.py: -------------------------------------------------------------------------------- 1 | from .craft import * 2 | -------------------------------------------------------------------------------- /ptlflow/models/csflow/LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2022 Hao 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. -------------------------------------------------------------------------------- /ptlflow/models/csflow/README.md: -------------------------------------------------------------------------------- 1 | # CSFlow 2 | 3 | ## Original code 4 | 5 | [https://github.com/MasterHow/CSFlow](https://github.com/MasterHow/CSFlow) 6 | 7 | ## Code license 8 | 9 | See [LICENSE](LICENSE). 10 | 11 | ## Pretrained weights license 12 | 13 | Not specified. 14 | 15 | ## Citation 16 | 17 | ``` 18 | @article{shi2022csflow, 19 | title={CSFlow: Learning optical flow via cross strip correlation for autonomous driving}, 20 | author={Shi, Hao and Zhou, Yifan and Yang, Kailun and Yin, Xiaoting and Wang, Kaiwei}, 21 | journal={arXiv preprint arXiv:2202.00909}, 22 | year={2022} 23 | } 24 | ``` -------------------------------------------------------------------------------- /ptlflow/models/csflow/__init__.py: -------------------------------------------------------------------------------- 1 | from .csflow import * 2 | -------------------------------------------------------------------------------- /ptlflow/models/dicl/README.md: -------------------------------------------------------------------------------- 1 | # DICL-Flow 2 | 3 | ## Original code 4 | 5 | [https://github.com/jytime/DICL-Flow](https://github.com/jytime/DICL-Flow) 6 | 7 | ## Code license 8 | 9 | See [https://github.com/jytime/DICL-Flow/issues/4](https://github.com/jytime/DICL-Flow/issues/4). 10 | 11 | ## Pretrained weights license 12 | 13 | Not specified. 14 | 15 | ## Citation 16 | 17 | ``` 18 | @article{wang2020displacement, 19 | title={Displacement-Invariant Matching Cost Learning for Accurate Optical Flow Estimation}, 20 | author={Wang, Jianyuan and Zhong, Yiran and Dai, Yuchao and Zhang, Kaihao and Ji, Pan and Li, Hongdong}, 21 | journal={Advances in Neural Information Processing Systems}, 22 | volume={33}, 23 | year={2020} 24 | } 25 | ``` -------------------------------------------------------------------------------- /ptlflow/models/dicl/__init__.py: -------------------------------------------------------------------------------- 1 | from .dicl import * 2 | -------------------------------------------------------------------------------- /ptlflow/models/dip/LICENSE: -------------------------------------------------------------------------------- 1 | BSD 3-Clause License 2 | 3 | Copyright (c) 2022, zihuazheng 4 | 5 | Redistribution and use in source and binary forms, with or without 6 | modification, are permitted provided that the following conditions are met: 7 | 8 | 1. Redistributions of source code must retain the above copyright notice, this 9 | list of conditions and the following disclaimer. 10 | 11 | 2. Redistributions in binary form must reproduce the above copyright notice, 12 | this list of conditions and the following disclaimer in the documentation 13 | and/or other materials provided with the distribution. 14 | 15 | 3. Neither the name of the copyright holder nor the names of its 16 | contributors may be used to endorse or promote products derived from 17 | this software without specific prior written permission. 18 | 19 | THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" 20 | AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 21 | IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE 22 | DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE 23 | FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL 24 | DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR 25 | SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER 26 | CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, 27 | OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 28 | OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. -------------------------------------------------------------------------------- /ptlflow/models/dip/README.md: -------------------------------------------------------------------------------- 1 | # DIP 2 | 3 | ## Original code 4 | 5 | [https://github.com/zihuazheng/DIP](https://github.com/zihuazheng/DIP) 6 | 7 | ## Code license 8 | 9 | See [LICENSE](LICENSE). 10 | 11 | ## Pretrained weights license 12 | 13 | Not specified. 14 | 15 | ## Citation 16 | 17 | ``` 18 | @InProceedings{Zheng2022DIPDeepInverse, 19 | author = {Zihua Zheng and Ni Nie and Zhi Ling and Pengfei Xiong and Jiangyu Liu and Hongya Wang and Jiankun Li}, 20 | booktitle = {CVPR}, 21 | title = {{DIP}: Deep Inverse Patchmatch for High-Resolution Optical Flow}, 22 | year = {2022}, 23 | pages = {8915-8924}, 24 | } 25 | ``` -------------------------------------------------------------------------------- /ptlflow/models/dip/__init__.py: -------------------------------------------------------------------------------- 1 | from .dip import * 2 | -------------------------------------------------------------------------------- /ptlflow/models/dpflow/__init__.py: -------------------------------------------------------------------------------- 1 | from .dpflow import dpflow 2 | -------------------------------------------------------------------------------- /ptlflow/models/dpflow/assets/teaser.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hmorimitsu/ptlflow/fb6d753c3b16e74481258891b45668dcba63e54e/ptlflow/models/dpflow/assets/teaser.jpg -------------------------------------------------------------------------------- /ptlflow/models/dpflow/configs/dpflow-test.yaml: -------------------------------------------------------------------------------- 1 | # lightning.pytorch==2.4.0 2 | output_path: outputs/test 3 | show: false 4 | max_forward_side: null 5 | scale_factor: null 6 | max_show_side: 1000 7 | save_viz: true 8 | model: 9 | class_path: ptlflow.models.dpflow 10 | init_args: 11 | iters_per_level: 4 12 | detach_flow: true 13 | use_norm_affine: false 14 | group_norm_num_groups: 8 15 | corr_mode: allpairs 16 | corr_levels: 1 17 | corr_range: 4 18 | activation_function: orig 19 | enc_network: cgu_bidir_dual 20 | enc_norm_type: group 21 | enc_depth: 4 22 | enc_mlp_ratio: 2.0 23 | enc_mlp_in_kernel_size: 1 24 | enc_mlp_out_kernel_size: 1 25 | enc_hidden_chs: 26 | - 64 27 | - 96 28 | - 128 29 | enc_num_out_stages: 1 30 | enc_out_1x1_chs: '384' 31 | dec_gru_norm_type: layer 32 | dec_gru_iters: 1 33 | dec_gru_depth: 4 34 | dec_gru_mlp_ratio: 2.0 35 | dec_gru_mlp_in_kernel_size: 1 36 | dec_gru_mlp_out_kernel_size: 1 37 | dec_net_chs: 128 38 | dec_inp_chs: 128 39 | dec_motion_chs: 128 40 | dec_flow_kernel_size: 7 41 | dec_flow_head_chs: 256 42 | dec_motenc_corr_hidden_chs: 256 43 | dec_motenc_corr_out_chs: 192 44 | dec_motenc_flow_hidden_chs: 128 45 | dec_motenc_flow_out_chs: 64 46 | use_upsample_mask: true 47 | upmask_gradient_scale: 1.0 48 | cgu_mlp_dw_kernel_size: 7 49 | cgu_fusion_gate_activation: gelu 50 | cgu_mlp_use_dw_conv: true 51 | cgu_mlp_activation_function: gelu 52 | cgu_layer_scale_init_value: 0.01 53 | loss: laplace 54 | gamma: 0.8 55 | max_flow: 400.0 56 | use_var: true 57 | var_min: 0.0 58 | var_max: 10.0 59 | warm_start: true 60 | -------------------------------------------------------------------------------- /ptlflow/models/dpflow/configs/dpflow-train1-chairs.yaml: -------------------------------------------------------------------------------- 1 | # lightning.pytorch==2.4.0 2 | seed_everything: true 3 | lr: 2.5e-4 4 | wdecay: 1.0e-4 5 | trainer: 6 | max_epochs: 45 7 | accumulate_grad_batches: 1 8 | gradient_clip_val: 1.0 9 | devices: 0,1 10 | model: 11 | class_path: ptlflow.models.dpflow 12 | init_args: 13 | pyramid_levels: 3 14 | iters_per_level: 4 15 | detach_flow: true 16 | use_norm_affine: false 17 | group_norm_num_groups: 8 18 | corr_mode: allpairs 19 | corr_levels: 1 20 | corr_range: 4 21 | activation_function: orig 22 | enc_network: cgu_bidir_dual 23 | enc_norm_type: group 24 | enc_depth: 4 25 | enc_mlp_ratio: 2.0 26 | enc_mlp_in_kernel_size: 1 27 | enc_mlp_out_kernel_size: 1 28 | enc_hidden_chs: 29 | - 64 30 | - 96 31 | - 128 32 | enc_num_out_stages: 1 33 | enc_out_1x1_chs: '384' 34 | dec_gru_norm_type: layer 35 | dec_gru_iters: 1 36 | dec_gru_depth: 4 37 | dec_gru_mlp_ratio: 2.0 38 | dec_gru_mlp_in_kernel_size: 1 39 | dec_gru_mlp_out_kernel_size: 1 40 | dec_net_chs: 128 41 | dec_inp_chs: 128 42 | dec_motion_chs: 128 43 | dec_flow_kernel_size: 7 44 | dec_flow_head_chs: 256 45 | dec_motenc_corr_hidden_chs: 256 46 | dec_motenc_corr_out_chs: 192 47 | dec_motenc_flow_hidden_chs: 128 48 | dec_motenc_flow_out_chs: 64 49 | use_upsample_mask: true 50 | upmask_gradient_scale: 1.0 51 | cgu_mlp_dw_kernel_size: 7 52 | cgu_fusion_gate_activation: gelu 53 | cgu_mlp_use_dw_conv: true 54 | cgu_mlp_activation_function: gelu 55 | cgu_layer_scale_init_value: 0.01 56 | loss: laplace 57 | gamma: 0.8 58 | max_flow: 400.0 59 | use_var: true 60 | var_min: 0.0 61 | var_max: 10.0 62 | warm_start: false 63 | data: 64 | train_dataset: chairs 65 | val_dataset: sintel-final-val+kitti-2015-val 66 | train_batch_size: 5 67 | train_num_workers: 5 68 | train_crop_size: [352, 480] 69 | train_transform_cuda: false 70 | train_transform_fp16: false 71 | -------------------------------------------------------------------------------- /ptlflow/models/dpflow/configs/dpflow-train4b-spring.yaml: -------------------------------------------------------------------------------- 1 | # lightning.pytorch==2.4.0 2 | seed_everything: true 3 | ckpt_path: sintel # Change to the ckpt resulting from dpflow-train3-sintel 4 | lr: 1.0e-4 5 | wdecay: 1.0e-5 6 | trainer: 7 | max_epochs: 100 8 | accumulate_grad_batches: 4 9 | gradient_clip_val: 1.0 10 | devices: 0,1 11 | model: 12 | class_path: ptlflow.models.dpflow 13 | init_args: 14 | pyramid_levels: 4 15 | iters_per_level: 4 16 | detach_flow: true 17 | use_norm_affine: false 18 | group_norm_num_groups: 8 19 | corr_mode: allpairs 20 | corr_levels: 1 21 | corr_range: 4 22 | activation_function: orig 23 | enc_network: cgu_bidir_dual 24 | enc_norm_type: group 25 | enc_depth: 4 26 | enc_mlp_ratio: 2.0 27 | enc_mlp_in_kernel_size: 1 28 | enc_mlp_out_kernel_size: 1 29 | enc_hidden_chs: 30 | - 64 31 | - 96 32 | - 128 33 | enc_num_out_stages: 1 34 | enc_out_1x1_chs: '384' 35 | dec_gru_norm_type: layer 36 | dec_gru_iters: 1 37 | dec_gru_depth: 4 38 | dec_gru_mlp_ratio: 2.0 39 | dec_gru_mlp_in_kernel_size: 1 40 | dec_gru_mlp_out_kernel_size: 1 41 | dec_net_chs: 128 42 | dec_inp_chs: 128 43 | dec_motion_chs: 128 44 | dec_flow_kernel_size: 7 45 | dec_flow_head_chs: 256 46 | dec_motenc_corr_hidden_chs: 256 47 | dec_motenc_corr_out_chs: 192 48 | dec_motenc_flow_hidden_chs: 128 49 | dec_motenc_flow_out_chs: 64 50 | use_upsample_mask: true 51 | upmask_gradient_scale: 1.0 52 | cgu_mlp_dw_kernel_size: 7 53 | cgu_fusion_gate_activation: gelu 54 | cgu_mlp_use_dw_conv: true 55 | cgu_mlp_activation_function: gelu 56 | cgu_layer_scale_init_value: 0.01 57 | loss: laplace 58 | gamma: 0.85 59 | max_flow: 400.0 60 | use_var: true 61 | var_min: 0.0 62 | var_max: 10.0 63 | warm_start: false 64 | data: 65 | train_dataset: spring-sub-rev 66 | val_dataset: sintel-final-val+kitti-2015-val 67 | train_batch_size: 2 68 | train_num_workers: 2 69 | train_crop_size: [540, 960] 70 | train_transform_cuda: false 71 | train_transform_fp16: false 72 | -------------------------------------------------------------------------------- /ptlflow/models/dpflow/configs/dpflow-validate.yaml: -------------------------------------------------------------------------------- 1 | # lightning.pytorch==2.4.0 2 | all: false 3 | select: null 4 | exclude: null 5 | output_path: outputs/validate 6 | write_outputs: false 7 | show: false 8 | flow_format: original 9 | max_forward_side: null 10 | scale_factor: null 11 | max_show_side: 1000 12 | max_samples: null 13 | reversed: false 14 | fp16: false 15 | seq_val_mode: all 16 | write_individual_metrics: false 17 | epe_clip: 5.0 18 | seed_everything: true 19 | model: 20 | class_path: ptlflow.models.dpflow 21 | init_args: 22 | iters_per_level: 4 23 | detach_flow: true 24 | use_norm_affine: false 25 | group_norm_num_groups: 8 26 | corr_mode: allpairs 27 | corr_levels: 1 28 | corr_range: 4 29 | activation_function: orig 30 | enc_network: cgu_bidir_dual 31 | enc_norm_type: group 32 | enc_depth: 4 33 | enc_mlp_ratio: 2.0 34 | enc_mlp_in_kernel_size: 1 35 | enc_mlp_out_kernel_size: 1 36 | enc_hidden_chs: 37 | - 64 38 | - 96 39 | - 128 40 | enc_num_out_stages: 1 41 | enc_out_1x1_chs: '384' 42 | dec_gru_norm_type: layer 43 | dec_gru_iters: 1 44 | dec_gru_depth: 4 45 | dec_gru_mlp_ratio: 2.0 46 | dec_gru_mlp_in_kernel_size: 1 47 | dec_gru_mlp_out_kernel_size: 1 48 | dec_net_chs: 128 49 | dec_inp_chs: 128 50 | dec_motion_chs: 128 51 | dec_flow_kernel_size: 7 52 | dec_flow_head_chs: 256 53 | dec_motenc_corr_hidden_chs: 256 54 | dec_motenc_corr_out_chs: 192 55 | dec_motenc_flow_hidden_chs: 128 56 | dec_motenc_flow_out_chs: 64 57 | use_upsample_mask: true 58 | upmask_gradient_scale: 1.0 59 | cgu_mlp_dw_kernel_size: 7 60 | cgu_fusion_gate_activation: gelu 61 | cgu_mlp_use_dw_conv: true 62 | cgu_mlp_activation_function: gelu 63 | cgu_layer_scale_init_value: 0.01 64 | loss: laplace 65 | gamma: 0.8 66 | max_flow: 400.0 67 | use_var: true 68 | var_min: 0.0 69 | var_max: 10.0 70 | warm_start: false 71 | -------------------------------------------------------------------------------- /ptlflow/models/dpflow/local_timm/conv2d_same.py: -------------------------------------------------------------------------------- 1 | """ Conv2d w/ Same Padding 2 | 3 | Hacked together by / Copyright 2020 Ross Wightman 4 | """ 5 | 6 | import torch 7 | import torch.nn as nn 8 | import torch.nn.functional as F 9 | from typing import Tuple, Optional 10 | 11 | from .padding import pad_same, get_padding_value 12 | 13 | 14 | def conv2d_same( 15 | x, 16 | weight: torch.Tensor, 17 | bias: Optional[torch.Tensor] = None, 18 | stride: Tuple[int, int] = (1, 1), 19 | padding: Tuple[int, int] = (0, 0), 20 | dilation: Tuple[int, int] = (1, 1), 21 | groups: int = 1, 22 | ): 23 | x = pad_same(x, weight.shape[-2:], stride, dilation) 24 | return F.conv2d(x, weight, bias, stride, (0, 0), dilation, groups) 25 | 26 | 27 | class Conv2dSame(nn.Conv2d): 28 | """Tensorflow like 'SAME' convolution wrapper for 2D convolutions""" 29 | 30 | def __init__( 31 | self, 32 | in_channels, 33 | out_channels, 34 | kernel_size, 35 | stride=1, 36 | padding=0, 37 | dilation=1, 38 | groups=1, 39 | bias=True, 40 | ): 41 | super(Conv2dSame, self).__init__( 42 | in_channels, out_channels, kernel_size, stride, 0, dilation, groups, bias 43 | ) 44 | 45 | def forward(self, x): 46 | return conv2d_same( 47 | x, 48 | self.weight, 49 | self.bias, 50 | self.stride, 51 | self.padding, 52 | self.dilation, 53 | self.groups, 54 | ) 55 | 56 | 57 | def create_conv2d_pad(in_chs, out_chs, kernel_size, **kwargs): 58 | padding = kwargs.pop("padding", "") 59 | kwargs.setdefault("bias", False) 60 | padding, is_dynamic = get_padding_value(padding, kernel_size, **kwargs) 61 | if is_dynamic: 62 | return Conv2dSame(in_chs, out_chs, kernel_size, **kwargs) 63 | else: 64 | return nn.Conv2d(in_chs, out_chs, kernel_size, padding=padding, **kwargs) 65 | -------------------------------------------------------------------------------- /ptlflow/models/dpflow/local_timm/create_conv2d.py: -------------------------------------------------------------------------------- 1 | """ Create Conv2d Factory Method 2 | 3 | Hacked together by / Copyright 2020 Ross Wightman 4 | """ 5 | 6 | from .mixed_conv2d import MixedConv2d 7 | from .cond_conv2d import CondConv2d 8 | from .conv2d_same import create_conv2d_pad 9 | 10 | 11 | def create_conv2d(in_channels, out_channels, kernel_size, **kwargs): 12 | """Select a 2d convolution implementation based on arguments 13 | Creates and returns one of torch.nn.Conv2d, Conv2dSame, MixedConv2d, or CondConv2d. 14 | 15 | Used extensively by EfficientNet, MobileNetv3 and related networks. 16 | """ 17 | if isinstance(kernel_size, list): 18 | assert ( 19 | "num_experts" not in kwargs 20 | ) # MixNet + CondConv combo not supported currently 21 | if "groups" in kwargs: 22 | groups = kwargs.pop("groups") 23 | if groups == in_channels: 24 | kwargs["depthwise"] = True 25 | else: 26 | assert groups == 1 27 | # We're going to use only lists for defining the MixedConv2d kernel groups, 28 | # ints, tuples, other iterables will continue to pass to normal conv and specify h, w. 29 | m = MixedConv2d(in_channels, out_channels, kernel_size, **kwargs) 30 | else: 31 | depthwise = kwargs.pop("depthwise", False) 32 | # for DW out_channels must be multiple of in_channels as must have out_channels % groups == 0 33 | groups = in_channels if depthwise else kwargs.pop("groups", 1) 34 | if "num_experts" in kwargs and kwargs["num_experts"] > 0: 35 | m = CondConv2d( 36 | in_channels, out_channels, kernel_size, groups=groups, **kwargs 37 | ) 38 | else: 39 | m = create_conv2d_pad( 40 | in_channels, out_channels, kernel_size, groups=groups, **kwargs 41 | ) 42 | return m 43 | -------------------------------------------------------------------------------- /ptlflow/models/dpflow/local_timm/gelu.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | 5 | 6 | class GELU(nn.Module): 7 | """Applies the Gaussian Error Linear Units function (w/ dummy inplace arg)""" 8 | 9 | def __init__(self, inplace: bool = False): 10 | super(GELU, self).__init__() 11 | 12 | def forward(self, input: torch.Tensor) -> torch.Tensor: 13 | return F.gelu(input) 14 | -------------------------------------------------------------------------------- /ptlflow/models/dpflow/local_timm/grn.py: -------------------------------------------------------------------------------- 1 | """ Global Response Normalization Module 2 | 3 | Based on the GRN layer presented in 4 | `ConvNeXt-V2 - Co-designing and Scaling ConvNets with Masked Autoencoders` - https://arxiv.org/abs/2301.00808 5 | 6 | This implementation 7 | * works for both NCHW and NHWC tensor layouts 8 | * uses affine param names matching existing torch norm layers 9 | * slightly improves eager mode performance via fused addcmul 10 | 11 | Hacked together by / Copyright 2023 Ross Wightman 12 | """ 13 | 14 | import torch 15 | from torch import nn as nn 16 | 17 | 18 | class GlobalResponseNorm(nn.Module): 19 | """Global Response Normalization layer""" 20 | 21 | def __init__(self, dim, eps=1e-6, channels_last=True): 22 | super().__init__() 23 | self.eps = eps 24 | if channels_last: 25 | self.spatial_dim = (1, 2) 26 | self.channel_dim = -1 27 | self.wb_shape = (1, 1, 1, -1) 28 | else: 29 | self.spatial_dim = (2, 3) 30 | self.channel_dim = 1 31 | self.wb_shape = (1, -1, 1, 1) 32 | 33 | self.weight = nn.Parameter(torch.zeros(dim)) 34 | self.bias = nn.Parameter(torch.zeros(dim)) 35 | 36 | def forward(self, x): 37 | x_g = x.norm(p=2, dim=self.spatial_dim, keepdim=True) 38 | x_n = x_g / (x_g.mean(dim=self.channel_dim, keepdim=True) + self.eps) 39 | return x + torch.addcmul( 40 | self.bias.view(self.wb_shape), self.weight.view(self.wb_shape), x * x_n 41 | ) 42 | -------------------------------------------------------------------------------- /ptlflow/models/dpflow/local_timm/layer_helpers.py: -------------------------------------------------------------------------------- 1 | """ Layer/Module Helpers 2 | 3 | Hacked together by / Copyright 2020 Ross Wightman 4 | """ 5 | 6 | from itertools import repeat 7 | import collections.abc 8 | 9 | 10 | # From PyTorch internals 11 | def _ntuple(n): 12 | def parse(x): 13 | if isinstance(x, collections.abc.Iterable) and not isinstance(x, str): 14 | return x 15 | return tuple(repeat(x, n)) 16 | 17 | return parse 18 | 19 | 20 | to_1tuple = _ntuple(1) 21 | to_2tuple = _ntuple(2) 22 | to_3tuple = _ntuple(3) 23 | to_4tuple = _ntuple(4) 24 | to_ntuple = _ntuple 25 | 26 | 27 | def make_divisible(v, divisor=8, min_value=None, round_limit=0.9): 28 | min_value = min_value or divisor 29 | new_v = max(min_value, int(v + divisor / 2) // divisor * divisor) 30 | # Make sure that round down does not go down by more than 10%. 31 | if new_v < round_limit * v: 32 | new_v += divisor 33 | return new_v 34 | 35 | 36 | def extend_tuple(x, n): 37 | # pdas a tuple to specified n by padding with last value 38 | if not isinstance(x, (tuple, list)): 39 | x = (x,) 40 | else: 41 | x = tuple(x) 42 | pad_n = n - len(x) 43 | if pad_n <= 0: 44 | return x[:n] 45 | return x + (x[-1],) * pad_n 46 | -------------------------------------------------------------------------------- /ptlflow/models/dpflow/local_timm/norm.py: -------------------------------------------------------------------------------- 1 | """ Normalization layers and wrappers 2 | 3 | Norm layer definitions that support fast norm and consistent channel arg order (always first arg). 4 | 5 | Hacked together by / Copyright 2022 Ross Wightman 6 | """ 7 | 8 | import torch 9 | import torch.nn as nn 10 | import torch.nn.functional as F 11 | 12 | 13 | class LayerNorm(nn.LayerNorm): 14 | """LayerNorm w/ fast norm option""" 15 | 16 | def __init__(self, num_channels, eps=1e-6, affine=True): 17 | super().__init__(num_channels, eps=eps, elementwise_affine=affine) 18 | 19 | def forward(self, x: torch.Tensor) -> torch.Tensor: 20 | x = F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps) 21 | return x 22 | 23 | 24 | class LayerNorm2d(nn.LayerNorm): 25 | """LayerNorm for channels of '2D' spatial NCHW tensors""" 26 | 27 | def __init__(self, num_channels, eps=1e-6, affine=True): 28 | super().__init__(num_channels, eps=eps, elementwise_affine=affine) 29 | 30 | def forward(self, x: torch.Tensor) -> torch.Tensor: 31 | x = x.permute(0, 2, 3, 1) 32 | x = F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps) 33 | x = x.permute(0, 3, 1, 2) 34 | return x 35 | -------------------------------------------------------------------------------- /ptlflow/models/dpflow/pwc_modules.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import, division, print_function 2 | 3 | import torch 4 | import torch.nn as nn 5 | import torch.nn.functional as F 6 | 7 | 8 | def upsample2d_as(inputs, target_as, mode="bilinear"): 9 | _, _, h, w = target_as.size() 10 | if inputs.shape[-2] != h or inputs.shape[-1] != w: 11 | inputs = F.interpolate(inputs, [h, w], mode=mode, align_corners=True) 12 | return inputs 13 | 14 | 15 | def rescale_flow(flow, width_im, height_im, to_local=True): 16 | if to_local: 17 | u_scale = float(flow.size(3) / width_im) 18 | v_scale = float(flow.size(2) / height_im) 19 | else: 20 | u_scale = float(width_im / flow.size(3)) 21 | v_scale = float(height_im / flow.size(2)) 22 | 23 | u, v = flow.chunk(2, dim=1) 24 | u = u * u_scale 25 | v = v * v_scale 26 | 27 | return torch.cat([u, v], dim=1) 28 | -------------------------------------------------------------------------------- /ptlflow/models/fastflownet/LICENSE: -------------------------------------------------------------------------------- 1 | This software and associated documentation files (the "Software"), and the research paper (FastFlowNet: A Lightweight Network for Fast Optical Flow Estimation) including but not limited to the figures, and tables (the "Paper") are provided for academic research purposes only and without any warranty. Any commercial use requires my consent. When using any parts of the Software or the Paper in your work, please cite the paper. -------------------------------------------------------------------------------- /ptlflow/models/fastflownet/README.md: -------------------------------------------------------------------------------- 1 | # FastFlowNet 2 | 3 | ## Original code 4 | 5 | [https://github.com/ltkong218/FastFlowNet](https://github.com/ltkong218/FastFlowNet) 6 | 7 | ## Code license 8 | 9 | See [LICENSE](LICENSE). 10 | 11 | ## Pretrained weights license 12 | 13 | Not specified. 14 | 15 | ## Citation 16 | 17 | ``` 18 | @inproceedings{Kong:2021:FastFlowNet, 19 | title = {FastFlowNet: A Lightweight Network for Fast Optical Flow Estimation}, 20 | author = {Lingtong Kong and Chunhua Shen and Jie Yang}, 21 | booktitle = {2021 IEEE International Conference on Robotics and Automation (ICRA)}, 22 | year = {2021} 23 | } 24 | ``` 25 | -------------------------------------------------------------------------------- /ptlflow/models/fastflownet/__init__.py: -------------------------------------------------------------------------------- 1 | from .fastflownet import * 2 | -------------------------------------------------------------------------------- /ptlflow/models/flow1d/LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2022 Haofei Xu 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. -------------------------------------------------------------------------------- /ptlflow/models/flow1d/README.md: -------------------------------------------------------------------------------- 1 | # Flow1D 2 | 3 | ## Original code 4 | 5 | [https://github.com/haofeixu/flow1d](https://github.com/haofeixu/flow1d) 6 | 7 | ## Code license 8 | 9 | See [LICENSE](LICENSE). 10 | 11 | ## Pretrained weights license 12 | 13 | Not specified. 14 | 15 | ## Citation 16 | 17 | ``` 18 | @inproceedings{xu2021high, 19 | title={High-Resolution Optical Flow from 1D Attention and Correlation}, 20 | author={Xu, Haofei and Yang, Jiaolong and Cai, Jianfei and Zhang, Juyong and Tong, Xin}, 21 | booktitle={Proceedings of the IEEE/CVF International Conference on Computer Vision}, 22 | pages={10498--10507}, 23 | year={2021} 24 | } 25 | ``` -------------------------------------------------------------------------------- /ptlflow/models/flow1d/__init__.py: -------------------------------------------------------------------------------- 1 | from .flow1d import * 2 | -------------------------------------------------------------------------------- /ptlflow/models/flowformer/README.md: -------------------------------------------------------------------------------- 1 | # FlowFormer 2 | 3 | ## Original code 4 | 5 | [https://github.com/drinkingcoder/FlowFormer-Official](https://github.com/drinkingcoder/FlowFormer-Official) 6 | 7 | ## Code license 8 | 9 | See [LICENSE](LICENSE). 10 | 11 | ## Pretrained weights license 12 | 13 | Not specified. 14 | 15 | ## Citation 16 | 17 | ``` 18 | @article{huang2022flowformer, 19 | title={{FlowFormer}: A Transformer Architecture for Optical Flow}, 20 | author={Huang, Zhaoyang and Shi, Xiaoyu and Zhang, Chao and Wang, Qiang and Cheung, Ka Chun and Qin, Hongwei and Dai, Jifeng and Li, Hongsheng}, 21 | journal={{ECCV}}, 22 | year={2022} 23 | } 24 | ``` -------------------------------------------------------------------------------- /ptlflow/models/flowformer/__init__.py: -------------------------------------------------------------------------------- 1 | from .flowformer import * 2 | -------------------------------------------------------------------------------- /ptlflow/models/flowformer/mlpmixer.py: -------------------------------------------------------------------------------- 1 | from torch import nn 2 | from functools import partial 3 | import numpy as np 4 | 5 | 6 | class PreNormResidual(nn.Module): 7 | def __init__(self, dim, fn): 8 | super().__init__() 9 | self.fn = fn 10 | self.norm = nn.LayerNorm(dim) 11 | 12 | def forward(self, x): 13 | return self.fn(self.norm(x)) + x 14 | 15 | 16 | def FeedForward(dim, expansion_factor=4, dropout=0.0, dense=nn.Linear): 17 | return nn.Sequential( 18 | dense(dim, dim * expansion_factor), 19 | nn.GELU(), 20 | nn.Dropout(dropout), 21 | dense(dim * expansion_factor, dim), 22 | nn.Dropout(dropout), 23 | ) 24 | 25 | 26 | class MLPMixerLayer(nn.Module): 27 | def __init__(self, dim, cost_latent_token_num, mlp_expansion_factor, dropout=0.0): 28 | super(MLPMixerLayer, self).__init__() 29 | 30 | # print(f"use mlp mixer layer") 31 | K = cost_latent_token_num 32 | expansion_factor = mlp_expansion_factor 33 | chan_first, chan_last = partial(nn.Conv1d, kernel_size=1), nn.Linear 34 | 35 | self.mlpmixer = nn.Sequential( 36 | PreNormResidual(dim, FeedForward(K, expansion_factor, dropout, chan_first)), 37 | PreNormResidual( 38 | dim, FeedForward(dim, expansion_factor, dropout, chan_last) 39 | ), 40 | ) 41 | 42 | def compute_params(self): 43 | num = 0 44 | for param in self.mlpmixer.parameters(): 45 | num += np.prod(param.size()) 46 | 47 | return num 48 | 49 | def forward(self, x): 50 | """ 51 | x: [BH1W1, K, D] 52 | """ 53 | 54 | return self.mlpmixer(x) 55 | -------------------------------------------------------------------------------- /ptlflow/models/flowformerplusplus/FlowFormer/PerCostFormer3/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hmorimitsu/ptlflow/fb6d753c3b16e74481258891b45668dcba63e54e/ptlflow/models/flowformerplusplus/FlowFormer/PerCostFormer3/__init__.py -------------------------------------------------------------------------------- /ptlflow/models/flowformerplusplus/FlowFormer/__init__.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | 4 | def build_flowformer(cfg): 5 | name = cfg.transformer 6 | if name == "percostformer3": 7 | from .PerCostFormer3.transformer import FlowFormer 8 | else: 9 | raise ValueError(f"FlowFormer = {name} is not a valid optimizer!") 10 | 11 | return FlowFormer(cfg[name]) 12 | -------------------------------------------------------------------------------- /ptlflow/models/flowformerplusplus/README.md: -------------------------------------------------------------------------------- 1 | # FlowFormer++ 2 | 3 | ## Original code 4 | 5 | [https://github.com/XiaoyuShi97/FlowFormerPlusPlus](https://github.com/XiaoyuShi97/FlowFormerPlusPlus) 6 | 7 | ## Code license 8 | 9 | Apache License. 10 | 11 | ## Pretrained weights license 12 | 13 | Not specified. 14 | 15 | ## Citation 16 | 17 | ``` 18 | @inproceedings{shi2023flowformer++, 19 | title={Flowformer++: Masked cost volume autoencoding for pretraining optical flow estimation}, 20 | author={Shi, Xiaoyu and Huang, Zhaoyang and Li, Dasong and Zhang, Manyuan and Cheung, Ka Chun and See, Simon and Qin, Hongwei and Dai, Jifeng and Li, Hongsheng}, 21 | booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition}, 22 | pages={1599--1610}, 23 | year={2023} 24 | } 25 | ``` -------------------------------------------------------------------------------- /ptlflow/models/flowformerplusplus/__init__.py: -------------------------------------------------------------------------------- 1 | from .flowformerplusplus import * 2 | -------------------------------------------------------------------------------- /ptlflow/models/flowformerplusplus/timm0412/data/__init__.py: -------------------------------------------------------------------------------- 1 | from .constants import * 2 | 3 | # from .auto_augment import RandAugment, AutoAugment, rand_augment_ops, auto_augment_policy,\ 4 | # rand_augment_transform, auto_augment_transform 5 | # from .config import resolve_data_config 6 | # from .dataset import ImageDataset, IterableImageDataset, AugMixDataset 7 | # from .dataset_factory import create_dataset 8 | # from .loader import create_loader 9 | # from .mixup import Mixup, FastCollateMixup 10 | # from .parsers import create_parser 11 | # from .real_labels import RealLabelsImagenet 12 | # from .transforms import * 13 | # from .transforms_factory import create_transform 14 | -------------------------------------------------------------------------------- /ptlflow/models/flowformerplusplus/timm0412/data/constants.py: -------------------------------------------------------------------------------- 1 | DEFAULT_CROP_PCT = 0.875 2 | IMAGENET_DEFAULT_MEAN = (0.485, 0.456, 0.406) 3 | IMAGENET_DEFAULT_STD = (0.229, 0.224, 0.225) 4 | IMAGENET_INCEPTION_MEAN = (0.5, 0.5, 0.5) 5 | IMAGENET_INCEPTION_STD = (0.5, 0.5, 0.5) 6 | IMAGENET_DPN_MEAN = (124 / 255, 117 / 255, 104 / 255) 7 | IMAGENET_DPN_STD = tuple([1 / (0.0167 * 255)] * 3) 8 | -------------------------------------------------------------------------------- /ptlflow/models/flowformerplusplus/timm0412/models/layers/conv2d_same.py: -------------------------------------------------------------------------------- 1 | """ Conv2d w/ Same Padding 2 | 3 | Hacked together by / Copyright 2020 Ross Wightman 4 | """ 5 | 6 | import torch 7 | import torch.nn as nn 8 | import torch.nn.functional as F 9 | from typing import Tuple, Optional 10 | 11 | from .padding import pad_same, get_padding_value 12 | 13 | 14 | def conv2d_same( 15 | x, 16 | weight: torch.Tensor, 17 | bias: Optional[torch.Tensor] = None, 18 | stride: Tuple[int, int] = (1, 1), 19 | padding: Tuple[int, int] = (0, 0), 20 | dilation: Tuple[int, int] = (1, 1), 21 | groups: int = 1, 22 | ): 23 | x = pad_same(x, weight.shape[-2:], stride, dilation) 24 | return F.conv2d(x, weight, bias, stride, (0, 0), dilation, groups) 25 | 26 | 27 | class Conv2dSame(nn.Conv2d): 28 | """Tensorflow like 'SAME' convolution wrapper for 2D convolutions""" 29 | 30 | def __init__( 31 | self, 32 | in_channels, 33 | out_channels, 34 | kernel_size, 35 | stride=1, 36 | padding=0, 37 | dilation=1, 38 | groups=1, 39 | bias=True, 40 | ): 41 | super(Conv2dSame, self).__init__( 42 | in_channels, out_channels, kernel_size, stride, 0, dilation, groups, bias 43 | ) 44 | 45 | def forward(self, x): 46 | return conv2d_same( 47 | x, 48 | self.weight, 49 | self.bias, 50 | self.stride, 51 | self.padding, 52 | self.dilation, 53 | self.groups, 54 | ) 55 | 56 | 57 | def create_conv2d_pad(in_chs, out_chs, kernel_size, **kwargs): 58 | padding = kwargs.pop("padding", "") 59 | kwargs.setdefault("bias", False) 60 | padding, is_dynamic = get_padding_value(padding, kernel_size, **kwargs) 61 | if is_dynamic: 62 | return Conv2dSame(in_chs, out_chs, kernel_size, **kwargs) 63 | else: 64 | return nn.Conv2d(in_chs, out_chs, kernel_size, padding=padding, **kwargs) 65 | -------------------------------------------------------------------------------- /ptlflow/models/flowformerplusplus/timm0412/models/layers/helpers.py: -------------------------------------------------------------------------------- 1 | """ Layer/Module Helpers 2 | 3 | Hacked together by / Copyright 2020 Ross Wightman 4 | """ 5 | 6 | from itertools import repeat 7 | import collections.abc 8 | 9 | 10 | # From PyTorch internals 11 | def _ntuple(n): 12 | def parse(x): 13 | if isinstance(x, collections.abc.Iterable): 14 | return x 15 | return tuple(repeat(x, n)) 16 | 17 | return parse 18 | 19 | 20 | to_1tuple = _ntuple(1) 21 | to_2tuple = _ntuple(2) 22 | to_3tuple = _ntuple(3) 23 | to_4tuple = _ntuple(4) 24 | to_ntuple = _ntuple 25 | 26 | 27 | def make_divisible(v, divisor=8, min_value=None, round_limit=0.9): 28 | min_value = min_value or divisor 29 | new_v = max(min_value, int(v + divisor / 2) // divisor * divisor) 30 | # Make sure that round down does not go down by more than 10%. 31 | if new_v < round_limit * v: 32 | new_v += divisor 33 | return new_v 34 | -------------------------------------------------------------------------------- /ptlflow/models/flowformerplusplus/timm0412/models/layers/linear.py: -------------------------------------------------------------------------------- 1 | """ Linear layer (alternate definition) 2 | """ 3 | 4 | import torch 5 | import torch.nn.functional as F 6 | from torch import nn as nn 7 | 8 | 9 | class Linear(nn.Linear): 10 | r"""Applies a linear transformation to the incoming data: :math:`y = xA^T + b` 11 | 12 | Wraps torch.nn.Linear to support AMP + torchscript usage by manually casting 13 | weight & bias to input.dtype to work around an issue w/ torch.addmm in this use case. 14 | """ 15 | 16 | def forward(self, input: torch.Tensor) -> torch.Tensor: 17 | if torch.jit.is_scripting(): 18 | bias = self.bias.to(dtype=input.dtype) if self.bias is not None else None 19 | return F.linear(input, self.weight.to(dtype=input.dtype), bias=bias) 20 | else: 21 | return F.linear(input, self.weight, self.bias) 22 | -------------------------------------------------------------------------------- /ptlflow/models/flowformerplusplus/timm0412/models/layers/patch_embed.py: -------------------------------------------------------------------------------- 1 | """ Image to Patch Embedding using Conv2d 2 | 3 | A convolution based approach to patchifying a 2D image w/ embedding projection. 4 | 5 | Based on the impl in https://github.com/google-research/vision_transformer 6 | 7 | Hacked together by / Copyright 2020 Ross Wightman 8 | """ 9 | 10 | from torch import nn as nn 11 | 12 | from .helpers import to_2tuple 13 | 14 | 15 | class PatchEmbed(nn.Module): 16 | """2D Image to Patch Embedding""" 17 | 18 | def __init__( 19 | self, 20 | img_size=224, 21 | patch_size=16, 22 | in_chans=3, 23 | embed_dim=768, 24 | norm_layer=None, 25 | flatten=True, 26 | ): 27 | super().__init__() 28 | img_size = to_2tuple(img_size) 29 | patch_size = to_2tuple(patch_size) 30 | self.img_size = img_size 31 | self.patch_size = patch_size 32 | self.grid_size = (img_size[0] // patch_size[0], img_size[1] // patch_size[1]) 33 | self.num_patches = self.grid_size[0] * self.grid_size[1] 34 | self.flatten = flatten 35 | 36 | self.proj = nn.Conv2d( 37 | in_chans, embed_dim, kernel_size=patch_size, stride=patch_size 38 | ) 39 | self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity() 40 | 41 | def forward(self, x): 42 | B, C, H, W = x.shape 43 | assert ( 44 | H == self.img_size[0] and W == self.img_size[1] 45 | ), f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})." 46 | x = self.proj(x) 47 | if self.flatten: 48 | x = x.flatten(2).transpose(1, 2) # BCHW -> BNC 49 | x = self.norm(x) 50 | return x 51 | -------------------------------------------------------------------------------- /ptlflow/models/flownet/LICENSE: -------------------------------------------------------------------------------- 1 | Copyright 2017 NVIDIA CORPORATION 2 | 3 | Licensed under the Apache License, Version 2.0 (the "License"); 4 | you may not use this file except in compliance with the License. 5 | You may obtain a copy of the License at 6 | 7 | http://www.apache.org/licenses/LICENSE-2.0 8 | 9 | Unless required by applicable law or agreed to in writing, software 10 | distributed under the License is distributed on an "AS IS" BASIS, 11 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | See the License for the specific language governing permissions and 13 | limitations under the License. -------------------------------------------------------------------------------- /ptlflow/models/flownet/README.md: -------------------------------------------------------------------------------- 1 | # Flownet 2 | 3 | ## Original code 4 | 5 | [https://github.com/NVIDIA/flownet2-pytorch](https://github.com/NVIDIA/flownet2-pytorch) 6 | 7 | ## Code license 8 | 9 | See [LICENSE](LICENSE). 10 | 11 | ## Pretrained weights license 12 | 13 | See [WEIGHTS_LICENSE](WEIGHTS_LICENSE). 14 | 15 | ## Citation 16 | 17 | ``` 18 | @InProceedings{IMKDB17, 19 | author = "E. Ilg and N. Mayer and T. Saikia and M. Keuper and A. Dosovitskiy and T. Brox", 20 | title = "FlowNet 2.0: Evolution of Optical Flow Estimation with Deep Networks", 21 | booktitle = "IEEE Conference on Computer Vision and Pattern Recognition (CVPR)", 22 | month = "Jul", 23 | year = "2017", 24 | url = "http://lmb.informatik.uni-freiburg.de//Publications/2017/IMKDB17" 25 | } 26 | ``` 27 | 28 | ``` 29 | @misc{flownet2-pytorch, 30 | author = {Fitsum Reda and Robert Pottorff and Jon Barker and Bryan Catanzaro}, 31 | title = {flownet2-pytorch: Pytorch implementation of FlowNet 2.0: Evolution of Optical Flow Estimation with Deep Networks}, 32 | year = {2017}, 33 | publisher = {GitHub}, 34 | journal = {GitHub repository}, 35 | howpublished = {\url{https://github.com/NVIDIA/flownet2-pytorch}} 36 | } 37 | ``` -------------------------------------------------------------------------------- /ptlflow/models/flownet/WEIGHTS_LICENSE: -------------------------------------------------------------------------------- 1 | Pre-trained wieghts are provided for research purposes only and without any warranty. 2 | 3 | Any commercial use of the pre-trained wieghts requires FlowNet2 authors consent. 4 | When using the the pre-trained weights in your research work, please cite the following paper: 5 | 6 | @InProceedings{IMKDB17, 7 | author = "E. Ilg and N. Mayer and T. Saikia and M. Keuper and A. Dosovitskiy and T. Brox", 8 | title = "FlowNet 2.0: Evolution of Optical Flow Estimation with Deep Networks", 9 | booktitle = "IEEE Conference on Computer Vision and Pattern Recognition (CVPR)", 10 | month = "Jul", 11 | year = "2017", 12 | url = "http://lmb.informatik.uni-freiburg.de//Publications/2017/IMKDB17" 13 | } -------------------------------------------------------------------------------- /ptlflow/models/flownet/__init__.py: -------------------------------------------------------------------------------- 1 | from .flownet2 import * 2 | from .flownetc import * 3 | from .flownetcs import * 4 | from .flownetcss import * 5 | from .flownets import * 6 | from .flownetsd import * 7 | -------------------------------------------------------------------------------- /ptlflow/models/flownet/flownet_base.py: -------------------------------------------------------------------------------- 1 | from ..base_model.base_model import BaseModel 2 | from .losses import MultiScale 3 | 4 | 5 | class FlowNetBase(BaseModel): 6 | def __init__( 7 | self, 8 | div_flow: float = 20.0, 9 | input_channels: int = 6, 10 | batch_norm: bool = False, 11 | loss_start_scale: int = 4, 12 | loss_num_scales: int = 5, 13 | loss_base_weight: float = 0.32, 14 | loss_norm: str = "L2", 15 | **kwargs, 16 | ): 17 | super(FlowNetBase, self).__init__( 18 | loss_fn=MultiScale( 19 | startScale=loss_start_scale, 20 | numScales=loss_num_scales, 21 | l_weight=loss_base_weight, 22 | norm=loss_norm, 23 | ), 24 | output_stride=64, 25 | **kwargs, 26 | ) 27 | self.div_flow = div_flow 28 | self.input_channels = input_channels 29 | self.batch_norm = batch_norm 30 | -------------------------------------------------------------------------------- /ptlflow/models/gma/LICENSE: -------------------------------------------------------------------------------- 1 | 2 | DO WHAT THE FUCK YOU WANT TO PUBLIC LICENSE 3 | Version 2, December 2004 4 | 5 | Copyright (C) 2004 Sam Hocevar 6 | 7 | Everyone is permitted to copy and distribute verbatim or modified 8 | copies of this license document, and changing it is allowed as long 9 | as the name is changed. 10 | 11 | DO WHAT THE FUCK YOU WANT TO PUBLIC LICENSE 12 | TERMS AND CONDITIONS FOR COPYING, DISTRIBUTION AND MODIFICATION 13 | 14 | 0. You just DO WHAT THE FUCK YOU WANT TO. -------------------------------------------------------------------------------- /ptlflow/models/gma/README.md: -------------------------------------------------------------------------------- 1 | # GMA 2 | 3 | ## Original code 4 | 5 | [https://github.com/zacjiang/GMA](https://github.com/zacjiang/GMA) 6 | 7 | ## Code license 8 | 9 | See [LICENSE](LICENSE). 10 | 11 | ## Pretrained weights license 12 | 13 | Not specified. 14 | 15 | ## Citation 16 | 17 | ``` 18 | @inproceedings{jiang2021learning, 19 | title={Learning to Estimate Hidden Motions with Global Motion Aggregation}, 20 | author={Jiang, Shihao and Campbell, Dylan and Lu, Yao and Li, Hongdong and Hartley, Richard}, 21 | booktitle={ICCV}, 22 | year={2021} 23 | } 24 | ``` -------------------------------------------------------------------------------- /ptlflow/models/gma/__init__.py: -------------------------------------------------------------------------------- 1 | from .gma import * 2 | -------------------------------------------------------------------------------- /ptlflow/models/gmflow/README.md: -------------------------------------------------------------------------------- 1 | # GMFlow 2 | 3 | ## Original code 4 | 5 | [https://github.com/haofeixu/gmflow](https://github.com/haofeixu/gmflow) 6 | 7 | ## Code license 8 | 9 | See [LICENSE](LICENSE). 10 | 11 | ## Pretrained weights license 12 | 13 | Not specified. 14 | 15 | ## Citation 16 | 17 | ``` 18 | @inproceedings{xu2022gmflow, 19 | title={GMFlow: Learning Optical Flow via Global Matching}, 20 | author={Xu, Haofei and Zhang, Jing and Cai, Jianfei and Rezatofighi, Hamid and Tao, Dacheng}, 21 | booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition}, 22 | pages={8121-8130}, 23 | year={2022} 24 | } 25 | ``` -------------------------------------------------------------------------------- /ptlflow/models/gmflow/__init__.py: -------------------------------------------------------------------------------- 1 | from .gmflow import * 2 | -------------------------------------------------------------------------------- /ptlflow/models/gmflownet/LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2022 Shiyu Zhao 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. -------------------------------------------------------------------------------- /ptlflow/models/gmflownet/README.md: -------------------------------------------------------------------------------- 1 | # GMFlowNet 2 | 3 | ## Original code 4 | 5 | [https://github.com/xiaofeng94/GMFlowNet](https://github.com/xiaofeng94/GMFlowNet) 6 | 7 | ## Code license 8 | 9 | See [LICENSE](LICENSE). 10 | 11 | ## Pretrained weights license 12 | 13 | Not specified. 14 | 15 | ## Citation 16 | 17 | ``` 18 | @inproceedings{xu2022gmflow, 19 | title={Global Matching with Overlapping Attention for Optical Flow Estimation}, 20 | author={Zhao, Shiyu and Zhao, Long and Zhang, Zhixing and Zhou, Enyu and Metaxas, Dimitris}, 21 | booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition}, 22 | year={2022} 23 | } 24 | ``` -------------------------------------------------------------------------------- /ptlflow/models/gmflownet/__init__.py: -------------------------------------------------------------------------------- 1 | from .gmflownet import * 2 | -------------------------------------------------------------------------------- /ptlflow/models/gmflownet/utils/__init__.py: -------------------------------------------------------------------------------- 1 | # functions from timm 2 | from .drop import DropBlock2d, DropPath, drop_block_2d, drop_path 3 | from .helpers import to_ntuple, to_2tuple, to_3tuple, to_4tuple, make_divisible 4 | from .weight_init import trunc_normal_, variance_scaling_, lecun_normal_ 5 | -------------------------------------------------------------------------------- /ptlflow/models/gmflownet/utils/helpers.py: -------------------------------------------------------------------------------- 1 | """ Layer/Module Helpers 2 | 3 | Hacked together by / Copyright 2020 Ross Wightman 4 | """ 5 | 6 | from itertools import repeat 7 | import collections.abc 8 | 9 | 10 | # From PyTorch internals 11 | def _ntuple(n): 12 | def parse(x): 13 | if isinstance(x, collections.abc.Iterable): 14 | return x 15 | return tuple(repeat(x, n)) 16 | 17 | return parse 18 | 19 | 20 | to_1tuple = _ntuple(1) 21 | to_2tuple = _ntuple(2) 22 | to_3tuple = _ntuple(3) 23 | to_4tuple = _ntuple(4) 24 | to_ntuple = _ntuple 25 | 26 | 27 | def make_divisible(v, divisor=8, min_value=None, round_limit=0.9): 28 | min_value = min_value or divisor 29 | new_v = max(min_value, int(v + divisor / 2) // divisor * divisor) 30 | # Make sure that round down does not go down by more than 10%. 31 | if new_v < round_limit * v: 32 | new_v += divisor 33 | return new_v 34 | -------------------------------------------------------------------------------- /ptlflow/models/hd3/LICENSE: -------------------------------------------------------------------------------- 1 | BSD 3-Clause License 2 | 3 | Copyright (c) 2019, Berkeley DeepDrive 4 | All rights reserved. 5 | 6 | Redistribution and use in source and binary forms, with or without 7 | modification, are permitted provided that the following conditions are met: 8 | 9 | 1. Redistributions of source code must retain the above copyright notice, this 10 | list of conditions and the following disclaimer. 11 | 12 | 2. Redistributions in binary form must reproduce the above copyright notice, 13 | this list of conditions and the following disclaimer in the documentation 14 | and/or other materials provided with the distribution. 15 | 16 | 3. Neither the name of the copyright holder nor the names of its 17 | contributors may be used to endorse or promote products derived from 18 | this software without specific prior written permission. 19 | 20 | THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" 21 | AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 22 | IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE 23 | DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE 24 | FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL 25 | DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR 26 | SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER 27 | CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, 28 | OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 29 | OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. -------------------------------------------------------------------------------- /ptlflow/models/hd3/README.md: -------------------------------------------------------------------------------- 1 | # HD3 2 | 3 | ## Original code 4 | 5 | [https://github.com/ucbdrive/hd3](https://github.com/ucbdrive/hd3) 6 | 7 | ## Code license 8 | 9 | See [LICENSE](LICENSE). 10 | 11 | ## Pretrained weights license 12 | 13 | Not specified. 14 | 15 | ## Citation 16 | 17 | ``` 18 | @InProceedings{Yin_2019_CVPR, 19 | author = {Yin, Zhichao and Darrell, Trevor and Yu, Fisher}, 20 | title = {Hierarchical Discrete Distribution Decomposition for Match Density Estimation}, 21 | booktitle = {The IEEE Conference on Computer Vision and Pattern Recognition (CVPR)}, 22 | month = {June}, 23 | year = {2019} 24 | } 25 | ``` -------------------------------------------------------------------------------- /ptlflow/models/hd3/__init__.py: -------------------------------------------------------------------------------- 1 | from .hd3 import * 2 | -------------------------------------------------------------------------------- /ptlflow/models/irr/README.md: -------------------------------------------------------------------------------- 1 | # IRR 2 | 3 | ## Original code 4 | 5 | [https://github.com/visinf/irr](https://github.com/visinf/irr) 6 | 7 | ## Code license 8 | 9 | See [LICENSE](LICENSE). 10 | 11 | ## Pretrained weights license 12 | 13 | Not specified. 14 | 15 | ## Citation 16 | 17 | ``` 18 | @inproceedings{Hur:2019:IRR, 19 | Author = {Junhwa Hur and Stefan Roth}, 20 | Booktitle = {CVPR}, 21 | Title = {Iterative Residual Refinement for Joint Optical Flow and Occlusion Estimation}, 22 | Year = {2019} 23 | } 24 | ``` -------------------------------------------------------------------------------- /ptlflow/models/irr/__init__.py: -------------------------------------------------------------------------------- 1 | from .irr_pwc import * 2 | from .pwcnet import * 3 | from .pwcnet_irr import * 4 | -------------------------------------------------------------------------------- /ptlflow/models/lcv/LICENSE_LCV: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2021 Taihong Xiao 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. -------------------------------------------------------------------------------- /ptlflow/models/lcv/LICENSE_RAFT: -------------------------------------------------------------------------------- 1 | BSD 3-Clause License 2 | 3 | Copyright (c) 2020, princeton-vl 4 | All rights reserved. 5 | 6 | Redistribution and use in source and binary forms, with or without 7 | modification, are permitted provided that the following conditions are met: 8 | 9 | * Redistributions of source code must retain the above copyright notice, this 10 | list of conditions and the following disclaimer. 11 | 12 | * Redistributions in binary form must reproduce the above copyright notice, 13 | this list of conditions and the following disclaimer in the documentation 14 | and/or other materials provided with the distribution. 15 | 16 | * Neither the name of the copyright holder nor the names of its 17 | contributors may be used to endorse or promote products derived from 18 | this software without specific prior written permission. 19 | 20 | THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" 21 | AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 22 | IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE 23 | DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE 24 | FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL 25 | DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR 26 | SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER 27 | CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, 28 | OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 29 | OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. -------------------------------------------------------------------------------- /ptlflow/models/lcv/README.md: -------------------------------------------------------------------------------- 1 | # LCV + RAFT 2 | 3 | ## Original code 4 | 5 | - LCV: [https://github.com/Prinsphield/LCV](https://github.com/Prinsphield/LCV) 6 | - RAFT: [https://github.com/princeton-vl/RAFT](https://github.com/princeton-vl/RAFT) 7 | 8 | ## Code license 9 | 10 | - LCV: [LICENSE_LCV](LICENSE_LCV). 11 | - RAFT: [LICENSE_RAFT](LICENSE_RAFT). 12 | 13 | ## Pretrained weights license 14 | 15 | See [LICENSE_WEIGHTS](LICENSE_WEIGHTS). 16 | 17 | ## Citation 18 | 19 | ### LCV 20 | 21 | ``` 22 | @inproceedings{xiao2020learnable, 23 | title={Learnable cost volume using the cayley representation}, 24 | author={Xiao, Taihong and Yuan, Jinwei and Sun, Deqing and Wang, Qifei and Zhang, Xin-Yu and Xu, Kehan and Yang, Ming-Hsuan}, 25 | booktitle={European Conference on Computer Vision (ECCV)}, 26 | pages={483--499}, 27 | year={2020}, 28 | organization={Springer} 29 | } 30 | ``` 31 | 32 | ### RAFT 33 | 34 | ``` 35 | @inproceedings{teed2020raft, 36 | title={Raft: Recurrent all-pairs field transforms for optical flow}, 37 | author={Teed, Zachary and Deng, Jia}, 38 | booktitle={European Conference on Computer Vision}, 39 | pages={402--419}, 40 | year={2020}, 41 | organization={Springer} 42 | } 43 | ``` -------------------------------------------------------------------------------- /ptlflow/models/lcv/__init__.py: -------------------------------------------------------------------------------- 1 | from .lcv_raft import * 2 | -------------------------------------------------------------------------------- /ptlflow/models/liteflownet/LICENSE_LFN1: -------------------------------------------------------------------------------- 1 | Copyright (c) 2018 Tak-Wai Hui 2 | 3 | This software and associated documentation files (the "Software"), and the research paper 4 | (LiteFlowNet: A Lightweight Convolutional Neural Network for Optical Flow Estimation) including 5 | but not limited to the figures, and tables (the "Paper") are provided for academic research purposes 6 | only and without any warranty. Any commercial use requires my consent. When using any parts 7 | of the Software or the Paper in your work, please cite the following paper: 8 | 9 | @InProceedings{hui18liteflownet, 10 | author = {Tak-Wai Hui and Xiaoou Tang and Chen Change Loy}, 11 | title = {LiteFlowNet: A Lightweight Convolutional Neural Network for Optical Flow Estimation}, 12 | booktitle = {Proceedings of IEEE Conference on Computer Vision and Pattern Recognition (CVPR)}, 13 | year = {2018}, 14 | pages = {8981--8989}, 15 | } 16 | 17 | The above copyright notice and this permission notice shall be included in all copies or 18 | substantial portions of the Software. 19 | -------------------------------------------------------------------------------- /ptlflow/models/liteflownet/LICENSE_LFN2: -------------------------------------------------------------------------------- 1 | Copyright (c) 2019 Tak-Wai Hui 2 | 3 | This software and associated documentation files (the "Software"), and the research paper 4 | (A Lightweight Optical Flow CNN - Revisiting Data Fidelity and Regularization) including 5 | but not limited to the figures, and tables (the "Paper") are provided for academic research 6 | purposes only and without any warranty. Any commercial use requires my consent. When using 7 | any parts of the Software or the Paper in your work, please cite the following papers: 8 | 9 | @InProceedings{hui20liteflownet2, 10 | author = {Tak-Wai Hui and Xiaoou Tang and Chen Change Loy}, 11 | title = {A {L}ightweight {O}ptical {F}low {CNN} - {R}evisiting {D}ata {F}idelity and {R}egularization}, 12 | journal = {IEEE Transactions on Pattern Analysis and Machine Intelligence}, 13 | year = {2020} 14 | } 15 | 16 | @InProceedings{hui18liteflownet, 17 | author = {Tak-Wai Hui and Xiaoou Tang and Chen Change Loy}, 18 | title = {LiteFlowNet: A Lightweight Convolutional Neural Network for Optical Flow Estimation}, 19 | booktitle = {Proceedings of IEEE Conference on Computer Vision and Pattern Recognition (CVPR)}, 20 | year = {2018}, 21 | pages = {8981--8989} 22 | } 23 | 24 | The above copyright notice and this permission notice shall be included in all copies or 25 | substantial portions of the Software. -------------------------------------------------------------------------------- /ptlflow/models/liteflownet/LICENSE_LFN3: -------------------------------------------------------------------------------- 1 | Copyright (c) 2020 Tak-Wai Hui 2 | 3 | This software and associated documentation files (the "Software"), and the research paper 4 | (LiteFlowNet3: Resolving Correspondence Ambiguity for More Accurate Optical Flow Estimation) 5 | including but not limited to the figures, and tables (the "Paper") are provided for academic 6 | research purposes only and without any warranty. Any commercial use requires my consent. When 7 | using any parts of the Software or the Paper in your work, please cite the following papers: 8 | 9 | @InProceedings{hui20liteflownet3, 10 | author = {Tak-Wai Hui and Chen Change Loy}, 11 | title = {{LiteFlowNet3: Resolving Correspondence Ambiguity for More Accurate Optical Flow Estimation}}, 12 | journal = {{Proceedings of the European Conference on Computer Vision (ECCV)}}, 13 | pages = {169--184}, 14 | year = {2020}, 15 | } 16 | 17 | @InProceedings{hui20liteflownet2, 18 | author = {Tak-Wai Hui and Xiaoou Tang and Chen Change Loy}, 19 | title = {{A Lightweight Optical Flow CNN - Revisiting Data Fidelity and Regularization}, 20 | journal = {{IEEE Transactions on Pattern Analysis and Machine Intelligence}}, 21 | year = {2020}, 22 | url = {http://mmlab.ie.cuhk.edu.hk/projects/LiteFlowNet/} 23 | } 24 | 25 | @InProceedings{hui18liteflownet, 26 | author = {Tak-Wai Hui and Xiaoou Tang and Chen Change Loy}, 27 | title = {{LiteFlowNet: A Lightweight Convolutional Neural Network for Optical Flow Estimation}}, 28 | booktitle = {{Proceedings of IEEE Conference on Computer Vision and Pattern Recognition (CVPR)}}, 29 | year = {2018}, 30 | pages = {8981--8989}, 31 | url = {http://mmlab.ie.cuhk.edu.hk/projects/LiteFlowNet/} 32 | } 33 | 34 | The above copyright notice and this permission notice shall be included in all copies or 35 | substantial portions of the Software. -------------------------------------------------------------------------------- /ptlflow/models/liteflownet/README.md: -------------------------------------------------------------------------------- 1 | # LiteFlowNet 2 | 3 | ## Original code 4 | 5 | - LiteFlowNet: [https://github.com/twhui/LiteFlowNet](https://github.com/twhui/LiteFlowNet) 6 | - LiteFlowNet2: [https://github.com/twhui/LiteFlowNet2](https://github.com/twhui/LiteFlowNet2) 7 | - LiteFlowNet3: [https://github.com/twhui/LiteFlowNet3](https://github.com/twhui/LiteFlowNet3) 8 | 9 | ## Code licenses 10 | 11 | See 12 | 13 | - [LICENSE_LFN1](LICENSE_LFN1). 14 | - [LICENSE_LFN2](LICENSE_LFN2). 15 | - [LICENSE_LFN3](LICENSE_LFN3). 16 | 17 | The warping layer comes from IRR (https://github.com/visinf/irr). See [LICENSE_IRR](LICENSE_IRR). 18 | 19 | ## Pretrained weights license 20 | 21 | Not specified. 22 | 23 | ## Citation 24 | 25 | ### LiteFlowNet 26 | 27 | ``` 28 | @InProceedings{hui18liteflownet, 29 | author = {Tak-Wai Hui and Xiaoou Tang and Chen Change Loy}, 30 | title = {{LiteFlowNet: A Lightweight Convolutional Neural Network for Optical Flow Estimation}}, 31 | booktitle = {{Proceedings of IEEE Conference on Computer Vision and Pattern Recognition (CVPR)}}, 32 | year = {2018}, 33 | pages = {8981--8989}, 34 | url = {http://mmlab.ie.cuhk.edu.hk/projects/LiteFlowNet/} 35 | } 36 | ``` 37 | 38 | ### LiteFlowNet2 39 | 40 | ``` 41 | @InProceedings{hui20liteflownet2, 42 | author = {Tak-Wai Hui and Xiaoou Tang and Chen Change Loy}, 43 | title = {{A Lightweight Optical Flow CNN - Revisiting Data Fidelity and Regularization}}, 44 | journal = {{IEEE Transactions on Pattern Analysis and Machine Intelligence}}, 45 | year = {2020}, 46 | url = {http://mmlab.ie.cuhk.edu.hk/projects/LiteFlowNet/} 47 | } 48 | ``` 49 | 50 | ### LiteFlowNet3 51 | 52 | ``` 53 | @InProceedings{hui20liteflownet3,     54 | author = {Tak-Wai Hui and Chen Change Loy},     55 | title = {{LiteFlowNet3: Resolving Correspondence Ambiguity for More Accurate Optical Flow Estimation}},  56 | journal = {{Proceedings of the European Conference on Computer Vision (ECCV)}}, 57 | year = {2020},     58 | pages = {169--184}, 59 | } 60 | ``` 61 | -------------------------------------------------------------------------------- /ptlflow/models/liteflownet/__init__.py: -------------------------------------------------------------------------------- 1 | from .liteflownet import * 2 | from .liteflownet2 import * 3 | from .liteflownet3 import * 4 | -------------------------------------------------------------------------------- /ptlflow/models/liteflownet/warp.py: -------------------------------------------------------------------------------- 1 | # Code taken from IRR: https://github.com/visinf/irr 2 | # Licensed under the Apache 2.0 license (see LICENSE_IRR). 3 | 4 | import torch 5 | import torch.nn as nn 6 | import torch.nn.functional as F 7 | 8 | 9 | def get_grid(x): 10 | grid_H = ( 11 | torch.linspace(-1.0, 1.0, x.size(3)) 12 | .view(1, 1, 1, x.size(3)) 13 | .expand(x.size(0), 1, x.size(2), x.size(3)) 14 | ) 15 | grid_V = ( 16 | torch.linspace(-1.0, 1.0, x.size(2)) 17 | .view(1, 1, x.size(2), 1) 18 | .expand(x.size(0), 1, x.size(2), x.size(3)) 19 | ) 20 | grid = torch.cat([grid_H, grid_V], 1) 21 | grids_cuda = grid.requires_grad_(False).to(dtype=x.dtype, device=x.device) 22 | return grids_cuda 23 | 24 | 25 | class WarpingLayer(nn.Module): 26 | def __init__(self): 27 | super(WarpingLayer, self).__init__() 28 | 29 | def forward(self, x, flow, height_im, width_im, div_flow): 30 | flo_list = [] 31 | flo_w = flow[:, 0] * 2 / max(width_im - 1, 1) / div_flow 32 | flo_h = flow[:, 1] * 2 / max(height_im - 1, 1) / div_flow 33 | flo_list.append(flo_w) 34 | flo_list.append(flo_h) 35 | flow_for_grid = torch.stack(flo_list).transpose(0, 1) 36 | grid = torch.add(get_grid(x), flow_for_grid).transpose(1, 2).transpose(2, 3) 37 | x_warp = F.grid_sample(x, grid, align_corners=True) 38 | 39 | mask = torch.ones(x.size(), requires_grad=False).to( 40 | dtype=x.dtype, device=x.device 41 | ) 42 | mask = F.grid_sample(mask, grid, align_corners=True) 43 | mask = (mask >= 1.0).to(dtype=x.dtype) 44 | 45 | return x_warp * mask 46 | -------------------------------------------------------------------------------- /ptlflow/models/llaflow/LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2023 mansang127 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. -------------------------------------------------------------------------------- /ptlflow/models/llaflow/README.md: -------------------------------------------------------------------------------- 1 | # LLA-Flow 2 | 3 | ## Original code 4 | 5 | [https://github.com/mansang127/LLA-Flow](https://github.com/mansang127/LLA-Flow) 6 | 7 | ## Code license 8 | 9 | See [LICENSE](LICENSE). 10 | 11 | ## Pretrained weights license 12 | 13 | Not specified. 14 | 15 | ## Citation 16 | 17 | ``` 18 | @inproceedings{Xu2023LLAFlowAL, 19 | title={LLA-Flow: A Lightweight Local Aggregation on Cost Volume for Optical Flow Estimation}, 20 | author={Jiawei Xu and Zongqing Lu and Qingmin Liao}, 21 | booktitle={IEEE International Conference on Image Processing (ICIP)}, 22 | year={2023}, 23 | pages={3220-3224}, 24 | url={https://api.semanticscholar.org/CorpusID:258179047} 25 | } 26 | ``` -------------------------------------------------------------------------------- /ptlflow/models/llaflow/__init__.py: -------------------------------------------------------------------------------- 1 | from .llaflow import * 2 | -------------------------------------------------------------------------------- /ptlflow/models/maskflownet/LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2020 Daniele Cattaneo 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. -------------------------------------------------------------------------------- /ptlflow/models/maskflownet/README.md: -------------------------------------------------------------------------------- 1 | # MaskFlownet 2 | 3 | ## Original code 4 | 5 | [https://github.com/cattaneod/MaskFlownet-Pytorch](https://github.com/cattaneod/MaskFlownet-Pytorch) 6 | 7 | ## Official implementation (MXNet) 8 | 9 | [https://github.com/microsoft/MaskFlownet](https://github.com/microsoft/MaskFlownet) 10 | 11 | ## Code license 12 | 13 | See [LICENSE](LICENSE). 14 | 15 | ## Pretrained weights license 16 | 17 | Not specified. 18 | 19 | ## Citation 20 | 21 | ``` 22 | @inproceedings{zhao2020maskflownet, 23 | author = {Zhao, Shengyu and Sheng, Yilun and Dong, Yue and Chang, Eric I-Chao and Xu, Yan}, 24 | title = {MaskFlownet: Asymmetric Feature Matching with Learnable Occlusion Mask}, 25 | booktitle = {Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR)}, 26 | year = {2020} 27 | } 28 | ``` -------------------------------------------------------------------------------- /ptlflow/models/maskflownet/__init__.py: -------------------------------------------------------------------------------- 1 | from .maskflownet import * 2 | -------------------------------------------------------------------------------- /ptlflow/models/matchflow/QuadtreeAttention/QuadtreeAttention/src/score_computation.cpp: -------------------------------------------------------------------------------- 1 | #include "score_computation.h" 2 | #include 3 | #include 4 | #include 5 | #include 6 | #define CHECK_CUDA(x) TORCH_CHECK(x.type().is_cuda(), #x " must be a CUDA tensor") 7 | #define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous") 8 | #define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x) 9 | 10 | // == Forward 11 | std::vector score_cuda_forward(torch::Tensor input1, //parameter: K*group_num, C 12 | torch::Tensor input2, //tensor : B, N, C 13 | torch::Tensor index) //tensor: B, N, K 14 | { 15 | CHECK_INPUT(input1); 16 | CHECK_INPUT(input2); 17 | CHECK_INPUT(index); 18 | return ScoreData_ongpu(input1, input2, index); 19 | 20 | } 21 | 22 | std::vector score_cuda_backward(torch::Tensor grad_output1, //B,N,C,group_num 23 | torch::Tensor input1, //scene : N, H, W, C1 24 | torch::Tensor input2, // scene coords: N, H, W, 3 25 | torch::Tensor index) //tensor: B, N, K 26 | { 27 | CHECK_INPUT(grad_output1); 28 | CHECK_INPUT(input1); 29 | CHECK_INPUT(input2); 30 | CHECK_INPUT(index); 31 | return ScoreData_backward_ongpu(grad_output1, input1, input2, index); 32 | } 33 | 34 | 35 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 36 | m.def("score_forward", &score_cuda_forward, "score forward (CUDA)"); 37 | m.def("score_backward", &score_cuda_backward, "score forward (CUDA)"); 38 | } -------------------------------------------------------------------------------- /ptlflow/models/matchflow/QuadtreeAttention/QuadtreeAttention/src/score_computation.h: -------------------------------------------------------------------------------- 1 | #ifndef _Score_CUDA 2 | #define _Score_CUDA 3 | #include 4 | #include 5 | 6 | std::vector score_cuda_forward(torch::Tensor input1, //query t: N, H, W, C1 7 | torch::Tensor input2, //scene : N, H, W, C1 8 | torch::Tensor index); //scene : N, H, W, C1 9 | 10 | 11 | 12 | std::vector ScoreData_ongpu(at::Tensor input1, //query t: N, H, W, C1 13 | at::Tensor input2, //scene : N, H, W, C1 14 | at::Tensor index); //scene : N, H, W, C1 15 | 16 | 17 | std::vector ScoreData_backward_ongpu(torch::Tensor grad_output1, //B,N,C,group_num 18 | torch::Tensor input1, //scene : N, H, W, C1 19 | torch::Tensor input2, // scene coords: N, H, W, 3 20 | torch::Tensor index); //tensor: B, N, K 21 | 22 | #endif -------------------------------------------------------------------------------- /ptlflow/models/matchflow/QuadtreeAttention/QuadtreeAttention/src/utils.h: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | 5 | class Formatter { 6 | public: 7 | Formatter() {} 8 | ~Formatter() {} 9 | 10 | template Formatter &operator<<(const Type &value) { 11 | stream_ << value; 12 | return *this; 13 | } 14 | 15 | std::string str() const { return stream_.str(); } 16 | operator std::string() const { return stream_.str(); } 17 | 18 | enum ConvertToString { to_str }; 19 | 20 | std::string operator>>(ConvertToString) { return stream_.str(); } 21 | 22 | private: 23 | std::stringstream stream_; 24 | Formatter(const Formatter &); 25 | Formatter &operator=(Formatter &); 26 | }; 27 | 28 | -------------------------------------------------------------------------------- /ptlflow/models/matchflow/QuadtreeAttention/QuadtreeAttention/src/value_aggregation.h: -------------------------------------------------------------------------------- 1 | #ifndef _VALUE_AGGREGATION_ 2 | #define _VALUE_AGGREGATION_ 3 | #include 4 | #include 5 | 6 | void value_aggregation_forward_kernel(float* score, // B, N, K, H 7 | float* value, // B, M, H, D 8 | long* index, // B, N, K, H 9 | float* output, // B, N, H, D 10 | int B, int N, int K, int H, int M, int D, cudaStream_t stream 11 | ); 12 | 13 | void value_aggregation_cuda_forward(at::Tensor score, at::Tensor value, at::Tensor index, at::Tensor output); 14 | 15 | void value_aggregation_backward_kernel(float* grad_output, float* score, float* value,long* index, float* grad_score, float* grad_value, int B, int N, int K, int H, int M, int D, cudaStream_t stream); 16 | 17 | #endif // _VALUE_AGGREGATION_ 18 | -------------------------------------------------------------------------------- /ptlflow/models/matchflow/QuadtreeAttention/setup.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | 3 | import os 4 | 5 | from setuptools import setup, find_packages 6 | from torch.utils.cpp_extension import BuildExtension, CUDAExtension 7 | 8 | this_file = os.path.dirname(__file__) 9 | 10 | setup( 11 | name="quadtree_attention_package", 12 | ext_modules=[ 13 | CUDAExtension( 14 | "score_computation_cuda", 15 | [ 16 | "QuadtreeAttention/src/score_computation.cpp", 17 | "QuadtreeAttention/src/score_computation_kernal.cu", 18 | ], 19 | extra_compile_args={"cxx": ["-g"], "nvcc": ["-O2"]}, 20 | ), 21 | CUDAExtension( 22 | "value_aggregation_cuda", 23 | [ 24 | "QuadtreeAttention/src/value_aggregation.cpp", 25 | "QuadtreeAttention/src/value_aggregation_kernel.cu", 26 | ], 27 | extra_compile_args={"cxx": ["-g"], "nvcc": ["-O2"]}, 28 | ), 29 | ], 30 | cmdclass={"build_ext": BuildExtension}, 31 | ) 32 | -------------------------------------------------------------------------------- /ptlflow/models/matchflow/README.md: -------------------------------------------------------------------------------- 1 | # MatchFlow 2 | 3 | ## Original code 4 | 5 | [https://github.com/DQiaole/MatchFlow](https://github.com/DQiaole/MatchFlow) 6 | 7 | ## Additional requirements 8 | 9 | In order to use MatchFlow you need to have CUDA installed and compile the QuadTreeAttention package: 10 | 11 | - Download and install CUDA from [https://developer.nvidia.com/cuda-downloads](https://developer.nvidia.com/cuda-downloads) 12 | - IMPORTANT! Be sure to choose the same CUDA version as your PyTorch 13 | - Enter the QuadTreeAttention folder and run the setup: 14 | ```bash 15 | cd QuadTreeAttention 16 | python setup.py install 17 | ``` 18 | 19 | ## Code license 20 | 21 | See [LICENSE](LICENSE). 22 | 23 | ## Pretrained weights license 24 | 25 | Not specified. 26 | 27 | ## Citation 28 | 29 | ``` 30 | @inproceedings{dong2023rethinking, 31 | title={Rethinking Optical Flow from Geometric Matching Consistent Perspective}, 32 | author={Dong, Qiaole and Cao, Chenjie and Fu, Yanwei}, 33 | booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition}, 34 | year={2023} 35 | } 36 | ``` -------------------------------------------------------------------------------- /ptlflow/models/matchflow/__init__.py: -------------------------------------------------------------------------------- 1 | from .matchflow import * 2 | -------------------------------------------------------------------------------- /ptlflow/models/matchflow/common.py: -------------------------------------------------------------------------------- 1 | def torch_init_model(model, total_dict, key, rank=0): 2 | if key in total_dict: 3 | state_dict = total_dict[key] 4 | else: 5 | state_dict = total_dict 6 | missing_keys = [] 7 | unexpected_keys = [] 8 | error_msgs = [] 9 | # copy state_dict so _load_from_state_dict can modify it 10 | metadata = getattr(state_dict, "_metadata", None) 11 | state_dict = state_dict.copy() 12 | if metadata is not None: 13 | state_dict._metadata = metadata 14 | 15 | def load(module, prefix=""): 16 | local_metadata = {} if metadata is None else metadata.get(prefix[:-1], {}) 17 | module._load_from_state_dict( 18 | state_dict=state_dict, 19 | prefix=prefix, 20 | local_metadata=local_metadata, 21 | strict=True, 22 | missing_keys=missing_keys, 23 | unexpected_keys=unexpected_keys, 24 | error_msgs=error_msgs, 25 | ) 26 | for name, child in module._modules.items(): 27 | if child is not None: 28 | load(child, prefix + name + ".") 29 | 30 | load(model, prefix="") 31 | 32 | if rank == 0: 33 | print("missing keys:{}".format(missing_keys)) 34 | print("unexpected keys:{}".format(unexpected_keys)) 35 | print("error msgs:{}".format(error_msgs)) 36 | -------------------------------------------------------------------------------- /ptlflow/models/memflow/MemFlowNet/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hmorimitsu/ptlflow/fb6d753c3b16e74481258891b45668dcba63e54e/ptlflow/models/memflow/MemFlowNet/__init__.py -------------------------------------------------------------------------------- /ptlflow/models/memflow/README.md: -------------------------------------------------------------------------------- 1 | # MemFlow 2 | 3 | ## Original code 4 | 5 | [https://github.com/DQiaole/MemFlow](https://github.com/DQiaole/MemFlow) 6 | 7 | ## Code license 8 | 9 | See [LICENSE](LICENSE). 10 | 11 | ## Pretrained weights license 12 | 13 | Not specified. 14 | 15 | ## Citation 16 | 17 | ``` 18 | @inproceedings{dong2024memflow, 19 | title={MemFlow: Optical Flow Estimation and Prediction with Memory}, 20 | author={Dong, Qiaole and Fu, Yanwei}, 21 | booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition}, 22 | year={2024} 23 | } 24 | ``` -------------------------------------------------------------------------------- /ptlflow/models/memflow/__init__.py: -------------------------------------------------------------------------------- 1 | from .memflow import * 2 | -------------------------------------------------------------------------------- /ptlflow/models/ms_raft_plus/LICENSE: -------------------------------------------------------------------------------- 1 | Copyright (c) 2023 Azin Jahedi, Maximilian Luz. All Rights Reserved. 2 | 3 | Redistribution and use in source and binary forms, with or without modification, 4 | are permitted provided that the following conditions are met: 5 | 6 | 1. Redistribution of source code must retain the above copyright notice, this 7 | list of conditions and the following disclaimer. 8 | 9 | 2. Redistribution in binary form must reproduce the above copyright notice, 10 | this list of conditions and the following disclaimer in the documentation 11 | and/or other materials provided with the distribution. 12 | 13 | 3. Neither the name of the copyright holder nor the names of its contributors 14 | may be used to endorse or promote products derived from this software without 15 | specific prior written permission. 16 | 17 | THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND 18 | ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED 19 | WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. 20 | IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, 21 | INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, 22 | BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, 23 | DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF 24 | LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE 25 | OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED 26 | OF THE POSSIBILITY OF SUCH DAMAGE. 27 | 28 | You acknowledge that this software is not designed, licensed or intended for use 29 | in the design, construction, operation or maintenance of any military facility. -------------------------------------------------------------------------------- /ptlflow/models/ms_raft_plus/README.md: -------------------------------------------------------------------------------- 1 | # MS-RAFT+ 2 | 3 | ## Original code 4 | 5 | [https://github.com/cv-stuttgart/MS_RAFT_plus](https://github.com/cv-stuttgart/MS_RAFT_plus) 6 | 7 | ## Additional requirements 8 | 9 | In order to use MS-RAFT+ you need to install the alt_cuda_corr package. 10 | Follow the instructions at [ptlflow/utils/external/alt_cuda_corr/README.md](../../../ptlflow/utils/external/alt_cuda_corr/README.md). 11 | 12 | ## Code license 13 | 14 | See [LICENSE](LICENSE). 15 | 16 | ## Pretrained weights license 17 | 18 | [CC BY-NC-SA 3.0 license.](https://creativecommons.org/licenses/by-nc-sa/3.0/) 19 | 20 | ## Citation 21 | 22 | ``` 23 | @techreport{jahediHighResolutionMultiScaleRAFT2022, 24 | title = {High-Resolution Multi-Scale {{RAFT}}}, 25 | author = {Jahedi, Azin and Luz, Maximilian and Mehl, Lukas and Rivinius, Marc and Bruhn, Andr{\'e}s}, 26 | year = {2022}, 27 | month = oct, 28 | eprint = {2210.16900}, 29 | eprinttype = {arxiv}, 30 | primaryclass = {cs}, 31 | pages = {3}, 32 | archiveprefix = {arXiv}, 33 | langid = {english}, 34 | keywords = {Computer Science - Computer Vision and Pattern Recognition}, 35 | } 36 | ``` -------------------------------------------------------------------------------- /ptlflow/models/ms_raft_plus/__init__.py: -------------------------------------------------------------------------------- 1 | from .ms_raft_plus import * 2 | -------------------------------------------------------------------------------- /ptlflow/models/neuflow/README.md: -------------------------------------------------------------------------------- 1 | # NeuFlow 2 | 3 | ## Original code 4 | 5 | [https://github.com/neufieldrobotics/NeuFlow](https://github.com/neufieldrobotics/NeuFlow) 6 | 7 | ## Code license 8 | 9 | See [LICENSE](LICENSE). 10 | 11 | ## Pretrained weights license 12 | 13 | Not specified. 14 | 15 | ## Citation 16 | 17 | ``` 18 | @Article{Zhang2024NeuFlowRealtime, 19 | author = {Zhiyong Zhang and Huaizu Jiang and Hanumant Singh}, 20 | journal = {ArXiv}, 21 | title = {{NeuFlow}: Real-time, High-accuracy Optical Flow Estimation on Robots Using Edge Devices}, 22 | year = {2024}, 23 | } 24 | 25 | ``` -------------------------------------------------------------------------------- /ptlflow/models/neuflow/__init__.py: -------------------------------------------------------------------------------- 1 | from .neuflow import * 2 | -------------------------------------------------------------------------------- /ptlflow/models/neuflow/matching.py: -------------------------------------------------------------------------------- 1 | import torch.nn.functional as F 2 | 3 | from . import utils 4 | 5 | 6 | class Matching: 7 | def init_grid(self, batch_size, height, width, dtype, device): 8 | self.grid = utils.coords_grid( 9 | batch_size, height, width, dtype=dtype, device=device 10 | ) # [B, 2, H, W] 11 | self.flatten_grid = self.grid.view(batch_size, 2, -1).permute( 12 | 0, 2, 1 13 | ) # [B, H*W, 2] 14 | 15 | def global_correlation_softmax(self, feature0, feature1): 16 | b, c, h, w = feature0.shape 17 | 18 | feature0 = feature0.flatten(-2).permute(0, 2, 1) 19 | feature1 = feature1.flatten(-2).permute(0, 2, 1) 20 | 21 | correspondence = F.scaled_dot_product_attention( 22 | feature0, feature1, self.flatten_grid 23 | ) 24 | 25 | correspondence = correspondence.view(b, h, w, 2).permute( 26 | 0, 3, 1, 2 27 | ) # [B, 2, H, W] 28 | 29 | flow = correspondence - self.grid 30 | 31 | return flow 32 | -------------------------------------------------------------------------------- /ptlflow/models/neuflow/upsample.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn.functional as F 3 | 4 | 5 | class UpSample(torch.nn.Module): 6 | def __init__(self, feature_dim, upsample_factor): 7 | super(UpSample, self).__init__() 8 | 9 | self.upsample_factor = upsample_factor 10 | 11 | self.conv1 = torch.nn.Conv2d(2 + feature_dim, 256, 3, 1, 1) 12 | self.conv2 = torch.nn.Conv2d(256, 512, 3, 1, 1) 13 | self.conv3 = torch.nn.Conv2d(512, upsample_factor**2 * 9, 1, 1, 0) 14 | self.relu = torch.nn.ReLU(inplace=True) 15 | 16 | def forward(self, feature, flow): 17 | concat = torch.cat((flow, feature), dim=1) 18 | 19 | mask = self.conv3(self.relu(self.conv2(self.relu(self.conv1(concat))))) 20 | 21 | b, _, h, w = flow.shape 22 | 23 | mask = mask.view( 24 | b, 1, 9, self.upsample_factor, self.upsample_factor, h, w 25 | ) # [B, 1, 9, K, K, H, W] 26 | mask = torch.softmax(mask, dim=2) 27 | 28 | up_flow = F.unfold(self.upsample_factor * flow, [3, 3], padding=1) 29 | up_flow = up_flow.view(b, 2, 9, 1, 1, h, w) # [B, 2, 9, 1, 1, H, W] 30 | 31 | up_flow = torch.sum(mask * up_flow, dim=2) # [B, 2, K, K, H, W] 32 | up_flow = up_flow.permute(0, 1, 4, 2, 5, 3) # [B, 2, K, H, K, W] 33 | up_flow = up_flow.reshape( 34 | b, 2, self.upsample_factor * h, self.upsample_factor * w 35 | ) # [B, 2, K*H, K*W] 36 | 37 | return up_flow 38 | -------------------------------------------------------------------------------- /ptlflow/models/neuflow/utils.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn.functional as F 3 | 4 | 5 | def normalize_img(img, mean, std): 6 | return (img / 255.0 - mean) / std 7 | 8 | 9 | def coords_grid(b, h, w, dtype, device): 10 | ys, xs = torch.meshgrid( 11 | torch.arange(h, dtype=dtype, device=device), 12 | torch.arange(w, dtype=dtype, device=device), 13 | indexing="ij", 14 | ) # [H, W] 15 | 16 | stacks = [xs, ys] 17 | 18 | grid = torch.stack(stacks, dim=0) # [2, H, W] or [3, H, W] 19 | 20 | grid = grid[None].repeat(b, 1, 1, 1) # [B, 2, H, W] or [B, 3, H, W] 21 | 22 | return grid 23 | 24 | 25 | def bilinear_sample(img, sample_coords): 26 | b, _, h, w = sample_coords.shape 27 | 28 | # Normalize to [-1, 1] 29 | x_grid = 2 * sample_coords[:, 0] / (w - 1) - 1 30 | y_grid = 2 * sample_coords[:, 1] / (h - 1) - 1 31 | 32 | grid = torch.stack([x_grid, y_grid], dim=-1) # [B, H, W, 2] 33 | 34 | img = F.grid_sample( 35 | img, grid, mode="bilinear", padding_mode="zeros", align_corners=True 36 | ) 37 | 38 | return img 39 | 40 | 41 | def flow_warp(feature, flow): 42 | b, c, h, w = feature.size() 43 | 44 | grid = ( 45 | coords_grid(b, h, w, dtype=flow.dtype, device=flow.device) + flow 46 | ) # [B, 2, H, W] 47 | 48 | return bilinear_sample(feature, grid) 49 | -------------------------------------------------------------------------------- /ptlflow/models/neuflow2/README.md: -------------------------------------------------------------------------------- 1 | # NeuFlow v2 2 | 3 | ## Original code 4 | 5 | [https://github.com/neufieldrobotics/NeuFlow_v2](https://github.com/neufieldrobotics/NeuFlow_v2) 6 | 7 | ## Code license 8 | 9 | See [LICENSE](LICENSE). 10 | 11 | ## Pretrained weights license 12 | 13 | Not specified. 14 | 15 | ## Citation 16 | 17 | ``` 18 | @Article{Zhang2024NeuFlowRealtime, 19 | author = {Zhiyong Zhang and Aniket Gupta and Huaizu Jiang and Hanumant Singh}, 20 | journal = {ArXiv}, 21 | title = {{NeuFlow} v2: Push High-Efficiency Optical Flow To the Limit}, 22 | year = {2024}, 23 | } 24 | ``` -------------------------------------------------------------------------------- /ptlflow/models/neuflow2/__init__.py: -------------------------------------------------------------------------------- 1 | from .neuflow2 import * 2 | -------------------------------------------------------------------------------- /ptlflow/models/neuflow2/matching.py: -------------------------------------------------------------------------------- 1 | import torch.nn.functional as F 2 | 3 | from . import utils 4 | 5 | 6 | class Matching: 7 | def init_bhwd(self, batch_size, height, width, device, amp): 8 | self.grid = utils.coords_grid( 9 | batch_size, height, width, device, amp 10 | ) # [B, 2, H, W] 11 | self.flatten_grid = self.grid.view(batch_size, 2, -1).permute( 12 | 0, 2, 1 13 | ) # [B, H*W, 2] 14 | 15 | def global_correlation_softmax(self, feature0, feature1): 16 | b, c, h, w = feature0.shape 17 | 18 | feature0 = feature0.flatten(-2).permute(0, 2, 1) 19 | feature1 = feature1.flatten(-2).permute(0, 2, 1) 20 | 21 | correspondence = F.scaled_dot_product_attention( 22 | feature0, feature1, self.flatten_grid 23 | ) 24 | 25 | correspondence = correspondence.view(b, h, w, 2).permute( 26 | 0, 3, 1, 2 27 | ) # [B, 2, H, W] 28 | 29 | flow = correspondence - self.grid 30 | 31 | return flow 32 | -------------------------------------------------------------------------------- /ptlflow/models/neuflow2/upsample.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn.functional as F 3 | 4 | # from spatial_correlation_sampler import SpatialCorrelationSampler 5 | 6 | 7 | class UpSample(torch.nn.Module): 8 | def __init__(self, feature_dim, upsample_factor): 9 | super(UpSample, self).__init__() 10 | 11 | self.upsample_factor = upsample_factor 12 | 13 | self.conv1 = torch.nn.Conv2d(2 + feature_dim, 256, 3, 1, 1) 14 | self.conv2 = torch.nn.Conv2d(256, 512, 3, 1, 1) 15 | self.conv3 = torch.nn.Conv2d(512, upsample_factor**2 * 9, 1, 1, 0) 16 | self.relu = torch.nn.ReLU(inplace=True) 17 | 18 | def forward(self, feature, flow): 19 | concat = torch.cat((flow, feature), dim=1) 20 | 21 | mask = self.conv3(self.relu(self.conv2(self.relu(self.conv1(concat))))) 22 | 23 | b, _, h, w = flow.shape 24 | 25 | mask = mask.view( 26 | b, 1, 9, self.upsample_factor, self.upsample_factor, h, w 27 | ) # [B, 1, 9, K, K, H, W] 28 | mask = torch.softmax(mask, dim=2) 29 | 30 | # up_flow = F.unfold(self.upsample_factor * flow, [3, 3], padding=1) 31 | up_flow = F.unfold(flow, [3, 3], padding=1) 32 | up_flow = up_flow.view(b, 2, 9, 1, 1, h, w) # [B, 2, 9, 1, 1, H, W] 33 | 34 | up_flow = torch.sum(mask * up_flow, dim=2) # [B, 2, K, K, H, W] 35 | up_flow = up_flow.permute(0, 1, 4, 2, 5, 3) # [B, 2, K, H, K, W] 36 | up_flow = up_flow.reshape( 37 | b, 2, self.upsample_factor * h, self.upsample_factor * w 38 | ) # [B, 2, K*H, K*W] 39 | 40 | return up_flow 41 | -------------------------------------------------------------------------------- /ptlflow/models/neuflow2/utils.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn.functional as F 3 | 4 | # def normalize(x): 5 | # x_min = x.min() 6 | # return (x - x_min) / (x.max() - x_min) 7 | 8 | 9 | def coords_grid(b, h, w, device, amp): 10 | ys, xs = torch.meshgrid( 11 | torch.arange(h, dtype=torch.half if amp else torch.float, device=device), 12 | torch.arange(w, dtype=torch.half if amp else torch.float, device=device), 13 | indexing="ij", 14 | ) # [H, W] 15 | 16 | grid = torch.stack([xs, ys], dim=0) # [2, H, W] or [3, H, W] 17 | 18 | grid = grid[None].repeat(b, 1, 1, 1) # [B, 2, H, W] or [B, 3, H, W] 19 | 20 | return grid 21 | 22 | 23 | def bilinear_sample(img, sample_coords): 24 | b, _, h, w = sample_coords.shape 25 | 26 | # Normalize to [-1, 1] 27 | x_grid = 2 * sample_coords[:, 0] / (w - 1) - 1 28 | y_grid = 2 * sample_coords[:, 1] / (h - 1) - 1 29 | 30 | grid = torch.stack([x_grid, y_grid], dim=-1) # [B, H, W, 2] 31 | 32 | img = F.grid_sample( 33 | img, grid, mode="bilinear", padding_mode="zeros", align_corners=True 34 | ) 35 | 36 | return img 37 | 38 | 39 | def flow_warp(feature, flow): 40 | b, c, h, w = feature.size() 41 | 42 | grid = coords_grid(b, h, w).to(flow.device) + flow # [B, 2, H, W] 43 | 44 | return bilinear_sample(feature, grid) 45 | -------------------------------------------------------------------------------- /ptlflow/models/pwcnet/README.md: -------------------------------------------------------------------------------- 1 | # PWCNet 2 | 3 | ## Original code 4 | 5 | [https://github.com/NVlabs/PWC-Net](https://github.com/NVlabs/PWC-Net) 6 | 7 | ## Code license 8 | 9 | See [LICENSE](LICENSE). 10 | 11 | ## Pretrained weights license 12 | 13 | See [LICENSE](LICENSE). 14 | 15 | ## Citation 16 | 17 | ``` 18 | @InProceedings{Sun2018PWC-Net, 19 | author = {Deqing Sun and Xiaodong Yang and Ming-Yu Liu and Jan Kautz}, 20 | title = {{PWC-Net}: {CNNs} for Optical Flow Using Pyramid, Warping, and Cost Volume}, 21 | booktitle = CVPR, 22 | year = {2018}, 23 | } 24 | ``` -------------------------------------------------------------------------------- /ptlflow/models/pwcnet/__init__.py: -------------------------------------------------------------------------------- 1 | from .pwcnet import * 2 | -------------------------------------------------------------------------------- /ptlflow/models/raft/LICENSE: -------------------------------------------------------------------------------- 1 | BSD 3-Clause License 2 | 3 | Copyright (c) 2020, princeton-vl 4 | All rights reserved. 5 | 6 | Redistribution and use in source and binary forms, with or without 7 | modification, are permitted provided that the following conditions are met: 8 | 9 | * Redistributions of source code must retain the above copyright notice, this 10 | list of conditions and the following disclaimer. 11 | 12 | * Redistributions in binary form must reproduce the above copyright notice, 13 | this list of conditions and the following disclaimer in the documentation 14 | and/or other materials provided with the distribution. 15 | 16 | * Neither the name of the copyright holder nor the names of its 17 | contributors may be used to endorse or promote products derived from 18 | this software without specific prior written permission. 19 | 20 | THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" 21 | AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 22 | IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE 23 | DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE 24 | FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL 25 | DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR 26 | SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER 27 | CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, 28 | OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 29 | OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. -------------------------------------------------------------------------------- /ptlflow/models/raft/README.md: -------------------------------------------------------------------------------- 1 | # RAFT 2 | 3 | ## Original code 4 | 5 | [https://github.com/princeton-vl/RAFT](https://github.com/princeton-vl/RAFT) 6 | 7 | ## Code license 8 | 9 | See [LICENSE](LICENSE). 10 | 11 | ## Pretrained weights license 12 | 13 | Not specified. 14 | 15 | ## Citation 16 | 17 | ``` 18 | @inproceedings{teed2020raft, 19 | title={Raft: Recurrent all-pairs field transforms for optical flow}, 20 | author={Teed, Zachary and Deng, Jia}, 21 | booktitle={European Conference on Computer Vision}, 22 | pages={402--419}, 23 | year={2020}, 24 | organization={Springer} 25 | } 26 | ``` 27 | 28 | ## Training 29 | 30 | This model can be trained on PTLFlow by following the [PTLFlow training instructions](https://ptlflow.readthedocs.io/en/latest/starting/training.html). 31 | 32 | ### Stage 1: FlyingChairs 33 | 34 | ```bash 35 | python train.py --config ptlflow/models/raft/configs/raft-train1-chairs.yaml 36 | ``` 37 | 38 | ### Stage 2: FlyingThings3D 39 | 40 | ```bash 41 | python train.py --config ptlflow/models/raft/configs/raft-train2-things.yaml 42 | ``` 43 | 44 | ### Stage 3: FlyingThings3D+Sintel+KITTI+HD1K 45 | ```bash 46 | python train.py --config ptlflow/models/raft/configs/raft-train3-sintel.yaml 47 | ``` 48 | 49 | ### Stage 4: KITTI 2015 50 | ```bash 51 | python train.py --config ptlflow/models/raft/configs/raft-train4-kitti.yaml 52 | ``` -------------------------------------------------------------------------------- /ptlflow/models/raft/__init__.py: -------------------------------------------------------------------------------- 1 | from .raft import * 2 | -------------------------------------------------------------------------------- /ptlflow/models/raft/configs/raft-train1-chairs.yaml: -------------------------------------------------------------------------------- 1 | # lightning.pytorch==2.4.0 2 | seed_everything: true 3 | lr: 0.0004 4 | wdecay: 0.0001 5 | trainer: 6 | max_epochs: 45 7 | accumulate_grad_batches: 1 8 | gradient_clip_val: 1.0 9 | model: 10 | class_path: ptlflow.models.raft 11 | init_args: 12 | corr_levels: 4 13 | corr_radius: 4 14 | dropout: 0.0 15 | gamma: 0.8 16 | max_flow: 400.0 17 | iters: 12 18 | alternate_corr: false 19 | warm_start: false 20 | data: 21 | train_dataset: chairs 22 | val_dataset: sintel-final-val+kitti-2015-val 23 | train_batch_size: 10 24 | train_num_workers: 4 25 | train_crop_size: [368, 496] 26 | train_transform_cuda: false 27 | train_transform_fp16: false 28 | -------------------------------------------------------------------------------- /ptlflow/models/raft/configs/raft-train2-things.yaml: -------------------------------------------------------------------------------- 1 | # lightning.pytorch==2.4.0 2 | seed_everything: true 3 | lr: 0.000125 4 | wdecay: 0.0001 5 | trainer: 6 | max_epochs: 8 7 | accumulate_grad_batches: 1 8 | gradient_clip_val: 1.0 9 | ckpt_path: /path/to/chairs.ckpt # Change to the ckpt resulting from raft-train1-chairs 10 | model: 11 | class_path: ptlflow.models.raft 12 | init_args: 13 | corr_levels: 4 14 | corr_radius: 4 15 | dropout: 0.0 16 | gamma: 0.8 17 | max_flow: 400.0 18 | iters: 12 19 | alternate_corr: false 20 | warm_start: false 21 | data: 22 | train_dataset: things 23 | val_dataset: sintel-final-val+kitti-2015-val 24 | train_batch_size: 6 25 | train_num_workers: 4 26 | train_crop_size: [400, 720] 27 | train_transform_cuda: false 28 | train_transform_fp16: false -------------------------------------------------------------------------------- /ptlflow/models/raft/configs/raft-train3-sintel.yaml: -------------------------------------------------------------------------------- 1 | # lightning.pytorch==2.4.0 2 | seed_everything: true 3 | lr: 0.000125 4 | wdecay: 0.00001 5 | trainer: 6 | max_epochs: 2 7 | accumulate_grad_batches: 1 8 | gradient_clip_val: 1.0 9 | ckpt_path: /path/to/things.ckpt # Change to the ckpt resulting from raft-train2-things 10 | model: 11 | class_path: ptlflow.models.raft 12 | init_args: 13 | corr_levels: 4 14 | corr_radius: 4 15 | dropout: 0.0 16 | gamma: 0.85 17 | max_flow: 400.0 18 | iters: 12 19 | alternate_corr: false 20 | warm_start: false 21 | data: 22 | train_dataset: sintel_finetune 23 | val_dataset: sintel-final-val+kitti-2015-val 24 | train_batch_size: 6 25 | train_num_workers: 4 26 | train_crop_size: [368, 768] 27 | train_transform_cuda: false 28 | train_transform_fp16: false -------------------------------------------------------------------------------- /ptlflow/models/raft/configs/raft-train4-kitti.yaml: -------------------------------------------------------------------------------- 1 | # lightning.pytorch==2.4.0 2 | seed_everything: true 3 | lr: 0.0001 4 | wdecay: 0.00001 5 | trainer: 6 | max_epochs: 1500 7 | check_val_every_n_epoch: 50 8 | accumulate_grad_batches: 1 9 | gradient_clip_val: 1.0 10 | ckpt_path: /path/to/sintel.ckpt # Change to the ckpt resulting from raft-train3-sintel 11 | model: 12 | class_path: ptlflow.models.raft 13 | init_args: 14 | corr_levels: 4 15 | corr_radius: 4 16 | dropout: 0.0 17 | gamma: 0.85 18 | max_flow: 400.0 19 | iters: 12 20 | alternate_corr: false 21 | warm_start: false 22 | data: 23 | train_dataset: kitti-2015 24 | val_dataset: kitti-2015-val+sintel-final-val 25 | train_batch_size: 6 26 | train_num_workers: 4 27 | train_crop_size: [288, 960] 28 | train_transform_cuda: false 29 | train_transform_fp16: false -------------------------------------------------------------------------------- /ptlflow/models/rapidflow/__init__.py: -------------------------------------------------------------------------------- 1 | from .rapidflow import * 2 | -------------------------------------------------------------------------------- /ptlflow/models/rapidflow/configs/rapidflow-test-kitti.yaml: -------------------------------------------------------------------------------- 1 | # lightning.pytorch==2.4.0 2 | ckpt_path: kitti 3 | output_path: outputs/test 4 | show: false 5 | max_forward_side: null 6 | scale_factor: null 7 | max_show_side: 1000 8 | save_viz: true 9 | seed_everything: true 10 | model: 11 | class_path: ptlflow.models.rapidflow 12 | init_args: 13 | pyramid_ranges: 14 | - 32 15 | - 8 16 | iters: 12 17 | corr_mode: allpairs 18 | corr_levels: 1 19 | corr_range: 4 20 | enc_hidden_chs: 64 21 | enc_out_chs: 128 22 | enc_stem_stride: 4 23 | enc_mlp_ratio: 4.0 24 | enc_depth: 4 25 | dec_net_chs: 64 26 | dec_inp_chs: 64 27 | dec_motion_chs: 128 28 | dec_depth: 2 29 | dec_mlp_ratio: 4.0 30 | use_upsample_mask: true 31 | fuse_next1d_weights: false 32 | simple_io: false 33 | gamma: 0.8 34 | max_flow: 400 35 | lr: null 36 | wdecay: null 37 | warm_start: false 38 | data: 39 | test_dataset: kitti-2015 40 | -------------------------------------------------------------------------------- /ptlflow/models/rapidflow/configs/rapidflow-test-sintel.yaml: -------------------------------------------------------------------------------- 1 | # lightning.pytorch==2.4.0 2 | ckpt_path: sintel 3 | output_path: outputs/test 4 | show: false 5 | max_forward_side: null 6 | scale_factor: null 7 | max_show_side: 1000 8 | save_viz: true 9 | seed_everything: true 10 | model: 11 | class_path: ptlflow.models.rapidflow 12 | init_args: 13 | pyramid_ranges: 14 | - 32 15 | - 8 16 | iters: 12 17 | corr_mode: allpairs 18 | corr_levels: 1 19 | corr_range: 4 20 | enc_hidden_chs: 64 21 | enc_out_chs: 128 22 | enc_stem_stride: 4 23 | enc_mlp_ratio: 4.0 24 | enc_depth: 4 25 | dec_net_chs: 64 26 | dec_inp_chs: 64 27 | dec_motion_chs: 128 28 | dec_depth: 2 29 | dec_mlp_ratio: 4.0 30 | use_upsample_mask: true 31 | fuse_next1d_weights: false 32 | simple_io: false 33 | gamma: 0.8 34 | max_flow: 400 35 | lr: null 36 | wdecay: null 37 | warm_start: true 38 | data: 39 | test_dataset: sintel 40 | -------------------------------------------------------------------------------- /ptlflow/models/rapidflow/configs/rapidflow-train1-chairs.yaml: -------------------------------------------------------------------------------- 1 | # lightning.pytorch==2.4.0 2 | seed_everything: true 3 | lr: 0.0004 4 | wdecay: 0.0001 5 | trainer: 6 | max_epochs: 10 7 | accumulate_grad_batches: 1 8 | gradient_clip_val: 1.0 9 | devices: 0, 10 | model: 11 | class_path: ptlflow.models.rapidflow 12 | init_args: 13 | pyramid_ranges: 14 | - 32 15 | - 8 16 | iters: 12 17 | corr_mode: allpairs 18 | corr_levels: 1 19 | corr_range: 4 20 | enc_hidden_chs: 64 21 | enc_out_chs: 128 22 | enc_stem_stride: 4 23 | enc_mlp_ratio: 4.0 24 | enc_depth: 4 25 | dec_net_chs: 64 26 | dec_inp_chs: 64 27 | dec_motion_chs: 128 28 | dec_depth: 2 29 | dec_mlp_ratio: 4.0 30 | use_upsample_mask: true 31 | fuse_next1d_weights: false 32 | simple_io: false 33 | gamma: 0.8 34 | max_flow: 400 35 | warm_start: false 36 | data: 37 | train_dataset: chairs 38 | val_dataset: sintel-final-val+kitti-2015-val 39 | train_batch_size: 8 40 | train_num_workers: 4 41 | train_crop_size: [352, 480] 42 | train_transform_cuda: false 43 | train_transform_fp16: false 44 | -------------------------------------------------------------------------------- /ptlflow/models/rapidflow/configs/rapidflow-train2-things.yaml: -------------------------------------------------------------------------------- 1 | # lightning.pytorch==2.4.0 2 | seed_everything: true 3 | ckpt_path: /path/to/chairs.ckpt # Change to the ckpt resulting from rapidflow-train1-chairs 4 | lr: 0.000125 5 | wdecay: 0.0001 6 | trainer: 7 | max_epochs: 10 8 | accumulate_grad_batches: 1 9 | gradient_clip_val: 1.0 10 | devices: 0, 11 | model: 12 | class_path: ptlflow.models.rapidflow 13 | init_args: 14 | pyramid_ranges: 15 | - 32 16 | - 8 17 | iters: 12 18 | corr_mode: allpairs 19 | corr_levels: 1 20 | corr_range: 4 21 | enc_hidden_chs: 64 22 | enc_out_chs: 128 23 | enc_stem_stride: 4 24 | enc_mlp_ratio: 4.0 25 | enc_depth: 4 26 | dec_net_chs: 64 27 | dec_inp_chs: 64 28 | dec_motion_chs: 128 29 | dec_depth: 2 30 | dec_mlp_ratio: 4.0 31 | use_upsample_mask: true 32 | fuse_next1d_weights: false 33 | simple_io: false 34 | gamma: 0.8 35 | max_flow: 400 36 | warm_start: false 37 | data: 38 | train_dataset: things 39 | val_dataset: sintel-final-val+kitti-2015-val 40 | train_batch_size: 4 41 | train_num_workers: 2 42 | train_crop_size: [384, 704] 43 | train_transform_cuda: false 44 | train_transform_fp16: false 45 | -------------------------------------------------------------------------------- /ptlflow/models/rapidflow/configs/rapidflow-train3-sintel.yaml: -------------------------------------------------------------------------------- 1 | # lightning.pytorch==2.4.0 2 | seed_everything: true 3 | ckpt_path: /path/to/things.ckpt # Change to the ckpt resulting from rapidflow-train2-things 4 | lr: 0.000125 5 | wdecay: 0.00001 6 | trainer: 7 | max_epochs: 4 8 | accumulate_grad_batches: 1 9 | gradient_clip_val: 1.0 10 | devices: 0, 11 | model: 12 | class_path: ptlflow.models.rapidflow 13 | init_args: 14 | pyramid_ranges: 15 | - 32 16 | - 8 17 | iters: 12 18 | corr_mode: allpairs 19 | corr_levels: 1 20 | corr_range: 4 21 | enc_hidden_chs: 64 22 | enc_out_chs: 128 23 | enc_stem_stride: 4 24 | enc_mlp_ratio: 4.0 25 | enc_depth: 4 26 | dec_net_chs: 64 27 | dec_inp_chs: 64 28 | dec_motion_chs: 128 29 | dec_depth: 2 30 | dec_mlp_ratio: 4.0 31 | use_upsample_mask: true 32 | fuse_next1d_weights: false 33 | simple_io: false 34 | gamma: 0.85 35 | max_flow: 400 36 | warm_start: false 37 | data: 38 | train_dataset: sintel_finetune 39 | val_dataset: sintel-final-val+kitti-2015-val 40 | train_batch_size: 6 41 | train_num_workers: 3 42 | train_crop_size: [352, 768] 43 | train_transform_cuda: false 44 | train_transform_fp16: false 45 | -------------------------------------------------------------------------------- /ptlflow/models/rapidflow/configs/rapidflow-train4-kitti.yaml: -------------------------------------------------------------------------------- 1 | # lightning.pytorch==2.4.0 2 | seed_everything: true 3 | ckpt_path: /path/to/sintel.ckpt # Change to the ckpt resulting from rapidflow-train3-sintel 4 | lr: 0.000125 5 | wdecay: 0.00001 6 | trainer: 7 | max_epochs: 300 8 | check_val_every_n_epoch: 10 9 | accumulate_grad_batches: 1 10 | gradient_clip_val: 1.0 11 | devices: 0, 12 | model: 13 | class_path: ptlflow.models.rapidflow 14 | init_args: 15 | pyramid_ranges: 16 | - 32 17 | - 8 18 | iters: 12 19 | corr_mode: allpairs 20 | corr_levels: 1 21 | corr_range: 4 22 | enc_hidden_chs: 64 23 | enc_out_chs: 128 24 | enc_stem_stride: 4 25 | enc_mlp_ratio: 4.0 26 | enc_depth: 4 27 | dec_net_chs: 64 28 | dec_inp_chs: 64 29 | dec_motion_chs: 128 30 | dec_depth: 2 31 | dec_mlp_ratio: 4.0 32 | use_upsample_mask: true 33 | fuse_next1d_weights: false 34 | simple_io: false 35 | gamma: 0.85 36 | max_flow: 400 37 | warm_start: false 38 | data: 39 | train_dataset: kitti-2015 40 | val_dataset: kitti-2015-val+sintel-final-val 41 | train_batch_size: 6 42 | train_num_workers: 3 43 | train_crop_size: [288, 960] 44 | train_transform_cuda: false 45 | train_transform_fp16: false 46 | -------------------------------------------------------------------------------- /ptlflow/models/rapidflow/configs/rapidflow-validate-things.yaml: -------------------------------------------------------------------------------- 1 | # lightning.pytorch==2.4.0 2 | all: false 3 | select: null 4 | exclude: null 5 | ckpt_path: things 6 | output_path: outputs/validate 7 | write_outputs: false 8 | show: false 9 | flow_format: original 10 | max_forward_side: null 11 | scale_factor: null 12 | max_show_side: 1000 13 | max_samples: null 14 | reversed: false 15 | fp16: true 16 | seq_val_mode: all 17 | write_individual_metrics: false 18 | epe_clip: 5.0 19 | seed_everything: true 20 | model: 21 | class_path: ptlflow.models.rapidflow 22 | init_args: 23 | pyramid_ranges: 24 | - 32 25 | - 8 26 | iters: 12 27 | corr_mode: allpairs 28 | corr_levels: 1 29 | corr_range: 4 30 | enc_hidden_chs: 64 31 | enc_out_chs: 128 32 | enc_stem_stride: 4 33 | enc_mlp_ratio: 4.0 34 | enc_depth: 4 35 | dec_net_chs: 64 36 | dec_inp_chs: 64 37 | dec_motion_chs: 128 38 | dec_depth: 2 39 | dec_mlp_ratio: 4.0 40 | use_upsample_mask: true 41 | fuse_next1d_weights: false 42 | simple_io: false 43 | gamma: 0.8 44 | max_flow: 400 45 | lr: null 46 | wdecay: null 47 | warm_start: false 48 | data: 49 | val_dataset: sintel-clean+sintel-final+kitti-2012+kitti-2015 50 | -------------------------------------------------------------------------------- /ptlflow/models/rapidflow/image_samples/000000_10.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hmorimitsu/ptlflow/fb6d753c3b16e74481258891b45668dcba63e54e/ptlflow/models/rapidflow/image_samples/000000_10.png -------------------------------------------------------------------------------- /ptlflow/models/rapidflow/image_samples/000000_11.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hmorimitsu/ptlflow/fb6d753c3b16e74481258891b45668dcba63e54e/ptlflow/models/rapidflow/image_samples/000000_11.png -------------------------------------------------------------------------------- /ptlflow/models/rapidflow/image_samples/README.md: -------------------------------------------------------------------------------- 1 | These image samples correspond to the first image pair of the training set of the KITTI 2015 dataset. 2 | 3 | The complete dataset is available at: [https://www.cvlibs.net/datasets/kitti/eval_scene_flow.php?benchmark=flow](https://www.cvlibs.net/datasets/kitti/eval_scene_flow.php?benchmark=flow). -------------------------------------------------------------------------------- /ptlflow/models/rapidflow/local_timm/README.md: -------------------------------------------------------------------------------- 1 | These files are adapted from the timm library: [https://github.com/rwightman/pytorch-image-models](https://github.com/rwightman/pytorch-image-models). 2 | 3 | Most files are stripped down versions of the original ones, but they may also contain small modifications to adapt the modules to our model. 4 | 5 | timm is licensed under the Apache 2.0 license: [https://github.com/rwightman/pytorch-image-models/blob/main/LICENSE](https://github.com/rwightman/pytorch-image-models/blob/main/LICENSE). -------------------------------------------------------------------------------- /ptlflow/models/rapidflow/local_timm/conv2d_same.py: -------------------------------------------------------------------------------- 1 | """ Conv2d w/ Same Padding 2 | 3 | Hacked together by / Copyright 2020 Ross Wightman 4 | """ 5 | 6 | import torch 7 | import torch.nn as nn 8 | import torch.nn.functional as F 9 | from typing import Tuple, Optional 10 | 11 | from .padding import pad_same, get_padding_value 12 | 13 | 14 | def conv2d_same( 15 | x, 16 | weight: torch.Tensor, 17 | bias: Optional[torch.Tensor] = None, 18 | stride: Tuple[int, int] = (1, 1), 19 | padding: Tuple[int, int] = (0, 0), 20 | dilation: Tuple[int, int] = (1, 1), 21 | groups: int = 1, 22 | ): 23 | x = pad_same(x, weight.shape[-2:], stride, dilation) 24 | return F.conv2d(x, weight, bias, stride, (0, 0), dilation, groups) 25 | 26 | 27 | class Conv2dSame(nn.Conv2d): 28 | """Tensorflow like 'SAME' convolution wrapper for 2D convolutions""" 29 | 30 | def __init__( 31 | self, 32 | in_channels, 33 | out_channels, 34 | kernel_size, 35 | stride=1, 36 | padding=0, 37 | dilation=1, 38 | groups=1, 39 | bias=True, 40 | ): 41 | super(Conv2dSame, self).__init__( 42 | in_channels, out_channels, kernel_size, stride, 0, dilation, groups, bias 43 | ) 44 | 45 | def forward(self, x): 46 | return conv2d_same( 47 | x, 48 | self.weight, 49 | self.bias, 50 | self.stride, 51 | self.padding, 52 | self.dilation, 53 | self.groups, 54 | ) 55 | 56 | 57 | def create_conv2d_pad(in_chs, out_chs, kernel_size, **kwargs): 58 | padding = kwargs.pop("padding", "") 59 | kwargs.setdefault("bias", False) 60 | padding, is_dynamic = get_padding_value(padding, kernel_size, **kwargs) 61 | if is_dynamic: 62 | return Conv2dSame(in_chs, out_chs, kernel_size, **kwargs) 63 | else: 64 | return nn.Conv2d(in_chs, out_chs, kernel_size, padding=padding, **kwargs) 65 | -------------------------------------------------------------------------------- /ptlflow/models/rapidflow/local_timm/create_conv2d.py: -------------------------------------------------------------------------------- 1 | """ Create Conv2d Factory Method 2 | 3 | Hacked together by / Copyright 2020 Ross Wightman 4 | """ 5 | 6 | from .mixed_conv2d import MixedConv2d 7 | from .cond_conv2d import CondConv2d 8 | from .conv2d_same import create_conv2d_pad 9 | 10 | 11 | def create_conv2d(in_channels, out_channels, kernel_size, **kwargs): 12 | """Select a 2d convolution implementation based on arguments 13 | Creates and returns one of torch.nn.Conv2d, Conv2dSame, MixedConv2d, or CondConv2d. 14 | 15 | Used extensively by EfficientNet, MobileNetv3 and related networks. 16 | """ 17 | if isinstance(kernel_size, list): 18 | assert ( 19 | "num_experts" not in kwargs 20 | ) # MixNet + CondConv combo not supported currently 21 | if "groups" in kwargs: 22 | groups = kwargs.pop("groups") 23 | if groups == in_channels: 24 | kwargs["depthwise"] = True 25 | else: 26 | assert groups == 1 27 | # We're going to use only lists for defining the MixedConv2d kernel groups, 28 | # ints, tuples, other iterables will continue to pass to normal conv and specify h, w. 29 | m = MixedConv2d(in_channels, out_channels, kernel_size, **kwargs) 30 | else: 31 | depthwise = kwargs.pop("depthwise", False) 32 | # for DW out_channels must be multiple of in_channels as must have out_channels % groups == 0 33 | groups = in_channels if depthwise else kwargs.pop("groups", 1) 34 | if "num_experts" in kwargs and kwargs["num_experts"] > 0: 35 | m = CondConv2d( 36 | in_channels, out_channels, kernel_size, groups=groups, **kwargs 37 | ) 38 | else: 39 | m = create_conv2d_pad( 40 | in_channels, out_channels, kernel_size, groups=groups, **kwargs 41 | ) 42 | return m 43 | -------------------------------------------------------------------------------- /ptlflow/models/rapidflow/local_timm/gelu.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | 5 | 6 | class GELU(nn.Module): 7 | """Applies the Gaussian Error Linear Units function (w/ dummy inplace arg)""" 8 | 9 | def __init__(self, inplace: bool = False): 10 | super(GELU, self).__init__() 11 | 12 | def forward(self, input: torch.Tensor) -> torch.Tensor: 13 | return F.gelu(input) 14 | -------------------------------------------------------------------------------- /ptlflow/models/rapidflow/local_timm/layer_helpers.py: -------------------------------------------------------------------------------- 1 | """ Layer/Module Helpers 2 | 3 | Hacked together by / Copyright 2020 Ross Wightman 4 | """ 5 | 6 | from itertools import repeat 7 | import collections.abc 8 | 9 | 10 | # From PyTorch internals 11 | def _ntuple(n): 12 | def parse(x): 13 | if isinstance(x, collections.abc.Iterable) and not isinstance(x, str): 14 | return x 15 | return tuple(repeat(x, n)) 16 | 17 | return parse 18 | 19 | 20 | to_1tuple = _ntuple(1) 21 | to_2tuple = _ntuple(2) 22 | to_3tuple = _ntuple(3) 23 | to_4tuple = _ntuple(4) 24 | to_ntuple = _ntuple 25 | 26 | 27 | def make_divisible(v, divisor=8, min_value=None, round_limit=0.9): 28 | min_value = min_value or divisor 29 | new_v = max(min_value, int(v + divisor / 2) // divisor * divisor) 30 | # Make sure that round down does not go down by more than 10%. 31 | if new_v < round_limit * v: 32 | new_v += divisor 33 | return new_v 34 | 35 | 36 | def extend_tuple(x, n): 37 | # pdas a tuple to specified n by padding with last value 38 | if not isinstance(x, (tuple, list)): 39 | x = (x,) 40 | else: 41 | x = tuple(x) 42 | pad_n = n - len(x) 43 | if pad_n <= 0: 44 | return x[:n] 45 | return x + (x[-1],) * pad_n 46 | -------------------------------------------------------------------------------- /ptlflow/models/rapidflow/local_timm/norm.py: -------------------------------------------------------------------------------- 1 | """ Normalization layers and wrappers 2 | 3 | Norm layer definitions that support fast norm and consistent channel arg order (always first arg). 4 | 5 | Hacked together by / Copyright 2022 Ross Wightman 6 | """ 7 | 8 | import torch 9 | import torch.nn as nn 10 | import torch.nn.functional as F 11 | 12 | 13 | class LayerNorm(nn.LayerNorm): 14 | """LayerNorm w/ fast norm option""" 15 | 16 | def __init__(self, num_channels, eps=1e-6, affine=True): 17 | super().__init__(num_channels, eps=eps, elementwise_affine=affine) 18 | 19 | def forward(self, x: torch.Tensor) -> torch.Tensor: 20 | x = F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps) 21 | return x 22 | 23 | 24 | class LayerNorm2d(nn.LayerNorm): 25 | """LayerNorm for channels of '2D' spatial NCHW tensors""" 26 | 27 | def __init__(self, num_channels, eps=1e-6, affine=True): 28 | super().__init__(num_channels, eps=eps, elementwise_affine=affine) 29 | 30 | def forward(self, x: torch.Tensor) -> torch.Tensor: 31 | x = x.permute(0, 2, 3, 1) 32 | x = F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps) 33 | x = x.permute(0, 3, 1, 2) 34 | return x 35 | -------------------------------------------------------------------------------- /ptlflow/models/rapidflow/pwc_modules.py: -------------------------------------------------------------------------------- 1 | # ============================================================================= 2 | # Copyright 2024 Henrique Morimitsu 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | # Adapted from IRR: https://github.com/visinf/irr/blob/master/models/pwc_modules.py 17 | # 18 | # Modifications by Henrique Morimitsu: 19 | # - Remove operations not related to rescaling 20 | # ============================================================================= 21 | 22 | import torch 23 | import torch.nn.functional as F 24 | 25 | 26 | def upsample2d_as(inputs, target_as, mode="bilinear"): 27 | _, _, h, w = target_as.size() 28 | if inputs.shape[-2] != h or inputs.shape[-1] != w: 29 | inputs = F.interpolate(inputs, [h, w], mode=mode, align_corners=True) 30 | return inputs 31 | 32 | 33 | def rescale_flow(flow, width_im, height_im, to_local=True): 34 | if to_local: 35 | u_scale = float(flow.size(3) / width_im) 36 | v_scale = float(flow.size(2) / height_im) 37 | else: 38 | u_scale = float(width_im / flow.size(3)) 39 | v_scale = float(height_im / flow.size(2)) 40 | 41 | u, v = flow.chunk(2, dim=1) 42 | u = u * u_scale 43 | v = v * v_scale 44 | 45 | return torch.cat([u, v], dim=1) 46 | -------------------------------------------------------------------------------- /ptlflow/models/rapidflow/utils.py: -------------------------------------------------------------------------------- 1 | # ============================================================================= 2 | # Copyright 2024 Henrique Morimitsu 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | # Adapted from RAFT: https://github.com/princeton-vl/RAFT 17 | # 18 | # Original license BSD 3-clause: https://github.com/princeton-vl/RAFT/blob/master/LICENSE 19 | # 20 | # Modifications by Henrique Morimitsu: 21 | # - Adapt residual block 22 | # ============================================================================= 23 | 24 | import torch 25 | import torch.nn as nn 26 | import torch.nn.functional as F 27 | 28 | 29 | def bilinear_sampler(img, coords, mask=False): 30 | """Wrapper for grid_sample, uses pixel coordinates""" 31 | H, W = img.shape[-2:] 32 | xgrid, ygrid = coords.split([1, 1], dim=-1) 33 | xgrid = 2 * xgrid / (W - 1) - 1 34 | ygrid = 2 * ygrid / (H - 1) - 1 35 | 36 | grid = torch.cat([xgrid, ygrid], dim=-1) 37 | img = F.grid_sample(img, grid, align_corners=True) 38 | 39 | if mask: 40 | mask = (xgrid > -1) & (ygrid > -1) & (xgrid < 1) & (ygrid < 1) 41 | return img, mask.to(dtype=coords.dtype) 42 | 43 | return img 44 | -------------------------------------------------------------------------------- /ptlflow/models/rpknet/__init__.py: -------------------------------------------------------------------------------- 1 | from .rpknet import * 2 | -------------------------------------------------------------------------------- /ptlflow/models/rpknet/configs/rpknet-test-kitti.yaml: -------------------------------------------------------------------------------- 1 | # lightning.pytorch==2.4.0 2 | ckpt_path: kitti 3 | output_path: outputs/test 4 | show: false 5 | max_forward_side: null 6 | scale_factor: null 7 | max_show_side: 1000 8 | save_viz: true 9 | seed_everything: true 10 | model: 11 | class_path: ptlflow.models.rpknet 12 | init_args: 13 | pyramid_ranges: 14 | - 32 15 | - 8 16 | iters: 32 17 | input_pad_one_side: true 18 | input_bgr_to_rgb: false 19 | detach_flow: true 20 | corr_mode: allpairs 21 | upgate_norm_type: group 22 | use_norm_affine: false 23 | group_norm_num_groups: 4 24 | corr_levels: 1 25 | corr_range: 4 26 | enc_norm_type: group 27 | enc_stem_stride: 2 28 | enc_depth: 2 29 | enc_mlp_ratio: 4.0 30 | enc_hidden_chs: 31 | - 32 32 | - 64 33 | - 96 34 | enc_out_1x1_chs: '2.0' 35 | dec_gru_iters: 2 36 | dec_gru_depth: 2 37 | dec_gru_mlp_ratio: 4.0 38 | dec_net_chs: null 39 | dec_inp_chs: null 40 | dec_motion_chs: 128 41 | use_upsample_mask: true 42 | upmask_gradient_scale: 1.0 43 | cache_pkconv_weights: true 44 | gamma: 0.8 45 | max_flow: 400 46 | lr: null 47 | wdecay: null 48 | warm_start: false 49 | data: 50 | test_dataset: kitti-2015 51 | -------------------------------------------------------------------------------- /ptlflow/models/rpknet/configs/rpknet-test-sintel.yaml: -------------------------------------------------------------------------------- 1 | # lightning.pytorch==2.4.0 2 | ckpt_path: sintel 3 | output_path: outputs/test 4 | show: false 5 | max_forward_side: null 6 | scale_factor: null 7 | max_show_side: 1000 8 | save_viz: true 9 | seed_everything: true 10 | model: 11 | class_path: ptlflow.models.rpknet 12 | init_args: 13 | pyramid_ranges: 14 | - 32 15 | - 8 16 | iters: 32 17 | input_pad_one_side: false 18 | input_bgr_to_rgb: false 19 | detach_flow: true 20 | corr_mode: allpairs 21 | upgate_norm_type: group 22 | use_norm_affine: false 23 | group_norm_num_groups: 4 24 | corr_levels: 1 25 | corr_range: 4 26 | enc_norm_type: group 27 | enc_stem_stride: 2 28 | enc_depth: 2 29 | enc_mlp_ratio: 4.0 30 | enc_hidden_chs: 31 | - 32 32 | - 64 33 | - 96 34 | enc_out_1x1_chs: '2.0' 35 | dec_gru_iters: 2 36 | dec_gru_depth: 2 37 | dec_gru_mlp_ratio: 4.0 38 | dec_net_chs: null 39 | dec_inp_chs: null 40 | dec_motion_chs: 128 41 | use_upsample_mask: true 42 | upmask_gradient_scale: 1.0 43 | cache_pkconv_weights: true 44 | gamma: 0.8 45 | max_flow: 400 46 | lr: null 47 | wdecay: null 48 | warm_start: true 49 | data: 50 | test_dataset: sintel 51 | -------------------------------------------------------------------------------- /ptlflow/models/rpknet/configs/rpknet-test-spring.yaml: -------------------------------------------------------------------------------- 1 | # lightning.pytorch==2.4.0 2 | ckpt_path: sintel 3 | output_path: outputs/test 4 | show: false 5 | max_forward_side: null 6 | scale_factor: null 7 | max_show_side: 1000 8 | save_viz: true 9 | seed_everything: true 10 | model: 11 | class_path: ptlflow.models.rpknet 12 | init_args: 13 | pyramid_ranges: 14 | - 32 15 | - 8 16 | iters: 32 17 | input_pad_one_side: false 18 | input_bgr_to_rgb: true # There is no special reason to convert to RGB here. But this mode was used by accident when submitting our results 19 | detach_flow: true 20 | corr_mode: allpairs 21 | upgate_norm_type: group 22 | use_norm_affine: false 23 | group_norm_num_groups: 4 24 | corr_levels: 1 25 | corr_range: 4 26 | enc_norm_type: group 27 | enc_stem_stride: 2 28 | enc_depth: 2 29 | enc_mlp_ratio: 4.0 30 | enc_hidden_chs: 31 | - 32 32 | - 64 33 | - 96 34 | enc_out_1x1_chs: '2.0' 35 | dec_gru_iters: 2 36 | dec_gru_depth: 2 37 | dec_gru_mlp_ratio: 4.0 38 | dec_net_chs: null 39 | dec_inp_chs: null 40 | dec_motion_chs: 128 41 | use_upsample_mask: true 42 | upmask_gradient_scale: 1.0 43 | cache_pkconv_weights: true 44 | gamma: 0.8 45 | max_flow: 400 46 | lr: null 47 | wdecay: null 48 | warm_start: true 49 | data: 50 | test_dataset: spring 51 | -------------------------------------------------------------------------------- /ptlflow/models/rpknet/configs/rpknet-train1-chairs.yaml: -------------------------------------------------------------------------------- 1 | # lightning.pytorch==2.4.0 2 | seed_everything: true 3 | lr: 0.0004 4 | wdecay: 0.0001 5 | trainer: 6 | max_epochs: 45 7 | accumulate_grad_batches: 1 8 | gradient_clip_val: 1.0 9 | devices: 0, 10 | model: 11 | class_path: ptlflow.models.rpknet 12 | init_args: 13 | pyramid_ranges: 14 | - 32 15 | - 8 16 | iters: 12 17 | input_pad_one_side: false 18 | input_bgr_to_rgb: false 19 | detach_flow: true 20 | corr_mode: allpairs 21 | upgate_norm_type: group 22 | use_norm_affine: false 23 | group_norm_num_groups: 4 24 | corr_levels: 1 25 | corr_range: 4 26 | enc_norm_type: group 27 | enc_stem_stride: 2 28 | enc_depth: 2 29 | enc_mlp_ratio: 4.0 30 | enc_hidden_chs: 31 | - 32 32 | - 64 33 | - 96 34 | enc_out_1x1_chs: '2.0' 35 | dec_gru_iters: 2 36 | dec_gru_depth: 2 37 | dec_gru_mlp_ratio: 4.0 38 | dec_net_chs: null 39 | dec_inp_chs: null 40 | dec_motion_chs: 128 41 | use_upsample_mask: true 42 | upmask_gradient_scale: 1.0 43 | cache_pkconv_weights: false 44 | gamma: 0.8 45 | max_flow: 400 46 | warm_start: false 47 | data: 48 | train_dataset: chairs 49 | val_dataset: sintel-final-val+kitti-2015-val 50 | train_batch_size: 10 51 | train_num_workers: 4 52 | train_crop_size: [352, 480] 53 | train_transform_cuda: false 54 | train_transform_fp16: false -------------------------------------------------------------------------------- /ptlflow/models/rpknet/configs/rpknet-train2-things.yaml: -------------------------------------------------------------------------------- 1 | # lightning.pytorch==2.4.0 2 | seed_everything: true 3 | ckpt_path: /path/to/chairs.ckpt # Change to the ckpt resulting from rpknet-train1-chairs 4 | lr: 0.000125 5 | wdecay: 0.0001 6 | trainer: 7 | max_epochs: 80 8 | accumulate_grad_batches: 1 9 | gradient_clip_val: 1.0 10 | devices: 0, 11 | model: 12 | class_path: ptlflow.models.rpknet 13 | init_args: 14 | pyramid_ranges: 15 | - 32 16 | - 8 17 | iters: 12 18 | input_pad_one_side: false 19 | input_bgr_to_rgb: false 20 | detach_flow: true 21 | corr_mode: allpairs 22 | upgate_norm_type: group 23 | use_norm_affine: false 24 | group_norm_num_groups: 4 25 | corr_levels: 1 26 | corr_range: 4 27 | enc_norm_type: group 28 | enc_stem_stride: 2 29 | enc_depth: 2 30 | enc_mlp_ratio: 4.0 31 | enc_hidden_chs: 32 | - 32 33 | - 64 34 | - 96 35 | enc_out_1x1_chs: '2.0' 36 | dec_gru_iters: 2 37 | dec_gru_depth: 2 38 | dec_gru_mlp_ratio: 4.0 39 | dec_net_chs: null 40 | dec_inp_chs: null 41 | dec_motion_chs: 128 42 | use_upsample_mask: true 43 | upmask_gradient_scale: 1.0 44 | cache_pkconv_weights: false 45 | gamma: 0.8 46 | max_flow: 400 47 | warm_start: false 48 | data: 49 | train_dataset: things 50 | val_dataset: sintel-final-val+kitti-2015-val 51 | train_batch_size: 6 52 | train_num_workers: 3 53 | train_crop_size: [384, 704] 54 | train_transform_cuda: false 55 | train_transform_fp16: false -------------------------------------------------------------------------------- /ptlflow/models/rpknet/configs/rpknet-train3-sintel.yaml: -------------------------------------------------------------------------------- 1 | # lightning.pytorch==2.4.0 2 | seed_everything: true 3 | ckpt_path: /path/to/things.ckpt # Change to the ckpt resulting from rpknet-train2-things 4 | lr: 0.000125 5 | wdecay: 0.00001 6 | trainer: 7 | max_epochs: 5 8 | accumulate_grad_batches: 1 9 | gradient_clip_val: 1.0 10 | devices: 0, 11 | model: 12 | class_path: ptlflow.models.rpknet 13 | init_args: 14 | pyramid_ranges: 15 | - 32 16 | - 8 17 | iters: 12 18 | input_pad_one_side: false 19 | input_bgr_to_rgb: false 20 | detach_flow: true 21 | corr_mode: allpairs 22 | upgate_norm_type: group 23 | use_norm_affine: false 24 | group_norm_num_groups: 4 25 | corr_levels: 1 26 | corr_range: 4 27 | enc_norm_type: group 28 | enc_stem_stride: 2 29 | enc_depth: 2 30 | enc_mlp_ratio: 4.0 31 | enc_hidden_chs: 32 | - 32 33 | - 64 34 | - 96 35 | enc_out_1x1_chs: '2.0' 36 | dec_gru_iters: 2 37 | dec_gru_depth: 2 38 | dec_gru_mlp_ratio: 4.0 39 | dec_net_chs: null 40 | dec_inp_chs: null 41 | dec_motion_chs: 128 42 | use_upsample_mask: true 43 | upmask_gradient_scale: 1.0 44 | cache_pkconv_weights: false 45 | gamma: 0.85 46 | max_flow: 400 47 | warm_start: false 48 | data: 49 | train_dataset: sintel_finetune 50 | val_dataset: sintel-final-val+kitti-2015-val 51 | train_batch_size: 6 52 | train_num_workers: 3 53 | train_crop_size: [352, 768] 54 | train_transform_cuda: false 55 | train_transform_fp16: false -------------------------------------------------------------------------------- /ptlflow/models/rpknet/configs/rpknet-train4-kitti.yaml: -------------------------------------------------------------------------------- 1 | # lightning.pytorch==2.4.0 2 | seed_everything: true 3 | ckpt_path: /path/to/sintel.ckpt # Change to the ckpt resulting from rpknet-train3-sintel 4 | lr: 0.000125 5 | wdecay: 0.00001 6 | trainer: 7 | max_epochs: 150 8 | check_val_every_n_epoch: 10 9 | accumulate_grad_batches: 1 10 | gradient_clip_val: 1.0 11 | devices: 0, 12 | model: 13 | class_path: ptlflow.models.rpknet 14 | init_args: 15 | pyramid_ranges: 16 | - 32 17 | - 8 18 | iters: 12 19 | input_pad_one_side: false 20 | input_bgr_to_rgb: false 21 | detach_flow: true 22 | corr_mode: allpairs 23 | upgate_norm_type: group 24 | use_norm_affine: false 25 | group_norm_num_groups: 4 26 | corr_levels: 1 27 | corr_range: 4 28 | enc_norm_type: group 29 | enc_stem_stride: 2 30 | enc_depth: 2 31 | enc_mlp_ratio: 4.0 32 | enc_hidden_chs: 33 | - 32 34 | - 64 35 | - 96 36 | enc_out_1x1_chs: '2.0' 37 | dec_gru_iters: 2 38 | dec_gru_depth: 2 39 | dec_gru_mlp_ratio: 4.0 40 | dec_net_chs: null 41 | dec_inp_chs: null 42 | dec_motion_chs: 128 43 | use_upsample_mask: true 44 | upmask_gradient_scale: 1.0 45 | cache_pkconv_weights: false 46 | gamma: 0.85 47 | max_flow: 400 48 | warm_start: false 49 | data: 50 | train_dataset: kitti-2015 51 | val_dataset: kitti-2015-val+sintel-final-val 52 | train_batch_size: 6 53 | train_num_workers: 3 54 | train_crop_size: [288, 960] 55 | train_transform_cuda: false 56 | train_transform_fp16: false -------------------------------------------------------------------------------- /ptlflow/models/rpknet/configs/rpknet-validate-things.yaml: -------------------------------------------------------------------------------- 1 | # lightning.pytorch==2.4.0 2 | all: false 3 | select: null 4 | exclude: null 5 | ckpt_path: things 6 | output_path: outputs/validate 7 | write_outputs: false 8 | show: false 9 | flow_format: original 10 | max_forward_side: null 11 | scale_factor: null 12 | max_show_side: 1000 13 | max_samples: null 14 | reversed: false 15 | fp16: false 16 | seq_val_mode: all 17 | write_individual_metrics: false 18 | epe_clip: 5.0 19 | seed_everything: true 20 | model: 21 | class_path: ptlflow.models.rpknet 22 | init_args: 23 | pyramid_ranges: 24 | - 32 25 | - 8 26 | iters: 12 27 | input_pad_one_side: false 28 | input_bgr_to_rgb: false 29 | detach_flow: true 30 | corr_mode: allpairs 31 | upgate_norm_type: group 32 | use_norm_affine: false 33 | group_norm_num_groups: 4 34 | corr_levels: 1 35 | corr_range: 4 36 | enc_norm_type: group 37 | enc_stem_stride: 2 38 | enc_depth: 2 39 | enc_mlp_ratio: 4.0 40 | enc_hidden_chs: 41 | - 32 42 | - 64 43 | - 96 44 | enc_out_1x1_chs: '2.0' 45 | dec_gru_iters: 2 46 | dec_gru_depth: 2 47 | dec_gru_mlp_ratio: 4.0 48 | dec_net_chs: null 49 | dec_inp_chs: null 50 | dec_motion_chs: 128 51 | use_upsample_mask: true 52 | upmask_gradient_scale: 1.0 53 | cache_pkconv_weights: true 54 | gamma: 0.8 55 | max_flow: 400 56 | lr: null 57 | wdecay: null 58 | warm_start: false 59 | data: 60 | val_dataset: sintel-clean+sintel-final+kitti-2012+kitti-2015 61 | -------------------------------------------------------------------------------- /ptlflow/models/rpknet/local_timm/README.md: -------------------------------------------------------------------------------- 1 | # timm 2 | 3 | The code in this folder has been borrowed and adapted from the timm library: [https://github.com/huggingface/pytorch-image-models](https://github.com/huggingface/pytorch-image-models). -------------------------------------------------------------------------------- /ptlflow/models/rpknet/local_timm/gelu.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | 5 | 6 | class GELU(nn.Module): 7 | """Applies the Gaussian Error Linear Units function (w/ dummy inplace arg)""" 8 | 9 | def __init__(self, inplace: bool = False): 10 | super(GELU, self).__init__() 11 | 12 | def forward(self, input: torch.Tensor) -> torch.Tensor: 13 | return F.gelu(input) 14 | -------------------------------------------------------------------------------- /ptlflow/models/rpknet/local_timm/layer_helpers.py: -------------------------------------------------------------------------------- 1 | """ Layer/Module Helpers 2 | 3 | Hacked together by / Copyright 2020 Ross Wightman 4 | """ 5 | 6 | from itertools import repeat 7 | import collections.abc 8 | 9 | 10 | # From PyTorch internals 11 | def _ntuple(n): 12 | def parse(x): 13 | if isinstance(x, collections.abc.Iterable) and not isinstance(x, str): 14 | return x 15 | return tuple(repeat(x, n)) 16 | 17 | return parse 18 | 19 | 20 | to_1tuple = _ntuple(1) 21 | to_2tuple = _ntuple(2) 22 | to_3tuple = _ntuple(3) 23 | to_4tuple = _ntuple(4) 24 | to_ntuple = _ntuple 25 | 26 | 27 | def make_divisible(v, divisor=8, min_value=None, round_limit=0.9): 28 | min_value = min_value or divisor 29 | new_v = max(min_value, int(v + divisor / 2) // divisor * divisor) 30 | # Make sure that round down does not go down by more than 10%. 31 | if new_v < round_limit * v: 32 | new_v += divisor 33 | return new_v 34 | 35 | 36 | def extend_tuple(x, n): 37 | # pdas a tuple to specified n by padding with last value 38 | if not isinstance(x, (tuple, list)): 39 | x = (x,) 40 | else: 41 | x = tuple(x) 42 | pad_n = n - len(x) 43 | if pad_n <= 0: 44 | return x[:n] 45 | return x + (x[-1],) * pad_n 46 | -------------------------------------------------------------------------------- /ptlflow/models/rpknet/pwc_modules.py: -------------------------------------------------------------------------------- 1 | # ============================================================================= 2 | # Copyright 2023 Henrique Morimitsu 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | # Adapted from IRR: https://github.com/visinf/irr/blob/master/models/pwc_modules.py 17 | # ============================================================================= 18 | 19 | import torch 20 | import torch.nn.functional as F 21 | 22 | 23 | def upsample2d_as(inputs, target_as, mode="bilinear"): 24 | _, _, h, w = target_as.size() 25 | if inputs.shape[-2] != h or inputs.shape[-1] != w: 26 | inputs = F.interpolate(inputs, [h, w], mode=mode, align_corners=True) 27 | return inputs 28 | 29 | 30 | def rescale_flow(flow, width_im, height_im, to_local=True): 31 | if to_local: 32 | u_scale = float(flow.size(3) / width_im) 33 | v_scale = float(flow.size(2) / height_im) 34 | else: 35 | u_scale = float(width_im / flow.size(3)) 36 | v_scale = float(height_im / flow.size(2)) 37 | 38 | u, v = flow.chunk(2, dim=1) 39 | u = u * u_scale 40 | v = v * v_scale 41 | 42 | return torch.cat([u, v], dim=1) 43 | -------------------------------------------------------------------------------- /ptlflow/models/scopeflow/README.md: -------------------------------------------------------------------------------- 1 | # ScopeFlow 2 | 3 | ## Original code 4 | 5 | [https://github.com/avirambh/ScopeFlow](https://github.com/avirambh/ScopeFlow) 6 | 7 | ## Code license 8 | 9 | See [LICENSE](LICENSE). 10 | 11 | ## Pretrained weights license 12 | 13 | Not specified. 14 | 15 | ## Citation 16 | 17 | ``` 18 | @InProceedings{Bar-Haim_2020_CVPR, 19 | author = {Bar-Haim, Aviram and Wolf, Lior}, 20 | title = {ScopeFlow: Dynamic Scene Scoping for Optical Flow}, 21 | booktitle = {The IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)}, 22 | month = {June}, 23 | year = {2020} 24 | } 25 | ``` -------------------------------------------------------------------------------- /ptlflow/models/scopeflow/__init__.py: -------------------------------------------------------------------------------- 1 | from .irr_pwc_v2 import * 2 | -------------------------------------------------------------------------------- /ptlflow/models/scv/LICENSE: -------------------------------------------------------------------------------- 1 | 2 | DO WHAT THE FUCK YOU WANT TO PUBLIC LICENSE 3 | Version 2, December 2004 4 | 5 | Copyright (C) 2004 Sam Hocevar 6 | 7 | Everyone is permitted to copy and distribute verbatim or modified 8 | copies of this license document, and changing it is allowed as long 9 | as the name is changed. 10 | 11 | DO WHAT THE FUCK YOU WANT TO PUBLIC LICENSE 12 | TERMS AND CONDITIONS FOR COPYING, DISTRIBUTION AND MODIFICATION 13 | 14 | 0. You just DO WHAT THE FUCK YOU WANT TO. -------------------------------------------------------------------------------- /ptlflow/models/scv/README.md: -------------------------------------------------------------------------------- 1 | # SCV 2 | 3 | ## Original code 4 | 5 | [https://github.com/zacjiang/SCV](https://github.com/zacjiang/SCV) 6 | 7 | ## Additional requirements 8 | 9 | In order to use SCV you need to install two additional requirements: 10 | 11 | - torch_scatter: Check [https://github.com/rusty1s/pytorch_scatter](https://github.com/rusty1s/pytorch_scatter) for installation instructions. 12 | - faiss: install with `pip install faiss-gpu` 13 | 14 | ## Code license 15 | 16 | See [LICENSE](LICENSE). 17 | 18 | ## Pretrained weights license 19 | 20 | Not specified. 21 | 22 | ## Citation 23 | 24 | ``` 25 | @inproceedings{jiang2021learning, 26 | title={Learning Optical Flow from a Few Matches}, 27 | author={Jiang, Shihao and Lu, Yao and Li, Hongdong and Hartley, Richard}, 28 | booktitle={CVPR}, 29 | pages={16592--16600}, 30 | year={2021} 31 | } 32 | ``` -------------------------------------------------------------------------------- /ptlflow/models/scv/__init__.py: -------------------------------------------------------------------------------- 1 | from .scv import * 2 | -------------------------------------------------------------------------------- /ptlflow/models/sea_raft/LICENSE: -------------------------------------------------------------------------------- 1 | BSD 3-Clause License 2 | 3 | Copyright (c) 2024, Princeton Vision & Learning Lab 4 | 5 | Redistribution and use in source and binary forms, with or without 6 | modification, are permitted provided that the following conditions are met: 7 | 8 | 1. Redistributions of source code must retain the above copyright notice, this 9 | list of conditions and the following disclaimer. 10 | 11 | 2. Redistributions in binary form must reproduce the above copyright notice, 12 | this list of conditions and the following disclaimer in the documentation 13 | and/or other materials provided with the distribution. 14 | 15 | 3. Neither the name of the copyright holder nor the names of its 16 | contributors may be used to endorse or promote products derived from 17 | this software without specific prior written permission. 18 | 19 | THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" 20 | AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 21 | IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE 22 | DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE 23 | FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL 24 | DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR 25 | SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER 26 | CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, 27 | OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 28 | OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. -------------------------------------------------------------------------------- /ptlflow/models/sea_raft/README.md: -------------------------------------------------------------------------------- 1 | # SEA-RAFT 2 | 3 | ## Original code 4 | 5 | [https://github.com/princeton-vl/SEA-RAFT](https://github.com/princeton-vl/SEA-RAFT) 6 | 7 | ## Code license 8 | 9 | See [LICENSE](LICENSE). 10 | 11 | ## Pretrained weights license 12 | 13 | Not specified. 14 | 15 | ## Citation 16 | 17 | ``` 18 | @article{wang2024sea, 19 | title={SEA-RAFT: Simple, Efficient, Accurate RAFT for Optical Flow}, 20 | author={Wang, Yihan and Lipson, Lahav and Deng, Jia}, 21 | journal={arXiv preprint arXiv:2405.14793}, 22 | year={2024} 23 | } 24 | ``` -------------------------------------------------------------------------------- /ptlflow/models/sea_raft/__init__.py: -------------------------------------------------------------------------------- 1 | from .sea_raft import * 2 | -------------------------------------------------------------------------------- /ptlflow/models/separableflow/LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2022 Feihu Zhang 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. -------------------------------------------------------------------------------- /ptlflow/models/separableflow/README.md: -------------------------------------------------------------------------------- 1 | # SeparableFlow 2 | 3 | ## Original code 4 | 5 | [https://github.com/feihuzhang/SeparableFlow](https://github.com/feihuzhang/SeparableFlow) 6 | 7 | ## Additional requirements 8 | 9 | In order to use SeparableFlow you need to have CUDA installed and compile the GANet package: 10 | 11 | - Download and install CUDA from [https://developer.nvidia.com/cuda-downloads](https://developer.nvidia.com/cuda-downloads) 12 | - IMPORTANT! Be sure to choose the same CUDA version as your PyTorch 13 | - Compile the package: 14 | ```bash 15 | bash compile.sh 16 | ``` 17 | 18 | ## Code license 19 | 20 | See [LICENSE](LICENSE). 21 | 22 | ## Pretrained weights license 23 | 24 | Not specified. 25 | 26 | ## Citation 27 | 28 | ``` 29 | @inproceedings{Zhang2021SepFlow, 30 | title={Separable Flow: Learning Motion Cost Volumes for Optical Flow Estimation}, 31 | author={Zhang, Feihu and Woodford, Oliver J. and Prisacariu, Victor Adrian and Torr, Philip H.S.}, 32 | booktitle={Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV)}, 33 | year={2021} 34 | pages={10807-10817} 35 | } 36 | ``` -------------------------------------------------------------------------------- /ptlflow/models/separableflow/__init__.py: -------------------------------------------------------------------------------- 1 | from .separableflow import * 2 | -------------------------------------------------------------------------------- /ptlflow/models/separableflow/compile.sh: -------------------------------------------------------------------------------- 1 | #export LD_LIBRARY_PATH="/home/feihu/anaconda3/lib:$LD_LIBRARY_PATH" 2 | #export LD_INCLUDE_PATH="/home/feihu/anaconda3/include:$LD_INCLUDE_PATH" 3 | #export CUDA_HOME="/usr/local/cuda-10.0" 4 | #export PATH="/home/feihu/anaconda3/bin:/usr/local/cuda-10.0/bin:$PATH" 5 | #export CPATH="/usr/local/cuda-10.0/include" 6 | #export CUDNN_INCLUDE_DIR="/usr/local/cuda-10.0/include" 7 | #export CUDNN_LIB_DIR="/usr/local/cuda-10.0/lib64" 8 | 9 | TORCH=$(python -c "import os; import torch; print(os.path.dirname(torch.__file__))") 10 | #echo $TORCH 11 | cd libs/GANet 12 | python setup.py clean 13 | rm -rf build 14 | python setup.py build 15 | cp -r build/lib* build/lib -------------------------------------------------------------------------------- /ptlflow/models/separableflow/libs/GANet/functions/__init__.py: -------------------------------------------------------------------------------- 1 | from .GANet import * 2 | -------------------------------------------------------------------------------- /ptlflow/models/separableflow/libs/GANet/modules/__init__.py: -------------------------------------------------------------------------------- 1 | from .GANet import * 2 | -------------------------------------------------------------------------------- /ptlflow/models/separableflow/libs/GANet/setup.py: -------------------------------------------------------------------------------- 1 | from setuptools import setup 2 | from torch.utils.cpp_extension import CppExtension, BuildExtension, CUDAExtension 3 | 4 | 5 | setup( 6 | name="GANet", 7 | ext_modules=[ 8 | CUDAExtension( 9 | "GANet", 10 | [ 11 | "src/GANet_cuda.cpp", 12 | "src/GANet_kernel_share.cu", 13 | "src/NLF_kernel.cu", 14 | ], 15 | ) 16 | ], 17 | cmdclass={"build_ext": BuildExtension}, 18 | ) 19 | -------------------------------------------------------------------------------- /ptlflow/models/skflow/LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2022 sunshk1227 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. -------------------------------------------------------------------------------- /ptlflow/models/skflow/README.md: -------------------------------------------------------------------------------- 1 | # SKFlow 2 | 3 | ## Original code 4 | 5 | [https://github.com/littlespray/SKFlow](https://github.com/littlespray/SKFlow) 6 | 7 | ## Code license 8 | 9 | See [LICENSE](LICENSE). 10 | 11 | ## Pretrained weights license 12 | 13 | Not specified. 14 | 15 | ## Citation 16 | 17 | ``` 18 | 19 | @InProceedings{Sun2022SKFlowLearningOptical, 20 | author = {Shangkun Sun and Yuanqi Chen and Y. Zhu and Guodong Guo and Gezhong Li}, 21 | booktitle = {NeurIPS}, 22 | title = {SKFlow: Learning Optical Flow with Super Kernels}, 23 | year = {2022}, 24 | } 25 | ``` -------------------------------------------------------------------------------- /ptlflow/models/skflow/__init__.py: -------------------------------------------------------------------------------- 1 | from .skflow import * 2 | -------------------------------------------------------------------------------- /ptlflow/models/splatflow/LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2024, Bo Wang 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. -------------------------------------------------------------------------------- /ptlflow/models/splatflow/README.md: -------------------------------------------------------------------------------- 1 | # SplatFlow 2 | 3 | ## Original code 4 | 5 | [https://github.com/wwsource/SplatFlow](https://github.com/wwsource/SplatFlow)## Additional requirements 6 | 7 | In order to use SplatFlow you need to install cupy: 8 | 9 | ```bash 10 | pip install cupy-cuda12x 11 | ``` 12 | NOTE: replace "12" by the CUDA version of your PyTorch installation. 13 | 14 | ## Code license 15 | 16 | See [LICENSE](LICENSE). 17 | 18 | ## Pretrained weights license 19 | 20 | Not specified. 21 | 22 | ## Citation 23 | 24 | ``` 25 | @article{wang2024splatflow, 26 | title={SplatFlow: Learning Multi-frame Optical Flow via Splatting}, 27 | author={Wang, Bo and Zhang, Yifan and Li, Jian and Yu, Yang and Sun, Zhenping and Liu, Li and Hu, Dewen}, 28 | journal={International Journal of Computer Vision}, 29 | pages={1--23}, 30 | year={2024}, 31 | publisher={Springer} 32 | } 33 | ``` -------------------------------------------------------------------------------- /ptlflow/models/splatflow/__init__.py: -------------------------------------------------------------------------------- 1 | from .splatflow import * 2 | -------------------------------------------------------------------------------- /ptlflow/models/splatflow/attention.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import nn, einsum 3 | from einops import rearrange 4 | 5 | 6 | class Attention(nn.Module): 7 | def __init__( 8 | self, 9 | *, 10 | dim, 11 | heads=4, 12 | dim_head=128, 13 | ): 14 | super().__init__() 15 | self.heads = heads 16 | self.scale = dim_head**-0.5 17 | inner_dim = heads * dim_head 18 | 19 | self.to_qk = nn.Conv2d(dim, inner_dim * 2, 1, bias=False) 20 | 21 | def forward(self, fmap): 22 | heads, b, c, h, w = self.heads, *fmap.shape 23 | 24 | q, k = self.to_qk(fmap).chunk(2, dim=1) 25 | 26 | q, k = map(lambda t: rearrange(t, "b (h d) x y -> b h x y d", h=heads), (q, k)) 27 | q = self.scale * q 28 | 29 | sim = einsum("b h x y d, b h u v d -> b h x y u v", q, k) 30 | 31 | sim = rearrange(sim, "b h x y u v -> b h (x y) (u v)") 32 | attn = sim.softmax(dim=-1) 33 | 34 | return attn 35 | 36 | 37 | class Aggregate(nn.Module): 38 | def __init__( 39 | self, 40 | dim, 41 | heads=4, 42 | dim_head=128, 43 | ): 44 | super().__init__() 45 | self.heads = heads 46 | self.scale = dim_head**-0.5 47 | inner_dim = heads * dim_head 48 | 49 | self.to_v = nn.Conv2d(dim, inner_dim, 1, bias=False) 50 | 51 | self.gamma = nn.Parameter(torch.zeros(1)) 52 | 53 | if dim != inner_dim: 54 | self.project = nn.Conv2d(inner_dim, dim, 1, bias=False) 55 | else: 56 | self.project = None 57 | 58 | def forward(self, attn, fmap): 59 | heads, b, c, h, w = self.heads, *fmap.shape 60 | 61 | v = self.to_v(fmap) 62 | v = rearrange(v, "b (h d) x y -> b h (x y) d", h=heads) 63 | out = einsum("b h i j, b h j d -> b h i d", attn, v) 64 | out = rearrange(out, "b h (x y) d -> b (h d) x y", x=h, y=w) 65 | 66 | if self.project is not None: 67 | out = self.project(out) 68 | 69 | out = fmap + self.gamma * out 70 | 71 | return out 72 | -------------------------------------------------------------------------------- /ptlflow/models/splatflow/util.py: -------------------------------------------------------------------------------- 1 | import torch.nn.functional as F 2 | 3 | 4 | class InputPadder: 5 | def __init__(self, dims, base=8): 6 | self.ht, self.wd = dims[-2:] 7 | pad_ht = (((self.ht // base) + 1) * base - self.ht) % base 8 | pad_wd = (((self.wd // base) + 1) * base - self.wd) % base 9 | self._pad = [0, pad_wd, 0, pad_ht] 10 | 11 | def pad(self, *inputs): 12 | outputs = [] 13 | 14 | for x in inputs: 15 | bhw_mode = 0 16 | 17 | if len(x.shape) == 3: 18 | bhw_mode = 1 19 | x = x.unsqueeze(1) 20 | 21 | x = F.pad(x, self._pad, mode="replicate") 22 | if bhw_mode: 23 | x = x.squeeze(1) 24 | outputs.append(x) 25 | 26 | return outputs 27 | 28 | def unpad(self, *inputs): 29 | ht, wd = inputs[0].shape[-2:] 30 | c = [self._pad[2], ht - self._pad[3], self._pad[0], wd - self._pad[1]] 31 | 32 | return [x[..., c[0] : c[1], c[2] : c[3]] for x in inputs] 33 | -------------------------------------------------------------------------------- /ptlflow/models/starflow/README.md: -------------------------------------------------------------------------------- 1 | # StarFlow 2 | 3 | ## Original code 4 | 5 | [https://github.com/pgodet/star_flow](https://github.com/pgodet/star_flow) 6 | 7 | ## Code license 8 | 9 | See [LICENSE](LICENSE). 10 | 11 | ## Pretrained weights license 12 | 13 | Not specified. 14 | 15 | ## Citation 16 | 17 | ``` 18 | @article{godet2020starflow, 19 | title={STaRFlow: A SpatioTemporal Recurrent Cell for Lightweight Multi-Frame Optical Flow Estimation}, 20 | author={Godet, Pierre and Boulch, Alexandre and Plyer, Aur{\'e}lien and Le Besnerais, Guy}, 21 | journal={arXiv preprint arXiv:2007.05481}, 22 | year={2020} 23 | } 24 | ``` 25 | -------------------------------------------------------------------------------- /ptlflow/models/starflow/__init__.py: -------------------------------------------------------------------------------- 1 | from .starflow import * 2 | -------------------------------------------------------------------------------- /ptlflow/models/streamflow/LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2023 sunshk1227 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. -------------------------------------------------------------------------------- /ptlflow/models/streamflow/README.md: -------------------------------------------------------------------------------- 1 | # StreamFlow 2 | 3 | ## Original code 4 | 5 | [https://github.com/littlespray/StreamFlow](https://github.com/littlespray/StreamFlow) 6 | 7 | ## Usage 8 | 9 | This model uses four frames as input for evaluation. To evaluate StreamFlow on MPI-Sintel, use the command: 10 | 11 | ```bash 12 | python validate.py --model streamflow --ckpt things --data.val_dataset sintel-clean-seqlen_4-seqpos_all+sintel-final-seqlen_4-seqpos_all 13 | ``` 14 | 15 | ## Code license 16 | 17 | See [LICENSE](LICENSE). 18 | 19 | ## Pretrained weights license 20 | 21 | Not specified. 22 | 23 | ## Citation 24 | 25 | ``` 26 | @inproceedings{sun2023streamflowstreamlinedmultiframeoptical, 27 | title={StreamFlow: Streamlined Multi-Frame Optical Flow Estimation for Video Sequences}, 28 | author={Shangkun Sun and Jiaming Liu and Thomas H. Li and Huaxia Li and Guoqing Liu and Wei Gao}, 29 | year={2024}, 30 | booktitle={NeurIPS}, 31 | } 32 | ``` -------------------------------------------------------------------------------- /ptlflow/models/streamflow/__init__.py: -------------------------------------------------------------------------------- 1 | from .streamflow import * 2 | -------------------------------------------------------------------------------- /ptlflow/models/unimatch/LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2022 autonomousvision 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. -------------------------------------------------------------------------------- /ptlflow/models/unimatch/README.md: -------------------------------------------------------------------------------- 1 | # UniMatch 2 | 3 | ## Original code 4 | 5 | [https://github.com/autonomousvision/unimatch](https://github.com/autonomousvision/unimatch) 6 | 7 | ## Code license 8 | 9 | See [LICENSE](LICENSE). 10 | 11 | ## Pretrained weights license 12 | 13 | Not specified. 14 | 15 | ## Citation 16 | 17 | ``` 18 | @article{xu2022unifying, 19 | title={Unifying Flow, Stereo and Depth Estimation}, 20 | author={Xu, Haofei and Zhang, Jing and Cai, Jianfei and Rezatofighi, Hamid and Yu, Fisher and Tao, Dacheng and Geiger, Andreas}, 21 | journal={arXiv preprint arXiv:2211.05783}, 22 | year={2022} 23 | } 24 | ``` -------------------------------------------------------------------------------- /ptlflow/models/unimatch/__init__.py: -------------------------------------------------------------------------------- 1 | from .unimatch import * 2 | -------------------------------------------------------------------------------- /ptlflow/models/vcn/LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2019 Carnegie Mellon University 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. -------------------------------------------------------------------------------- /ptlflow/models/vcn/README.md: -------------------------------------------------------------------------------- 1 | # VCN 2 | 3 | ## Original code 4 | 5 | [https://github.com/gengshan-y/VCN](https://github.com/gengshan-y/VCN) 6 | 7 | ## Code license 8 | 9 | See [LICENSE](LICENSE). 10 | 11 | ## Pretrained weights license 12 | 13 | Not specified. 14 | 15 | ## Citation 16 | 17 | ``` 18 | @inproceedings{yang2019vcn, 19 | title={Volumetric Correspondence Networks for Optical Flow}, 20 | author={Yang, Gengshan and Ramanan, Deva}, 21 | booktitle={NeurIPS}, 22 | year={2019} 23 | } 24 | ``` -------------------------------------------------------------------------------- /ptlflow/models/vcn/__init__.py: -------------------------------------------------------------------------------- 1 | from .vcn import * 2 | -------------------------------------------------------------------------------- /ptlflow/models/videoflow/Networks/BOFNet/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hmorimitsu/ptlflow/fb6d753c3b16e74481258891b45668dcba63e54e/ptlflow/models/videoflow/Networks/BOFNet/__init__.py -------------------------------------------------------------------------------- /ptlflow/models/videoflow/Networks/MOFNetStack/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hmorimitsu/ptlflow/fb6d753c3b16e74481258891b45668dcba63e54e/ptlflow/models/videoflow/Networks/MOFNetStack/__init__.py -------------------------------------------------------------------------------- /ptlflow/models/videoflow/Networks/__init__.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | 4 | def build_network(cfg): 5 | name = cfg.network 6 | if name == "MOFNetStack": 7 | from .MOFNetStack.network import MOFNet as network 8 | elif name == "BOFNet": 9 | from .BOFNet.network import BOFNet as network 10 | else: 11 | raise ValueError(f"Network = {name} is not a valid optimizer!") 12 | 13 | return network(cfg[name]) 14 | -------------------------------------------------------------------------------- /ptlflow/models/videoflow/README.md: -------------------------------------------------------------------------------- 1 | # VideoFlow 2 | 3 | ## Original code 4 | 5 | [https://github.com/XiaoyuShi97/VideoFlow](https://github.com/XiaoyuShi97/VideoFlow) 6 | 7 | ## Code license 8 | 9 | Apache License. 10 | 11 | ## Pretrained weights license 12 | 13 | Not specified. 14 | 15 | ## Citation 16 | 17 | ``` 18 | @inproceedings{shi2023videoflow, 19 | title={Videoflow: Exploiting temporal cues for multi-frame optical flow estimation}, 20 | author={Shi, Xiaoyu and Huang, Zhaoyang and Bian, Weikang and Li, Dasong and Zhang, Manyuan and Cheung, Ka Chun and See, Simon and Qin, Hongwei and Dai, Jifeng and Li, Hongsheng}, 21 | booktitle={ICCV}, 22 | year={2023} 23 | } 24 | ``` -------------------------------------------------------------------------------- /ptlflow/models/videoflow/__init__.py: -------------------------------------------------------------------------------- 1 | from .videoflow_bof import * 2 | from .videoflow_mof import * 3 | -------------------------------------------------------------------------------- /ptlflow/utils/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hmorimitsu/ptlflow/fb6d753c3b16e74481258891b45668dcba63e54e/ptlflow/utils/__init__.py -------------------------------------------------------------------------------- /ptlflow/utils/callbacks/RobotoMono-Regular.ttf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hmorimitsu/ptlflow/fb6d753c3b16e74481258891b45668dcba63e54e/ptlflow/utils/callbacks/RobotoMono-Regular.ttf -------------------------------------------------------------------------------- /ptlflow/utils/callbacks/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hmorimitsu/ptlflow/fb6d753c3b16e74481258891b45668dcba63e54e/ptlflow/utils/callbacks/__init__.py -------------------------------------------------------------------------------- /ptlflow/utils/external/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hmorimitsu/ptlflow/fb6d753c3b16e74481258891b45668dcba63e54e/ptlflow/utils/external/__init__.py -------------------------------------------------------------------------------- /ptlflow/utils/external/alt_cuda_corr/LICENSE: -------------------------------------------------------------------------------- 1 | BSD 3-Clause License 2 | 3 | Copyright (c) 2020, princeton-vl 4 | All rights reserved. 5 | 6 | Redistribution and use in source and binary forms, with or without 7 | modification, are permitted provided that the following conditions are met: 8 | 9 | * Redistributions of source code must retain the above copyright notice, this 10 | list of conditions and the following disclaimer. 11 | 12 | * Redistributions in binary form must reproduce the above copyright notice, 13 | this list of conditions and the following disclaimer in the documentation 14 | and/or other materials provided with the distribution. 15 | 16 | * Neither the name of the copyright holder nor the names of its 17 | contributors may be used to endorse or promote products derived from 18 | this software without specific prior written permission. 19 | 20 | THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" 21 | AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 22 | IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE 23 | DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE 24 | FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL 25 | DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR 26 | SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER 27 | CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, 28 | OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 29 | OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. -------------------------------------------------------------------------------- /ptlflow/utils/external/alt_cuda_corr/README.md: -------------------------------------------------------------------------------- 1 | # alt_cuda_corr 2 | 3 | CUDA implementation of local correlation calculation (cost volume). 4 | Originally implemented in RAFT to avoid computing a global cost volume. 5 | It decreases memory consumption, but increases running time. 6 | 7 | ## Installation instructions 8 | 9 | 1. Download and install CUDA from [https://developer.nvidia.com/cuda-downloads](https://developer.nvidia.com/cuda-downloads) 10 | - IMPORTANT! Be sure to choose the same CUDA version as your PyTorch 11 | 2. Enter this folder and then run the setup: 12 | ```bash 13 | cd ptlflow/utils/external/alt_cuda_corr/ 14 | python setup.py install 15 | ``` 16 | 17 | ## Original source 18 | 19 | [https://github.com/princeton-vl/RAFT/tree/master/alt_cuda_corr](https://github.com/princeton-vl/RAFT/tree/master/alt_cuda_corr) 20 | 21 | ## LICENSE 22 | 23 | See [LICENSE](LICENSE) -------------------------------------------------------------------------------- /ptlflow/utils/external/alt_cuda_corr/correlation.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | 4 | // CUDA forward declarations 5 | std::vector corr_cuda_forward( 6 | torch::Tensor fmap1, 7 | torch::Tensor fmap2, 8 | torch::Tensor coords, 9 | int radius); 10 | 11 | std::vector corr_cuda_backward( 12 | torch::Tensor fmap1, 13 | torch::Tensor fmap2, 14 | torch::Tensor coords, 15 | torch::Tensor corr_grad, 16 | int radius); 17 | 18 | // C++ interface 19 | #define CHECK_CUDA(x) TORCH_CHECK(x.type().is_cuda(), #x " must be a CUDA tensor") 20 | #define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous") 21 | #define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x) 22 | 23 | std::vector corr_forward( 24 | torch::Tensor fmap1, 25 | torch::Tensor fmap2, 26 | torch::Tensor coords, 27 | int radius) { 28 | CHECK_INPUT(fmap1); 29 | CHECK_INPUT(fmap2); 30 | CHECK_INPUT(coords); 31 | 32 | return corr_cuda_forward(fmap1, fmap2, coords, radius); 33 | } 34 | 35 | 36 | std::vector corr_backward( 37 | torch::Tensor fmap1, 38 | torch::Tensor fmap2, 39 | torch::Tensor coords, 40 | torch::Tensor corr_grad, 41 | int radius) { 42 | CHECK_INPUT(fmap1); 43 | CHECK_INPUT(fmap2); 44 | CHECK_INPUT(coords); 45 | CHECK_INPUT(corr_grad); 46 | 47 | return corr_cuda_backward(fmap1, fmap2, coords, corr_grad, radius); 48 | } 49 | 50 | 51 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 52 | m.def("forward", &corr_forward, "CORR forward"); 53 | m.def("backward", &corr_backward, "CORR backward"); 54 | } -------------------------------------------------------------------------------- /ptlflow/utils/external/alt_cuda_corr/setup.py: -------------------------------------------------------------------------------- 1 | from setuptools import setup 2 | from torch.utils.cpp_extension import BuildExtension, CUDAExtension 3 | 4 | 5 | setup( 6 | name="correlation", 7 | ext_modules=[ 8 | CUDAExtension( 9 | "alt_cuda_corr", 10 | sources=["correlation.cpp", "correlation_kernel.cu"], 11 | extra_compile_args={"cxx": [], "nvcc": ["-O3"]}, 12 | ), 13 | ], 14 | cmdclass={"build_ext": BuildExtension}, 15 | ) 16 | -------------------------------------------------------------------------------- /ptlflow/utils/registry.py: -------------------------------------------------------------------------------- 1 | # ============================================================================= 2 | # Copyright 2024 Henrique Morimitsu 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # ============================================================================= 16 | 17 | import sys 18 | 19 | import lightning.pytorch as pl 20 | 21 | from ptlflow.models.base_model.base_model import BaseModel 22 | 23 | 24 | _models_dict = {} 25 | _trainable_models = [] 26 | _ptlflow_trained_models = [] 27 | 28 | 29 | class RegisteredModel(pl.LightningModule): 30 | pass 31 | 32 | 33 | def register_model(model_class: BaseModel) -> BaseModel: 34 | # lookup containing module 35 | model_dir = ".".join(model_class.__module__.split(".")[:-1]) 36 | mod = sys.modules[model_dir] 37 | # add model to __all__ in module 38 | model_name = model_class.__name__ 39 | if hasattr(mod, "__all__"): 40 | mod.__all__.append(model_name) 41 | else: 42 | mod.__all__ = [model_name] # type: ignore 43 | 44 | _models_dict[model_class.__name__] = model_class 45 | registered_class = type(model_class.__name__, (model_class, RegisteredModel), {}) 46 | registered_class.__module__ = model_class.__module__ 47 | return registered_class 48 | 49 | 50 | def trainable(model_class: BaseModel) -> BaseModel: 51 | _trainable_models.append(model_class.__name__) 52 | return model_class 53 | 54 | 55 | def ptlflow_trained(model_class: BaseModel) -> BaseModel: 56 | _ptlflow_trained_models.append(model_class.__name__) 57 | return model_class 58 | -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- 1 | [build-system] 2 | requires = ["setuptools>=61.0"] 3 | build-backend = "setuptools.build_meta" 4 | 5 | [project] 6 | name = "ptlflow" 7 | version = "0.4.1" 8 | authors = [ 9 | {name = "Henrique Morimitsu", email = "henriquem87@gmail.com"}, 10 | ] 11 | description = "PyTorch Lightning Optical Flow" 12 | readme = "README.md" 13 | requires-python = ">=3.8,<3.13" 14 | keywords = ["torch", "pytorch lightning", "optical flow", "models"] 15 | license = {text = "Apache 2.0 License"} 16 | classifiers = [ 17 | "Development Status :: 3 - Alpha", 18 | "Intended Audience :: Education", 19 | "Intended Audience :: Science/Research", 20 | "License :: OSI Approved :: Apache Software License", 21 | "Programming Language :: Python :: 3", 22 | "Programming Language :: Python :: 3.10", 23 | "Programming Language :: Python :: 3.11", 24 | "Programming Language :: Python :: 3.12", 25 | "Topic :: Scientific/Engineering", 26 | "Topic :: Scientific/Engineering :: Artificial Intelligence", 27 | "Topic :: Software Development", 28 | "Topic :: Software Development :: Libraries", 29 | "Topic :: Software Development :: Libraries :: Python Modules", 30 | ] 31 | dependencies = [ 32 | "einops<0.9", 33 | "h5py<3.13", 34 | "kaleido<0.3", 35 | "lightning[pytorch-extra]>=2.1,<2.6", 36 | "loguru<0.8", 37 | "opencv-python<4.11", 38 | "pandas<2.3", 39 | "plotly<5.25", 40 | "pypng==0.20220715.0", 41 | "scipy<1.15", 42 | "tabulate<0.10", 43 | "tensorboard<2.19", 44 | "timm<1.1", 45 | ] 46 | 47 | [project.urls] 48 | "Source Code" = "https://github.com/hmorimitsu/ptlflow" 49 | "Documentation" = "https://ptlflow.readthedocs.io/" 50 | 51 | [tool.setuptools.packages.find] 52 | include = ["ptlflow*"] 53 | exclude = ["tests"] 54 | 55 | [options.package_data] 56 | "*" = ["*.md", "*.txt", "*LICENSE", "*.ttf"] -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | einops<0.9 2 | h5py<3.13 3 | kaleido<0.3 4 | lightning[pytorch-extra]>=2.1,<2.6 5 | loguru<0.8 6 | opencv-python<4.11 7 | pandas<2.3 8 | plotly<5.25 9 | pypng==0.20220715.0 10 | scipy<1.15 11 | tabulate<0.10 12 | tensorboard<2.19 13 | timm<1.1 -------------------------------------------------------------------------------- /tests/ptlflow/test_lightning.py: -------------------------------------------------------------------------------- 1 | # ============================================================================= 2 | # Copyright 2021 Henrique Morimitsu 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # ============================================================================= 16 | 17 | import sys 18 | 19 | from ptlflow.utils.lightning.ptlflow_cli import PTLFlowCLI 20 | from ptlflow.utils.lightning.ptlflow_checkpoint_connector import ( 21 | _PTLFlowCheckpointConnector, 22 | ) 23 | from ptlflow.utils.lightning.ptlflow_trainer import PTLFlowTrainer 24 | from ptlflow.utils.registry import RegisteredModel 25 | 26 | TEST_MODEL = "raft_small" 27 | 28 | 29 | def test_cli_no_model() -> None: 30 | sys.argv = sys.argv[:1] 31 | PTLFlowCLI( 32 | model_class=None, 33 | subclass_mode_model=False, 34 | parser_kwargs={"parents": []}, 35 | run=False, 36 | parse_only=False, 37 | auto_configure_optimizers=False, 38 | ) 39 | 40 | 41 | def test_cli() -> None: 42 | sys.argv = sys.argv[:1] 43 | sys.argv.extend(["--model", "rapidflow"]) 44 | PTLFlowCLI( 45 | model_class=RegisteredModel, 46 | subclass_mode_model=True, 47 | parser_kwargs={"parents": []}, 48 | run=False, 49 | parse_only=False, 50 | auto_configure_optimizers=False, 51 | ) 52 | 53 | 54 | def test_checkpoint_connector_and_trainer() -> None: 55 | _PTLFlowCheckpointConnector(PTLFlowTrainer()) 56 | -------------------------------------------------------------------------------- /tests/ptlflow/utils/test_flow_utils.py: -------------------------------------------------------------------------------- 1 | # ============================================================================= 2 | # Copyright 2021 Henrique Morimitsu 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # ============================================================================= 16 | 17 | from pathlib import Path 18 | import shutil 19 | 20 | import numpy as np 21 | 22 | from ptlflow.utils import flow_utils 23 | 24 | IMG_SIDE = 29 25 | IMG_MIDDLE = IMG_SIDE // 2 + 1 26 | 27 | 28 | def test_read_write_pfm(tmp_path: Path) -> None: 29 | flow = np.stack( 30 | np.meshgrid(np.arange(IMG_SIDE) - IMG_MIDDLE, np.arange(IMG_SIDE) - IMG_MIDDLE), 31 | axis=2, 32 | ).astype(np.float32) 33 | file_path = tmp_path / "flow.pfm" 34 | flow_utils.flow_write(file_path, flow) 35 | assert file_path.exists() 36 | 37 | loaded_flow = flow_utils.flow_read(file_path) 38 | assert np.array_equal(flow, loaded_flow) 39 | 40 | shutil.rmtree(tmp_path) 41 | -------------------------------------------------------------------------------- /tests/test_model_benchmark.py: -------------------------------------------------------------------------------- 1 | # ============================================================================= 2 | # Copyright 2021 Henrique Morimitsu 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # ============================================================================= 16 | 17 | from pathlib import Path 18 | import shutil 19 | 20 | from jsonargparse import ArgumentParser 21 | import ptlflow 22 | import model_benchmark 23 | 24 | TEST_MODEL = "raft_small" 25 | 26 | 27 | def test_benchmark(tmp_path: Path) -> None: 28 | model_ref = ptlflow.get_model_reference(TEST_MODEL) 29 | 30 | model_parser = ArgumentParser(parents=[model_benchmark._init_parser()]) 31 | model_parser.add_argument_group("model") 32 | model_parser.add_class_arguments(model_ref, "model.init_args") 33 | args = model_parser.parse_args([]) 34 | args.model.class_path = f"{model_ref.__module__}.{model_ref.__qualname__}" 35 | 36 | args.num_samples = 1 37 | args.output_path = tmp_path 38 | 39 | model_benchmark.benchmark(args, None) 40 | 41 | shutil.rmtree(tmp_path) 42 | -------------------------------------------------------------------------------- /tests/test_summary_metrics.py: -------------------------------------------------------------------------------- 1 | # ============================================================================= 2 | # Copyright 2021 Henrique Morimitsu 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # ============================================================================= 16 | 17 | from pathlib import Path 18 | import shutil 19 | 20 | import summary_metrics 21 | 22 | 23 | def test_summary(tmp_path: Path) -> None: 24 | parser = summary_metrics._init_parser() 25 | args = parser.parse_args([]) 26 | args.metrics_path = Path("docs/source/results/metrics_all.csv") 27 | args.output_dir = tmp_path 28 | summary_metrics.summarize(args) 29 | 30 | assert len(list(tmp_path.glob("**/*.md"))) > 0 31 | assert len(list(tmp_path.glob("**/*.csv"))) > 0 32 | 33 | shutil.rmtree(tmp_path) 34 | --------------------------------------------------------------------------------