Добавлены пропсы конвейера и стереодвижки, задействованные в прогоне
assets/conveyors (274 МБ) - ленты и угловая секция NVIDIA, на которые ссылается сцена относительным путём. Раньше исключались как перекачиваемые, но без них сцена не композится из коробки. cv/ - код стереодвижков, которые вызывает control_test, без весов: * defom-stereo - рабочий бейзлайн (DEFOM vitl, вход 480, iters 24) * crestereo - второй движок, точнее по габаритам (MAE 23.5 против 32.8 мм) * fast-foundationstereo - проверялся, в бейзлайн не вошёл * circular_section.py - показатель кругового сечения, перенесён в measure_plane.py: выравнивает облако по СОБСТВЕННЫМ главным осям и режет на пяти высотах вдоль каждой. Три самодельные версии (мировые оси, одно сечение) давали хуже; результаты проверки на эталонной геометрии - в circular_section_results.json Веса по-прежнему не в репозитории - источники в MODELS.md. Наборы кадров прежних прогонов (cv/flow_*, 1.26 ГБ) исключены: это выход, а не исходники. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
dasha_f
@@ -0,0 +1,21 @@
|
||||
MIT License
|
||||
|
||||
Copyright (c) 2024 Insta360 Research Team
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||
of this software and associated documentation files (the "Software"), to deal
|
||||
in the Software without restriction, including without limitation the rights
|
||||
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||
copies of the Software, and to permit persons to whom the Software is
|
||||
furnished to do so, subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in all
|
||||
copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||
SOFTWARE.
|
||||
@@ -0,0 +1,225 @@
|
||||
# DEFOM-Stereo [CVPR 2025]
|
||||
|
||||
The Official Pytorch Implementation for
|
||||
|
||||
> [**DEFOM-Stereo: Depth Foundation Model Based Stereo Matching**](https://arxiv.org/abs/2501.09466)
|
||||
>
|
||||
> Authors: Hualie Jiang, Zhiqiang Lou, Laiyan Ding, Rui Xu, Minglang Tan, Wenjie Jiang and Rui Huang
|
||||
|
||||
# Abstract
|
||||
Stereo matching is a key technique for metric depth estimation in computer vision and robotics.
|
||||
Real-world challenges like occlusion and non-texture hinder accurate disparity estimation from binocular matching cues. Recently, monocular relative depth estimation has shown remarkable generalization using vision foundation models. Thus, to facilitate robust stereo matching with monocular depth cues, we incorporate a robust monocular relative depth model into the recurrent stereo-matching framework, building a new framework for depth foundation model-based stereo-matching, DEFOM-Stereo.
|
||||
In the feature extraction stage, we construct the combined context and matching feature encoder by integrating features from conventional CNNs and DEFOM. In the update stage, we use the depth predicted by DEFOM to initialize the recurrent disparity and introduce a scale update module to refine the disparity at the correct scale. DEFOM-Stereo is verified to have much stronger zero-shot generalization compared with SOTA methods. Moreover, DEFOM-Stereo achieves top performance on the KITTI 2012, KITTI 2015, Middlebury, and ETH3D benchmarks, ranking $1^{st}$ on many metrics. In the joint evaluation under the robust vision challenge, our model simultaneously outperforms previous models on the individual benchmarks, further demonstrating its outstanding capabilities.
|
||||
|
||||
|
||||
# Pipeline
|
||||
<p align="center">
|
||||
<img src='assets/framework.svg' width=980>
|
||||
</p>
|
||||
|
||||
- We propose **a novel recurrent stereo-matching framework incorporating monocular depth cues** from a depth foundation model to improve robustness.
|
||||
- We develop a simple technique that utilizes **pre-trained DEFOM features** to construct stronger **combined feature and context encoders**.
|
||||
- We invent a **recurrent scale update** module empowered with the **scale lookup**, serving to recover accurate pixel-wise scales for the coarse DEFOM depth.
|
||||
|
||||
|
||||
|
||||
# Zero-Shot Perfomance
|
||||
|
||||
<p align="center">
|
||||
<img src='assets/zeroshot.jpg' width=1000>
|
||||
</p>
|
||||
|
||||
# Benchmark Performance
|
||||
|
||||
<p align="center">
|
||||
<img src='assets/benchmark.jpg' width=800>
|
||||
</p>
|
||||
|
||||
# Robust Vision Challange
|
||||
|
||||
<p align="center">
|
||||
<img src='assets/rvc.svg' width=1200>
|
||||
</p>
|
||||
|
||||
# Preparation
|
||||
|
||||
### Installation
|
||||
|
||||
Create the environment
|
||||
|
||||
```bash
|
||||
conda env create -f environment.yaml
|
||||
conda activate defomstereo
|
||||
pip install -r requirements.txt
|
||||
```
|
||||
|
||||
|
||||
### Datasets
|
||||
The project requires the follow datasets:
|
||||
|
||||
<table style="border-collapse: collapse; width: 80%;">
|
||||
<tr>
|
||||
<td style="border: 1px solid #000; padding: 8px;"><a class="custom-link" href="https://www.cvlibs.net/datasets/kitti/eval_stereo_flow.php?benchmark=stereo" target="_blank">KITTI-2012</a></td>
|
||||
<td style="border: 1px solid #000; padding: 8px;"><a class="custom-link" href="https://www.cvlibs.net/datasets/kitti/eval_scene_flow.php?benchmark=stereo" target="_blank">KITTI-2015</a></td>
|
||||
<td style="border: 1px solid #000; padding: 8px;"><a class="custom-link" href="https://vision.middlebury.edu/stereo/submit3/" target="_blank">Middlebury</a></td>
|
||||
<td style="border: 1px solid #000; padding: 8px;"><a class="custom-link" href="https://www.eth3d.net/datasets" target="_blank">ETH3D</a></td>
|
||||
<td style="border: 1px solid #000; padding: 8px;"><a class="custom-link" href="https://github.com/YuhuaXu/StereoDataset" target="_blank">InStereo2K</a></td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style="border: 1px solid #000; padding: 8px;"><a class="custom-link" href="https://europe.naverlabs.com/proxy-virtual-worlds-vkitti-2/" target="_blank">Virtual KITTI 2</a></td>
|
||||
<td style="border: 1px solid #000; padding: 8px;"><a class="custom-link" href="https://lmb.informatik.uni-freiburg.de/resources/datasets/SceneFlowDatasets.en.html" target="_blank">SceneFlow</a></td>
|
||||
<td style="border: 1px solid #000; padding: 8px;"><a class="custom-link" href="https://github.com/castacks/tartanair_tools" target="_blank">TartanAir</a>
|
||||
</td>
|
||||
<td style="border: 1px solid #000; padding: 8px;"><a class="custom-link" href="https://github.com/megvii-research/CREStereo" target="_blank">CREStereo Dataset</a>
|
||||
</td>
|
||||
<td style="border: 1px solid #000; padding: 8px;"><a class="custom-link" href="https://research.nvidia.com/publication/2018-06_falling-things-synthetic-dataset-3d-object-detection-and-pose-estimation" target="_blank">FallingThings</a>
|
||||
</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style="border: 1px solid #000; padding: 8px;"><a class="custom-link" href="http://sintel.is.tue.mpg.de/stereo" target="_blank">Sintel Stereo</a></td>
|
||||
<td style="border: 1px solid #000; padding: 8px;"><a class="custom-link" href="https://drive.google.com/file/d/1SgEIrH_IQTKJOToUwR1rx4-237sThUqX/view" target="_blank">HR-VS</a></td>
|
||||
<td style="border: 1px solid #000; padding: 8px;"><a class="custom-link" href="https://github.com/sniklaus/3d-ken-burns" target="_blank">3D Ken Burns</a></td>
|
||||
<td style="border: 1px solid #000; padding: 8px;"><a class="custom-link" href="https://github.com/HKBU-HPML/IRS" target="_blank">IRS Dataset</a></td>
|
||||
<td style="border: 1px solid #000; padding: 8px;"><a class="custom-link" href="https://cvlab-unibo.github.io/booster-web/" target="_blank">Booster Dataset</a></td>
|
||||
</tr>
|
||||
</table>
|
||||
The datasets are organized as follows,
|
||||
|
||||
```
|
||||
.
|
||||
└── datasets
|
||||
├── 3dkenburns
|
||||
│ ├── asdf-flying
|
||||
│ ├── asdf-flying-depth
|
||||
│ └── ...
|
||||
├── Booster_Dataset
|
||||
│ ├── test
|
||||
│ └── train
|
||||
├── CreStereo
|
||||
│ ├── hole
|
||||
│ ├── reflective
|
||||
│ ├── shapenet
|
||||
│ └── tree
|
||||
├── ETH3D
|
||||
│ ├── two_view_testing
|
||||
│ ├── two_view_training
|
||||
│ └── two_view_training_gt
|
||||
├── FallingThings
|
||||
│ └── fat
|
||||
├── HRVS
|
||||
│ └── carla-highres
|
||||
├── InStereo2K
|
||||
│ ├── part1
|
||||
│ ├── part2
|
||||
│ ├── part3
|
||||
│ ├── part4
|
||||
│ ├── part5
|
||||
│ └── test
|
||||
├── IRSDataset
|
||||
│ ├── Home
|
||||
│ ├── Office
|
||||
│ ├── Restaurant
|
||||
│ └── Store
|
||||
├── KITTI12
|
||||
│ ├── testing
|
||||
│ └── training
|
||||
├── KITTI15
|
||||
│ ├── testing
|
||||
│ └── training
|
||||
├── Middlebury
|
||||
│ ├── 2005
|
||||
│ ├── 2006
|
||||
│ ├── 2014
|
||||
│ ├── 2021
|
||||
│ └── MiddEval3
|
||||
├── SceneFlow
|
||||
│ ├── Driving
|
||||
│ ├── FlyingThings3D
|
||||
│ └── Monkaa
|
||||
├── SintelStereo
|
||||
│ └── training
|
||||
├── TartanAir
|
||||
│ ├── abandonedfactory
|
||||
│ ├── abandonedfactory_night
|
||||
│ └── ...
|
||||
└── VKITTI2
|
||||
├── Scene01
|
||||
├── Scene02
|
||||
├── Scene06
|
||||
├── Scene18
|
||||
└── Scene20
|
||||
```
|
||||
|
||||
|
||||
# Evaluation
|
||||
|
||||
### Download the pre-trained models
|
||||
```
|
||||
bash scripts/download_models.sh
|
||||
```
|
||||
|
||||
The pretrained models are available on [Google Drive](https://drive.google.com/drive/folders/1cZLcIjHlmUo986gkR6FbofG1cj5BT36x?usp=sharing) and can be downloaded mamanually.
|
||||
|
||||
### Perfom Evaluation
|
||||
```
|
||||
bash scripts/evaluate.sh
|
||||
```
|
||||
|
||||
# Make Benchmark Submission
|
||||
|
||||
```
|
||||
bash scripts/make_submission.sh
|
||||
```
|
||||
|
||||
|
||||
# Training
|
||||
|
||||
### Download DaV2 models
|
||||
```
|
||||
bash scripts/download_dav2.sh
|
||||
```
|
||||
|
||||
|
||||
### Train on SceneFlow
|
||||
|
||||
```
|
||||
bash scripts/train_sceneflow_vits.sh
|
||||
bash scripts/train_sceneflow_vitl.sh
|
||||
```
|
||||
|
||||
### Finetune for Benchmarks
|
||||
|
||||
```
|
||||
bash scripts/train_kitti.sh
|
||||
bash scripts/train_middlebury.sh
|
||||
bash scripts/train_eth3d.sh
|
||||
bash scripts/train_rvc.sh
|
||||
```
|
||||
|
||||
|
||||
# Domo on real samples
|
||||
|
||||
```
|
||||
python demo.py --restore_ckpt checkpoints/defomstereo_vitl_sceneflow.pth
|
||||
```
|
||||
|
||||
|
||||
|
||||
|
||||
## Acknowledgements
|
||||
|
||||
The project is based on [RAFT-Stereo](https://github.com/princeton-vl/RAFT-Stereo) and [Depth Anything V2](https://github.com/DepthAnything/Depth-Anything-V2) and we sincerely acknowledge their authors for opensourcing the excellent work. Besides, we would like to thank the CVPR reviewers and AC for their valuable feedback and recognition of our work.
|
||||
|
||||
|
||||
## Citation
|
||||
|
||||
Please cite our paper if you find our work useful in your research.
|
||||
|
||||
```
|
||||
@inproceedings{jiang2025defom,
|
||||
title={DEFOM-Stereo: Depth Foundation Model Based Stereo Matching},
|
||||
author={Jiang, Hualie and Lou, Zhiqiang and Ding, Laiyan and Xu, Rui and Tan, Minglang and Jiang, Wenjie and Huang, Rui},
|
||||
booktitle={IEEE International Conference on Computer Vision and Pattern Recognition (CVPR)},
|
||||
year={2025}
|
||||
}
|
||||
```
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 868 KiB |
File diff suppressed because one or more lines are too long
|
After Width: | Height: | Size: 518 KiB |
File diff suppressed because one or more lines are too long
|
After Width: | Height: | Size: 95 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 4.4 MiB |
@@ -0,0 +1,212 @@
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from core.utils.utils import bilinear_sampler
|
||||
|
||||
try:
|
||||
import corr_sampler
|
||||
except:
|
||||
pass
|
||||
|
||||
try:
|
||||
import alt_cuda_corr
|
||||
except:
|
||||
# alt_cuda_corr is not compiled
|
||||
pass
|
||||
|
||||
class CorrSampler(torch.autograd.Function):
|
||||
@staticmethod
|
||||
def forward(ctx, volume, coords, radius):
|
||||
ctx.save_for_backward(volume,coords)
|
||||
ctx.radius = radius
|
||||
corr, = corr_sampler.forward(volume, coords, radius)
|
||||
return corr
|
||||
@staticmethod
|
||||
def backward(ctx, grad_output):
|
||||
volume, coords = ctx.saved_tensors
|
||||
grad_output = grad_output.contiguous()
|
||||
grad_volume, = corr_sampler.backward(volume, coords, grad_output, ctx.radius)
|
||||
return grad_volume, None, None
|
||||
|
||||
|
||||
class CorrBlockFast1D:
|
||||
def __init__(self, fmap1, fmap2, num_levels=4, radius=4, **kwargs):
|
||||
self.num_levels = num_levels
|
||||
self.radius = radius
|
||||
self.corr_pyramid = []
|
||||
# all pairs correlation
|
||||
corr = CorrBlockFast1D.corr(fmap1, fmap2)
|
||||
batch, h1, w1, dim, w2 = corr.shape
|
||||
corr = corr.reshape(batch*h1*w1, dim, 1, w2)
|
||||
for i in range(self.num_levels):
|
||||
self.corr_pyramid.append(corr.view(batch, h1, w1, -1, w2//2**i))
|
||||
corr = F.avg_pool2d(corr, [1, 2], stride=[1, 2])
|
||||
|
||||
def __call__(self, coords):
|
||||
out_pyramid = []
|
||||
bz, _, ht, wd = coords.shape
|
||||
coords = coords[:, [0]]
|
||||
for i in range(self.num_levels):
|
||||
corr = CorrSampler.apply(self.corr_pyramid[i].squeeze(3), coords/2**i, self.radius)
|
||||
out_pyramid.append(corr.view(bz, -1, ht, wd))
|
||||
return torch.cat(out_pyramid, dim=1)
|
||||
|
||||
@staticmethod
|
||||
def corr(fmap1, fmap2):
|
||||
B, D, H, W1 = fmap1.shape
|
||||
_, _, _, W2 = fmap2.shape
|
||||
fmap1 = fmap1.view(B, D, H, W1)
|
||||
fmap2 = fmap2.view(B, D, H, W2)
|
||||
corr = torch.einsum('aijk,aijh->ajkh', fmap1, fmap2)
|
||||
corr = corr.reshape(B, H, W1, 1, W2).contiguous()
|
||||
return corr / torch.sqrt(torch.tensor(D).float())
|
||||
|
||||
|
||||
class PytorchAlternateCorrBlock1D:
|
||||
def __init__(self, fmap1, fmap2, num_levels=4, radius=4, **kwargs):
|
||||
self.num_levels = num_levels
|
||||
self.radius = radius
|
||||
self.corr_pyramid = []
|
||||
self.fmap1 = fmap1
|
||||
self.fmap2 = fmap2
|
||||
|
||||
def corr(self, fmap1, fmap2, coords):
|
||||
B, D, H, W = fmap2.shape
|
||||
# map grid coordinates to [-1,1]
|
||||
xgrid, ygrid = coords.split([1,1], dim=-1)
|
||||
xgrid = 2*xgrid/(W-1) - 1
|
||||
ygrid = 2*ygrid/(H-1) - 1
|
||||
|
||||
grid = torch.cat([xgrid, ygrid], dim=-1)
|
||||
output_corr = []
|
||||
for grid_slice in grid.unbind(3):
|
||||
fmapw_mini = F.grid_sample(fmap2, grid_slice, align_corners=True)
|
||||
corr = torch.sum(fmapw_mini * fmap1, dim=1)
|
||||
output_corr.append(corr)
|
||||
corr = torch.stack(output_corr, dim=1).permute(0,2,3,1)
|
||||
|
||||
return corr / torch.sqrt(torch.tensor(D).float())
|
||||
|
||||
def __call__(self, coords):
|
||||
r = self.radius
|
||||
coords = coords.permute(0, 2, 3, 1)
|
||||
batch, h1, w1, _ = coords.shape
|
||||
fmap1 = self.fmap1
|
||||
fmap2 = self.fmap2
|
||||
out_pyramid = []
|
||||
for i in range(self.num_levels):
|
||||
dx = torch.zeros(1)
|
||||
dy = torch.linspace(-r, r, 2*r+1)
|
||||
delta = torch.stack(torch.meshgrid(dy, dx), axis=-1).to(coords.device)
|
||||
centroid_lvl = coords.reshape(batch, h1, w1, 1, 2).clone()
|
||||
centroid_lvl[..., 0] = centroid_lvl[..., 0] / 2**i
|
||||
coords_lvl = centroid_lvl + delta.view(-1, 2)
|
||||
corr = self.corr(fmap1, fmap2, coords_lvl)
|
||||
fmap2 = F.avg_pool2d(fmap2, [1, 2], stride=[1, 2])
|
||||
out_pyramid.append(corr)
|
||||
out = torch.cat(out_pyramid, dim=-1)
|
||||
return out.permute(0, 3, 1, 2).contiguous().float()
|
||||
|
||||
|
||||
class CorrBlock1D:
|
||||
def __init__(self, fmap1, fmap2, coords, num_levels=4, radius=4,
|
||||
scale_list=[0.25, 0.5, 2.0, 4.0], scale_corr_radius=4):
|
||||
self.num_levels = num_levels
|
||||
self.radius = radius
|
||||
self.scale_list = scale_list
|
||||
self.scale_corr_radius = scale_corr_radius
|
||||
self.corr_pyramid = []
|
||||
self.coords_pyramid = []
|
||||
dx = torch.linspace(-radius, radius, 2*radius+1)
|
||||
self.dx = dx[:, None].to(coords.device)
|
||||
|
||||
sdx = torch.linspace(-scale_corr_radius, scale_corr_radius, 2*scale_corr_radius+1)
|
||||
self.sdx = sdx[:, None].to(coords.device)
|
||||
|
||||
# all pairs correlation
|
||||
corr = CorrBlock1D.corr(fmap1, fmap2)
|
||||
|
||||
batch, h1, w1, _, w2 = corr.shape
|
||||
self.batch = batch
|
||||
self.h1 = h1
|
||||
self.w1 = w1
|
||||
self.w2 = w2
|
||||
corr = corr.reshape(batch*h1*w1, 1, 1, w2)
|
||||
self.coords = coords.reshape(batch*h1*w1, 1, 1, 1)
|
||||
|
||||
self.corr_pyramid.append(corr)
|
||||
for i in range(1, self.num_levels):
|
||||
corr = F.avg_pool2d(corr, [1, 2], stride=[1, 2])
|
||||
self.corr_pyramid.append(corr)
|
||||
|
||||
def __call__(self, disp, scaling=False):
|
||||
batch, _, h1, w1 = disp.shape
|
||||
|
||||
disp = disp.reshape(self.batch*self.h1*self.w1, 1, 1, 1)
|
||||
out_pyramid = []
|
||||
|
||||
if scaling:
|
||||
corr = self.corr_pyramid[0]
|
||||
for scale in self.scale_list:
|
||||
x0 = self.sdx + self.coords - scale * disp
|
||||
y0 = torch.zeros_like(x0)
|
||||
coords_lvl = torch.cat([x0, y0], dim=-1)
|
||||
corr_s = bilinear_sampler(corr, coords_lvl)
|
||||
corr_s = corr_s.view(self.batch, self.h1, self.w1, -1)
|
||||
out_pyramid.append(corr_s)
|
||||
else:
|
||||
coords = self.coords - disp
|
||||
for i in range(self.num_levels):
|
||||
corr = self.corr_pyramid[i]
|
||||
x0 = self.dx + coords / 2**i
|
||||
y0 = torch.zeros_like(x0)
|
||||
coords_lvl = torch.cat([x0, y0], dim=-1)
|
||||
corr_s = bilinear_sampler(corr, coords_lvl)
|
||||
corr_s = corr_s.view(self.batch, self.h1, self.w1, -1)
|
||||
out_pyramid.append(corr_s)
|
||||
|
||||
out = torch.cat(out_pyramid, dim=-1)
|
||||
return out.permute(0, 3, 1, 2).contiguous().float()
|
||||
|
||||
@staticmethod
|
||||
def corr(fmap1, fmap2):
|
||||
B, D, H, W1 = fmap1.shape
|
||||
_, _, _, W2 = fmap2.shape
|
||||
fmap1 = fmap1.view(B, D, H, W1)
|
||||
fmap2 = fmap2.view(B, D, H, W2)
|
||||
corr = torch.einsum('aijk,aijh->ajkh', fmap1, fmap2)
|
||||
corr = corr.reshape(B, H, W1, 1, W2).contiguous()
|
||||
return corr / torch.sqrt(torch.tensor(D).float())
|
||||
|
||||
|
||||
class AlternateCorrBlock:
|
||||
def __init__(self, fmap1, fmap2, num_levels=4, radius=4, **kwargs):
|
||||
raise NotImplementedError
|
||||
self.num_levels = num_levels
|
||||
self.radius = radius
|
||||
|
||||
self.pyramid = [(fmap1, fmap2)]
|
||||
for i in range(1, self.num_levels):
|
||||
fmap1 = F.avg_pool2d(fmap1, 2, stride=2)
|
||||
fmap2 = F.avg_pool2d(fmap2, 2, stride=2)
|
||||
self.pyramid.append((fmap1, fmap2))
|
||||
|
||||
def __call__(self, coords):
|
||||
coords = coords.permute(0, 2, 3, 1)
|
||||
B, H, W, _ = coords.shape
|
||||
dim = self.pyramid[0][0].shape[1]
|
||||
|
||||
corr_list = []
|
||||
for i in range(self.num_levels):
|
||||
r = self.radius
|
||||
fmap1_i = self.pyramid[0][0].permute(0, 2, 3, 1).contiguous()
|
||||
fmap2_i = self.pyramid[i][1].permute(0, 2, 3, 1).contiguous()
|
||||
|
||||
coords_i = (coords / 2**i).reshape(B, 1, H, W, 2).contiguous()
|
||||
corr, = alt_cuda_corr.forward(fmap1_i, fmap2_i, coords_i, r)
|
||||
corr_list.append(corr.squeeze(1))
|
||||
|
||||
corr = torch.stack(corr_list, dim=1)
|
||||
corr = corr.reshape(B, -1, H, W)
|
||||
return corr / torch.sqrt(torch.tensor(dim).float())
|
||||
@@ -0,0 +1,142 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from core.update import BasicMultiUpdateBlock, ScaleBasicMultiUpdateBlock
|
||||
from core.extractor import BasicEncoder, MultiBasicEncoder, ResidualBlock, DefomEncoder
|
||||
from core.corr import CorrBlock1D, PytorchAlternateCorrBlock1D, CorrBlockFast1D, AlternateCorrBlock
|
||||
from core.utils.utils import coords_grid, upflow, get_danv2_io_size
|
||||
|
||||
|
||||
try:
|
||||
autocast = torch.cuda.amp.autocast
|
||||
except:
|
||||
# dummy autocast for PyTorch < 1.6
|
||||
class autocast:
|
||||
def __init__(self, enabled):
|
||||
pass
|
||||
def __enter__(self):
|
||||
pass
|
||||
def __exit__(self, *args):
|
||||
pass
|
||||
|
||||
|
||||
class DEFOMStereo(nn.Module):
|
||||
def __init__(self, args):
|
||||
super(DEFOMStereo, self).__init__()
|
||||
self.args = args
|
||||
|
||||
self.register_buffer('mean', torch.tensor([[0.485, 0.456, 0.406]])[..., None, None] * 255)
|
||||
self.register_buffer('std', torch.tensor([[0.229, 0.224, 0.225]])[..., None, None] * 255)
|
||||
|
||||
self.defomencoder = DefomEncoder(args.dinov2_encoder, idepth_scale=args.idepth_scale)
|
||||
|
||||
context_dims = args.hidden_dims
|
||||
|
||||
self.fnet = BasicEncoder(self.defomencoder.out_dim, output_dim=256, norm_fn='instance', downsample=args.n_downsample)
|
||||
|
||||
self.context_zqr_convs = nn.ModuleList([nn.Conv2d(context_dims[i], args.hidden_dims[i]*3, 3, padding=3//2) for i in range(self.args.n_gru_layers)])
|
||||
|
||||
self.update_block = BasicMultiUpdateBlock(self.args, hidden_dims=args.hidden_dims)
|
||||
self.scale_update_block = ScaleBasicMultiUpdateBlock(self.args, hidden_dims=args.hidden_dims)
|
||||
|
||||
self.cnet = MultiBasicEncoder(self.defomencoder.out_dim, output_dim=[args.hidden_dims, context_dims],
|
||||
norm_fn=args.context_norm, downsample=args.n_downsample)
|
||||
|
||||
def freeze_bn(self):
|
||||
for m in self.modules():
|
||||
if isinstance(m, nn.BatchNorm2d):
|
||||
m.eval()
|
||||
|
||||
def initialize_coords(self, img):
|
||||
""" Disparity is represented as difference between two vertical coordinate grids disp
|
||||
= coords0[:, :1] - coords1[:, :1] """
|
||||
N, _, H, W = img.shape
|
||||
|
||||
coords = coords_grid(N, H, W)[:, :1].to(img.device)
|
||||
|
||||
return coords
|
||||
|
||||
def upsample_flow(self, flow, mask):
|
||||
""" Upsample disparity field [H/scale, W/scale, 1] -> [H, W, 1] using convex combination """
|
||||
N, D, H, W = flow.shape
|
||||
factor = 2 ** self.args.n_downsample
|
||||
mask = mask.view(N, 1, 9, factor, factor, H, W)
|
||||
mask = torch.softmax(mask, dim=2)
|
||||
|
||||
up_flow = F.unfold(factor * flow, [3, 3], padding=1)
|
||||
up_flow = up_flow.view(N, D, 9, 1, 1, H, W)
|
||||
|
||||
up_flow = torch.sum(mask * up_flow, dim=2)
|
||||
up_flow = up_flow.permute(0, 1, 4, 2, 5, 3)
|
||||
return up_flow.reshape(N, D, factor * H, factor * W)
|
||||
|
||||
def forward(self, image1, image2, iters=12, scale_iters=3, test_mode=False):
|
||||
""" Estimate optical flow between pair of frames """
|
||||
|
||||
image1 = ((image1 - self.mean)/self.std).contiguous()
|
||||
image2 = ((image2 - self.mean)/self.std).contiguous()
|
||||
|
||||
bs, _, h, w = image1.shape
|
||||
danv2_io_sizes = get_danv2_io_size(h, w, self.args.n_downsample)
|
||||
|
||||
# run the context network
|
||||
with autocast(enabled=self.args.mixed_precision):
|
||||
d_features, dfeat1, dfeat2, disp = self.defomencoder([image1, image2], danv2_io_sizes)
|
||||
|
||||
cnet_list = self.cnet(image1, d_features)
|
||||
fmap1, fmap2 = self.fnet([image1, image2], [dfeat1, dfeat2])
|
||||
net_list = [torch.tanh(x[0]) for x in cnet_list]
|
||||
inp_list = [torch.relu(x[1]) for x in cnet_list]
|
||||
# Rather than running the GRU's conv layers on the context features multiple times, we do it once at the beginning
|
||||
inp_list = [list(conv(i).split(split_size=conv.out_channels//3, dim=1)) for i, conv in zip(inp_list, self.context_zqr_convs)]
|
||||
|
||||
coords = self.initialize_coords(net_list[0])
|
||||
|
||||
fmap1, fmap2 = fmap1.float(), fmap2.float()
|
||||
disp = disp.float()
|
||||
corr_fn = CorrBlock1D(fmap1, fmap2, coords, radius=self.args.corr_radius, num_levels=self.args.corr_levels,
|
||||
scale_list=self.args.scale_list, scale_corr_radius=self.args.scale_corr_radius)
|
||||
|
||||
disp_predictions = []
|
||||
for itr in range(iters):
|
||||
disp = disp.detach()
|
||||
|
||||
if itr < scale_iters:
|
||||
corr = corr_fn(disp, scaling=True) # index correlation volume
|
||||
with autocast(enabled=self.args.mixed_precision):
|
||||
net_list, up_mask, scale_disp = self.scale_update_block(net_list, inp_list, corr, disp,
|
||||
iter32=self.args.n_gru_layers == 3,
|
||||
iter16=self.args.n_gru_layers >= 2)
|
||||
|
||||
# F(t+1) = \Scale(t) x F(t)
|
||||
disp = scale_disp * disp
|
||||
else:
|
||||
corr = corr_fn(disp, scaling=False) # index correlation volume
|
||||
with autocast(enabled=self.args.mixed_precision):
|
||||
net_list, up_mask, delta_disp = self.update_block(net_list, inp_list, corr, disp,
|
||||
iter32=self.args.n_gru_layers == 3,
|
||||
iter16=self.args.n_gru_layers >= 2)
|
||||
|
||||
# To avoid unstability, we limit the disparity update within the searching range.
|
||||
delta_disp = torch.clip(delta_disp, min=-2**(self.args.corr_levels-1)*self.args.corr_radius,
|
||||
max=2**(self.args.corr_levels-1)*self.args.corr_radius)
|
||||
|
||||
# F(t+1) = F(t) + \Delta(t)
|
||||
disp = disp + delta_disp
|
||||
|
||||
# We do not need to upsample or output intermediate results in test_mode
|
||||
if test_mode and itr < iters - 1:
|
||||
continue
|
||||
|
||||
# upsample predictions
|
||||
if up_mask is None:
|
||||
disp_up = upflow(disp, factor=2 ** self.n_downsample)
|
||||
else:
|
||||
disp_up = self.upsample_flow(disp, up_mask)
|
||||
|
||||
disp_predictions.append(disp_up)
|
||||
|
||||
if test_mode:
|
||||
return disp_up
|
||||
|
||||
return disp_predictions
|
||||
@@ -0,0 +1,388 @@
|
||||
import os
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from timm.models.layers import DropPath
|
||||
|
||||
from depth_anything_v2.dpt import DepthAnythingV2
|
||||
|
||||
|
||||
class ConvBlock(nn.Module):
|
||||
def __init__(self, in_planes, planes, norm_fn='group', stride=1):
|
||||
super(ConvBlock, self).__init__()
|
||||
|
||||
self.conv = nn.Conv2d(in_planes, planes, kernel_size=3, padding=1, stride=stride)
|
||||
self.relu = nn.ReLU(inplace=True)
|
||||
|
||||
num_groups = planes // 8
|
||||
|
||||
if norm_fn == 'group':
|
||||
self.norm1 = nn.GroupNorm(num_groups=num_groups, num_channels=planes)
|
||||
self.norm2 = nn.GroupNorm(num_groups=num_groups, num_channels=planes)
|
||||
if not (stride == 1 and in_planes == planes):
|
||||
self.norm3 = nn.GroupNorm(num_groups=num_groups, num_channels=planes)
|
||||
|
||||
elif norm_fn == 'batch':
|
||||
self.norm1 = nn.BatchNorm2d(planes)
|
||||
self.norm2 = nn.BatchNorm2d(planes)
|
||||
if not (stride == 1 and in_planes == planes):
|
||||
self.norm3 = nn.BatchNorm2d(planes)
|
||||
|
||||
elif norm_fn == 'instance':
|
||||
self.norm1 = nn.InstanceNorm2d(planes)
|
||||
self.norm2 = nn.InstanceNorm2d(planes)
|
||||
if not (stride == 1 and in_planes == planes):
|
||||
self.norm3 = nn.InstanceNorm2d(planes)
|
||||
|
||||
elif norm_fn == 'none':
|
||||
self.norm1 = nn.Sequential()
|
||||
self.norm2 = nn.Sequential()
|
||||
if not (stride == 1 and in_planes == planes):
|
||||
self.norm3 = nn.Sequential()
|
||||
|
||||
def forward(self, x):
|
||||
|
||||
return self.relu(self.norm1(self.conv(x)))
|
||||
|
||||
|
||||
class ResidualBlock(nn.Module):
|
||||
def __init__(self, in_planes, planes, norm_fn='group', stride=1):
|
||||
super(ResidualBlock, self).__init__()
|
||||
|
||||
self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=3, padding=1, stride=stride)
|
||||
self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, padding=1)
|
||||
self.relu = nn.ReLU(inplace=True)
|
||||
|
||||
num_groups = planes // 8
|
||||
|
||||
if norm_fn == 'group':
|
||||
self.norm1 = nn.GroupNorm(num_groups=num_groups, num_channels=planes)
|
||||
self.norm2 = nn.GroupNorm(num_groups=num_groups, num_channels=planes)
|
||||
if not (stride == 1 and in_planes == planes):
|
||||
self.norm3 = nn.GroupNorm(num_groups=num_groups, num_channels=planes)
|
||||
|
||||
elif norm_fn == 'batch':
|
||||
self.norm1 = nn.BatchNorm2d(planes)
|
||||
self.norm2 = nn.BatchNorm2d(planes)
|
||||
if not (stride == 1 and in_planes == planes):
|
||||
self.norm3 = nn.BatchNorm2d(planes)
|
||||
|
||||
elif norm_fn == 'instance':
|
||||
self.norm1 = nn.InstanceNorm2d(planes)
|
||||
self.norm2 = nn.InstanceNorm2d(planes)
|
||||
if not (stride == 1 and in_planes == planes):
|
||||
self.norm3 = nn.InstanceNorm2d(planes)
|
||||
|
||||
elif norm_fn == 'none':
|
||||
self.norm1 = nn.Sequential()
|
||||
self.norm2 = nn.Sequential()
|
||||
if not (stride == 1 and in_planes == planes):
|
||||
self.norm3 = nn.Sequential()
|
||||
|
||||
if stride == 1 and in_planes == planes:
|
||||
self.downsample = None
|
||||
|
||||
else:
|
||||
self.downsample = nn.Sequential(
|
||||
nn.Conv2d(in_planes, planes, kernel_size=1, stride=stride), self.norm3)
|
||||
|
||||
def forward(self, x):
|
||||
y = x
|
||||
y = self.conv1(y)
|
||||
y = self.norm1(y)
|
||||
y = self.relu(y)
|
||||
y = self.conv2(y)
|
||||
y = self.norm2(y)
|
||||
y = self.relu(y)
|
||||
|
||||
if self.downsample is not None:
|
||||
x = self.downsample(x)
|
||||
|
||||
return self.relu(x+y)
|
||||
|
||||
|
||||
class BottleneckBlock(nn.Module):
|
||||
def __init__(self, in_planes, planes, norm_fn='group', stride=1, ratio=4):
|
||||
super(BottleneckBlock, self).__init__()
|
||||
|
||||
self.conv1 = nn.Conv2d(in_planes, planes // ratio, kernel_size=1, padding=0)
|
||||
self.conv2 = nn.Conv2d(planes // ratio, planes // ratio, kernel_size=3, padding=1, stride=stride)
|
||||
self.conv3 = nn.Conv2d(planes // ratio, planes, kernel_size=1, padding=0)
|
||||
self.relu = nn.ReLU(inplace=True)
|
||||
|
||||
num_groups = planes // 8
|
||||
|
||||
if norm_fn == 'group':
|
||||
self.norm1 = nn.GroupNorm(num_groups=num_groups, num_channels=planes // ratio)
|
||||
self.norm2 = nn.GroupNorm(num_groups=num_groups, num_channels=planes // ratio)
|
||||
self.norm3 = nn.GroupNorm(num_groups=num_groups, num_channels=planes)
|
||||
if not (stride == 1 and in_planes == planes):
|
||||
self.norm4 = nn.GroupNorm(num_groups=num_groups, num_channels=planes)
|
||||
|
||||
elif norm_fn == 'batch':
|
||||
self.norm1 = nn.BatchNorm2d(planes // ratio)
|
||||
self.norm2 = nn.BatchNorm2d(planes // ratio)
|
||||
self.norm3 = nn.BatchNorm2d(planes)
|
||||
if not (stride == 1 and in_planes == planes):
|
||||
self.norm4 = nn.BatchNorm2d(planes)
|
||||
|
||||
elif norm_fn == 'instance':
|
||||
self.norm1 = nn.InstanceNorm2d(planes // ratio)
|
||||
self.norm2 = nn.InstanceNorm2d(planes // ratio)
|
||||
self.norm3 = nn.InstanceNorm2d(planes)
|
||||
if not (stride == 1 and in_planes == planes):
|
||||
self.norm4 = nn.InstanceNorm2d(planes)
|
||||
|
||||
elif norm_fn == 'none':
|
||||
self.norm1 = nn.Sequential()
|
||||
self.norm2 = nn.Sequential()
|
||||
self.norm3 = nn.Sequential()
|
||||
if not (stride == 1 and in_planes == planes):
|
||||
self.norm4 = nn.Sequential()
|
||||
|
||||
if stride == 1 and in_planes == planes:
|
||||
self.downsample = None
|
||||
|
||||
else:
|
||||
self.downsample = nn.Sequential(
|
||||
nn.Conv2d(in_planes, planes, kernel_size=1, stride=stride), self.norm4)
|
||||
|
||||
def forward(self, x):
|
||||
y = x
|
||||
y = self.relu(self.norm1(self.conv1(y)))
|
||||
y = self.relu(self.norm2(self.conv2(y)))
|
||||
y = self.relu(self.norm3(self.conv3(y)))
|
||||
|
||||
if self.downsample is not None:
|
||||
x = self.downsample(x)
|
||||
|
||||
return self.relu(x + y)
|
||||
|
||||
|
||||
class BasicEncoder(nn.Module):
|
||||
def __init__(self, d_dim, output_dim=128, norm_fn='batch', downsample=3):
|
||||
super(BasicEncoder, self).__init__()
|
||||
self.norm_fn = norm_fn
|
||||
self.downsample = downsample
|
||||
|
||||
if self.norm_fn == 'group':
|
||||
self.norm1 = nn.GroupNorm(num_groups=8, num_channels=64)
|
||||
|
||||
elif self.norm_fn == 'batch':
|
||||
self.norm1 = nn.BatchNorm2d(64)
|
||||
|
||||
elif self.norm_fn == 'instance':
|
||||
self.norm1 = nn.InstanceNorm2d(64)
|
||||
|
||||
elif self.norm_fn == 'none':
|
||||
self.norm1 = nn.Sequential()
|
||||
|
||||
self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=1 + (downsample > 2), padding=3)
|
||||
self.relu1 = nn.ReLU(inplace=True)
|
||||
|
||||
self.in_planes = 64
|
||||
self.layer1 = self._make_layer(64, stride=1)
|
||||
self.layer2 = self._make_layer(96, stride=1 + (downsample > 1))
|
||||
self.layer3 = self._make_layer(128, stride=1 + (downsample > 0))
|
||||
|
||||
# depth feat convolution
|
||||
self.convd = ConvBlock(d_dim, 128, self.norm_fn)
|
||||
|
||||
# output convolution
|
||||
self.conv2 = nn.Conv2d(128, output_dim, kernel_size=1)
|
||||
|
||||
for m in self.modules():
|
||||
if isinstance(m, nn.Conv2d):
|
||||
nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
|
||||
elif isinstance(m, (nn.BatchNorm2d, nn.InstanceNorm2d, nn.GroupNorm)):
|
||||
if m.weight is not None:
|
||||
nn.init.constant_(m.weight, 1)
|
||||
if m.bias is not None:
|
||||
nn.init.constant_(m.bias, 0)
|
||||
|
||||
def _make_layer(self, dim, stride=1):
|
||||
layer1 = ResidualBlock(self.in_planes, dim, self.norm_fn, stride=stride)
|
||||
layer2 = ResidualBlock(dim, dim, self.norm_fn, stride=1)
|
||||
layers = (layer1, layer2)
|
||||
|
||||
self.in_planes = dim
|
||||
return nn.Sequential(*layers)
|
||||
|
||||
def forward(self, x, dfeats):
|
||||
|
||||
# if input is list, combine batch dimension
|
||||
is_list = isinstance(x, tuple) or isinstance(x, list)
|
||||
if is_list:
|
||||
batch_dim = x[0].shape[0]
|
||||
x = torch.cat(x, dim=0)
|
||||
|
||||
is_list = isinstance(dfeats, tuple) or isinstance(dfeats, list)
|
||||
if is_list:
|
||||
batch_dim = dfeats[0].shape[0]
|
||||
dfeats = torch.cat(dfeats, dim=0)
|
||||
|
||||
x = self.conv1(x)
|
||||
x = self.norm1(x)
|
||||
x = self.relu1(x)
|
||||
|
||||
x = self.layer1(x)
|
||||
x = self.layer2(x)
|
||||
x = self.layer3(x)
|
||||
|
||||
x = x + self.convd(dfeats)
|
||||
|
||||
x = self.conv2(x)
|
||||
|
||||
if is_list:
|
||||
x = x.split(split_size=batch_dim, dim=0)
|
||||
|
||||
return x
|
||||
|
||||
|
||||
class MultiBasicEncoder(nn.Module):
|
||||
def __init__(self, d_dim, output_dim=[128, 128, 128], norm_fn='batch', downsample=3, drop_path_rate=0.2):
|
||||
super(MultiBasicEncoder, self).__init__()
|
||||
self.d_dim = d_dim
|
||||
self.norm_fn = norm_fn
|
||||
self.downsample = downsample
|
||||
|
||||
if self.norm_fn == 'group':
|
||||
self.norm1 = nn.GroupNorm(num_groups=8, num_channels=64)
|
||||
|
||||
elif self.norm_fn == 'batch':
|
||||
self.norm1 = nn.BatchNorm2d(64)
|
||||
|
||||
elif self.norm_fn == 'instance':
|
||||
self.norm1 = nn.InstanceNorm2d(64)
|
||||
|
||||
elif self.norm_fn == 'none':
|
||||
self.norm1 = nn.Sequential()
|
||||
|
||||
self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=1 + (downsample > 2), padding=3)
|
||||
self.relu1 = nn.ReLU(inplace=True)
|
||||
|
||||
self.in_planes = 64
|
||||
self.layer1 = self._make_layer(64, stride=1)
|
||||
self.layer2 = self._make_layer(96, stride=1 + (downsample > 1))
|
||||
self.layer3 = self._make_layer(128, stride=1 + (downsample > 0))
|
||||
self.layer4 = self._make_layer(128, stride=2)
|
||||
self.layer5 = self._make_layer(128, stride=2)
|
||||
|
||||
self.drop_path = DropPath(drop_path_rate)
|
||||
|
||||
self.conv08 = ConvBlock(d_dim, 128, self.norm_fn)
|
||||
output_list = []
|
||||
for dim in output_dim:
|
||||
conv_out = nn.Sequential(
|
||||
ResidualBlock(128, 128, self.norm_fn, stride=1),
|
||||
nn.Conv2d(128, dim[2], 3, padding=1))
|
||||
output_list.append(conv_out)
|
||||
|
||||
self.outputs08 = nn.ModuleList(output_list)
|
||||
|
||||
self.conv16 = ConvBlock(d_dim, 128, self.norm_fn)
|
||||
output_list = []
|
||||
for dim in output_dim:
|
||||
conv_out = nn.Sequential(
|
||||
ResidualBlock(128, 128, self.norm_fn, stride=1),
|
||||
nn.Conv2d(128, dim[1], 3, padding=1))
|
||||
output_list.append(conv_out)
|
||||
|
||||
self.outputs16 = nn.ModuleList(output_list)
|
||||
|
||||
self.conv32 = ConvBlock(d_dim, 128, self.norm_fn)
|
||||
output_list = []
|
||||
for dim in output_dim:
|
||||
conv_out = nn.Conv2d(128, dim[0], 3, padding=1)
|
||||
output_list.append(conv_out)
|
||||
|
||||
self.outputs32 = nn.ModuleList(output_list)
|
||||
|
||||
for m in self.modules():
|
||||
if isinstance(m, nn.Conv2d):
|
||||
nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
|
||||
elif isinstance(m, (nn.BatchNorm2d, nn.InstanceNorm2d, nn.GroupNorm)):
|
||||
if m.weight is not None:
|
||||
nn.init.constant_(m.weight, 1)
|
||||
if m.bias is not None:
|
||||
nn.init.constant_(m.bias, 0)
|
||||
|
||||
def _make_layer(self, dim, stride=1):
|
||||
layer1 = ResidualBlock(self.in_planes, dim, self.norm_fn, stride=stride)
|
||||
layer2 = ResidualBlock(dim, dim, self.norm_fn, stride=1)
|
||||
layers = (layer1, layer2)
|
||||
|
||||
self.in_planes = dim
|
||||
return nn.Sequential(*layers)
|
||||
|
||||
def forward(self, x, d_feats, num_layers=3):
|
||||
|
||||
x = self.conv1(x)
|
||||
x = self.norm1(x)
|
||||
x = self.relu1(x)
|
||||
|
||||
x = self.layer1(x)
|
||||
x = self.layer2(x)
|
||||
x = self.layer3(x)
|
||||
|
||||
feat = x + self.drop_path(self.conv08(d_feats[0]))
|
||||
outputs08 = [f(feat) for f in self.outputs08]
|
||||
if num_layers == 1:
|
||||
return (outputs08,)
|
||||
|
||||
y = self.layer4(x)
|
||||
feat = y + self.drop_path(self.conv16(d_feats[1]))
|
||||
outputs16 = [f(feat) for f in self.outputs16]
|
||||
|
||||
if num_layers == 2:
|
||||
return (outputs08, outputs16)
|
||||
|
||||
z = self.layer5(y)
|
||||
feat = z + self.drop_path(self.conv32(d_feats[2]))
|
||||
outputs32 = [f(feat) for f in self.outputs32]
|
||||
|
||||
return (outputs08, outputs16, outputs32)
|
||||
|
||||
|
||||
class DefomEncoder(nn.Module):
|
||||
def __init__(self, dinov2_encoder, pretrained=True, freeze=True, idepth_scale=0.25):
|
||||
super(DefomEncoder, self).__init__()
|
||||
self.dinov2_encoder = dinov2_encoder
|
||||
self.idepth_scale = idepth_scale
|
||||
self.pretrained = pretrained
|
||||
self.freeze = freeze
|
||||
|
||||
model_configs = {
|
||||
'vits': {'encoder': 'vits', 'features': 64, 'out_channels': [48, 96, 192, 384]},
|
||||
'vitb': {'encoder': 'vitb', 'features': 128, 'out_channels': [96, 192, 384, 768]},
|
||||
'vitl': {'encoder': 'vitl', 'features': 256, 'out_channels': [256, 512, 1024, 1024]},
|
||||
'vitg': {'encoder': 'vitg', 'features': 384, 'out_channels': [1536, 1536, 1536, 1536]}
|
||||
}
|
||||
|
||||
self.depth_anything = DepthAnythingV2(**model_configs[self.dinov2_encoder])
|
||||
|
||||
if pretrained and os.path.exists(f'./checkpoints/depth_anything_v2_{dinov2_encoder}.pth'):
|
||||
self.depth_anything.load_state_dict(
|
||||
torch.load(f'./checkpoints/depth_anything_v2_{dinov2_encoder}.pth', map_location='cpu'), strict=False)
|
||||
if freeze:
|
||||
for param in self.depth_anything.pretrained.parameters():
|
||||
param.requires_grad = False
|
||||
for param in self.depth_anything.depth_head.parameters():
|
||||
param.requires_grad = False
|
||||
|
||||
self.out_dim = model_configs[self.dinov2_encoder]['features']
|
||||
|
||||
def forward(self, x, danv2_io_sizes):
|
||||
|
||||
x = torch.cat(x, dim=0)
|
||||
ih, iw, oh, ow = danv2_io_sizes
|
||||
x = F.interpolate(x, (ih, iw), mode="bilinear", align_corners=True)
|
||||
|
||||
features, left_feat, right_feat, idepth = self.depth_anything(x, oh, ow)
|
||||
|
||||
bs = idepth.shape[0]
|
||||
max_idepth, _ = torch.max(idepth.view(bs, -1), dim=1)
|
||||
max_idepth = max_idepth.detach().view(bs, 1, 1, 1) + 1e-8
|
||||
idepth = idepth / max_idepth * self.idepth_scale * ow + 0.01
|
||||
|
||||
return features, left_feat, right_feat, idepth
|
||||
@@ -0,0 +1,583 @@
|
||||
# Data loading based on https://github.com/NVIDIA/flownet2-pytorch
|
||||
|
||||
import numpy as np
|
||||
from numpy import linalg as LA
|
||||
import torch
|
||||
import torch.utils.data as data
|
||||
import torch.nn.functional as F
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
import copy
|
||||
import math
|
||||
import random
|
||||
from pathlib import Path
|
||||
from glob import glob
|
||||
import os.path as osp
|
||||
|
||||
from core.utils import frame_utils
|
||||
from core.utils.augmentor import DispAugmentor, SparseDispAugmentor
|
||||
|
||||
|
||||
class StereoDataset(data.Dataset):
|
||||
def __init__(self, aug_params=None, sparse=False, reader=None, is_eval=False, is_test=False):
|
||||
self.augmentor = None
|
||||
self.sparse = sparse
|
||||
if aug_params is not None and "crop_size" in aug_params:
|
||||
if sparse:
|
||||
self.augmentor = SparseDispAugmentor(**aug_params)
|
||||
else:
|
||||
self.augmentor = DispAugmentor(**aug_params)
|
||||
|
||||
if reader is None:
|
||||
self.disparity_reader = frame_utils.read_gen
|
||||
else:
|
||||
self.disparity_reader = reader
|
||||
|
||||
self.is_eval = is_eval
|
||||
self.is_test = is_test
|
||||
self.init_seed = False
|
||||
self.disparity_list = []
|
||||
self.image_list = []
|
||||
|
||||
# number of copies of the datasets
|
||||
self.v = 1
|
||||
|
||||
def __getitem__(self, index):
|
||||
|
||||
if self.is_test:
|
||||
img1 = frame_utils.read_gen(self.image_list[index][0])
|
||||
img2 = frame_utils.read_gen(self.image_list[index][1])
|
||||
img1 = np.array(img1).astype(np.uint8)
|
||||
img2 = np.array(img2).astype(np.uint8)
|
||||
if len(img1.shape) == 2:
|
||||
img1 = np.tile(img1[..., None], (1, 1, 3))
|
||||
img2 = np.tile(img2[..., None], (1, 1, 3))
|
||||
else:
|
||||
img1 = img1[..., :3]
|
||||
img2 = img2[..., :3]
|
||||
img1 = torch.from_numpy(img1).permute(2, 0, 1).float()
|
||||
img2 = torch.from_numpy(img2).permute(2, 0, 1).float()
|
||||
return img1, img2, self.image_list[index][0]
|
||||
|
||||
if not self.init_seed:
|
||||
worker_info = torch.utils.data.get_worker_info()
|
||||
if worker_info is not None:
|
||||
torch.manual_seed(worker_info.id)
|
||||
np.random.seed(worker_info.id)
|
||||
random.seed(worker_info.id)
|
||||
self.init_seed = True
|
||||
|
||||
index = index % (len(self.image_list)*self.v)
|
||||
index = index % len(self.image_list)
|
||||
|
||||
if not self.is_eval and len(self.disparity_list[index]) > 1 and np.random.rand() > 0.5:
|
||||
disp = self.disparity_reader(self.disparity_list[index][1])
|
||||
if isinstance(disp, tuple):
|
||||
disp, valid = disp
|
||||
else:
|
||||
valid = disp < 1024
|
||||
img1 = frame_utils.read_gen(self.image_list[index][1])
|
||||
img2 = frame_utils.read_gen(self.image_list[index][0])
|
||||
|
||||
img1 = np.array(img1).astype(np.uint8)[:, ::-1]
|
||||
img2 = np.array(img2).astype(np.uint8)[:, ::-1]
|
||||
disp = np.array(disp).astype(np.float32)[:, ::-1]
|
||||
valid = np.array(valid).astype(np.bool_)[:, ::-1]
|
||||
|
||||
else:
|
||||
disp = self.disparity_reader(self.disparity_list[index][0])
|
||||
if isinstance(disp, tuple):
|
||||
disp, valid = disp
|
||||
else:
|
||||
valid = disp < 1024
|
||||
|
||||
img1 = frame_utils.read_gen(self.image_list[index][0])
|
||||
img2 = frame_utils.read_gen(self.image_list[index][1])
|
||||
|
||||
img1 = np.array(img1).astype(np.uint8)
|
||||
img2 = np.array(img2).astype(np.uint8)
|
||||
disp = np.array(disp).astype(np.float32)
|
||||
valid = np.array(valid).astype(np.bool_)
|
||||
|
||||
# grayscale images
|
||||
if len(img1.shape) == 2:
|
||||
img1 = np.tile(img1[..., None], (1, 1, 3))
|
||||
img2 = np.tile(img2[..., None], (1, 1, 3))
|
||||
else:
|
||||
img1 = img1[..., :3]
|
||||
img2 = img2[..., :3]
|
||||
|
||||
if self.augmentor is not None:
|
||||
if self.sparse:
|
||||
img1, img2, disp, valid = self.augmentor(img1, img2, disp, valid)
|
||||
else:
|
||||
img1, img2, disp = self.augmentor(img1, img2, disp)
|
||||
|
||||
img1 = torch.from_numpy(img1.copy()).permute(2, 0, 1).float()
|
||||
img2 = torch.from_numpy(img2.copy()).permute(2, 0, 1).float()
|
||||
disp = torch.from_numpy(disp[..., np.newaxis].copy()).permute(2, 0, 1).float()
|
||||
if self.sparse:
|
||||
valid = torch.from_numpy(valid[..., np.newaxis].astype(np.bool_).copy()).permute(2, 0, 1)
|
||||
else:
|
||||
valid = disp < 512
|
||||
|
||||
return {"img1": img1, "img2": img2, "disp": disp, "valid": valid, "imageL_file": self.image_list[index][0], "disp_file": self.disparity_list[index][0]}
|
||||
|
||||
def __mul__(self, v):
|
||||
self.v = v
|
||||
return self
|
||||
|
||||
def __len__(self):
|
||||
return len(self.image_list)*self.v
|
||||
|
||||
|
||||
class SceneFlowDatasets(StereoDataset):
|
||||
def __init__(self, aug_params=None, root='./datasets/SceneFlow/', dstype='frames_cleanpass', things_test=False):
|
||||
super(SceneFlowDatasets, self).__init__(aug_params, is_eval=things_test)
|
||||
self.root = root
|
||||
self.dstype = dstype
|
||||
|
||||
if things_test:
|
||||
self._add_things("TEST")
|
||||
else:
|
||||
self._add_things("TRAIN")
|
||||
self._add_monkaa()
|
||||
self._add_driving()
|
||||
|
||||
def _add_things(self, split='TRAIN'):
|
||||
""" Add FlyingThings3D data """
|
||||
|
||||
original_length = len(self.disparity_list)
|
||||
root = osp.join(self.root, 'FlyingThings3D')
|
||||
left_images = sorted(glob(osp.join(root, self.dstype, split, '*/*/left/*.png')))
|
||||
right_images = [im.replace('left', 'right') for im in left_images]
|
||||
disparity_images = [im.replace(self.dstype, 'disparity').replace('.png', '.pfm') for im in left_images]
|
||||
|
||||
# Choose a random subset of 400 images for validation
|
||||
state = np.random.get_state()
|
||||
np.random.seed(1000)
|
||||
val_idxs = set(np.random.permutation(len(left_images))[:400])
|
||||
np.random.set_state(state)
|
||||
|
||||
for idx, (img1, img2, disp) in enumerate(zip(left_images, right_images, disparity_images)):
|
||||
if (split == 'TEST' and idx in val_idxs) or split == 'TRAIN':
|
||||
self.image_list += [[img1, img2]]
|
||||
self.disparity_list += [[disp, disp.replace('left', 'right')]]
|
||||
logging.info(f"Added {len(self.disparity_list) - original_length} from FlyingThings {self.dstype}")
|
||||
|
||||
def _add_monkaa(self):
|
||||
""" Add FlyingThings3D data """
|
||||
|
||||
original_length = len(self.disparity_list)
|
||||
root = osp.join(self.root, 'Monkaa')
|
||||
left_images = sorted(glob(osp.join(root, self.dstype, '*/left/*.png')) )
|
||||
right_images = [image_file.replace('left', 'right') for image_file in left_images ]
|
||||
disparity_images = [im.replace(self.dstype, 'disparity').replace('.png', '.pfm') for im in left_images ]
|
||||
|
||||
for img1, img2, disp in zip(left_images, right_images, disparity_images):
|
||||
self.image_list += [[img1, img2]]
|
||||
self.disparity_list += [[disp, disp.replace('left', 'right')]]
|
||||
logging.info(f"Added {len(self.disparity_list) - original_length} from Monkaa {self.dstype}")
|
||||
|
||||
def _add_driving(self):
|
||||
""" Add FlyingThings3D data """
|
||||
|
||||
original_length = len(self.disparity_list)
|
||||
root = osp.join(self.root, 'Driving')
|
||||
left_images = sorted(glob(osp.join(root, self.dstype, '*/*/*/left/*.png')) )
|
||||
right_images = [image_file.replace('left', 'right') for image_file in left_images ]
|
||||
disparity_images = [im.replace(self.dstype, 'disparity').replace('.png', '.pfm') for im in left_images ]
|
||||
|
||||
for img1, img2, disp in zip(left_images, right_images, disparity_images):
|
||||
self.image_list += [[img1, img2]]
|
||||
self.disparity_list += [[disp, disp.replace('left', 'right')]]
|
||||
logging.info(f"Added {len(self.disparity_list) - original_length} from Driving {self.dstype}")
|
||||
|
||||
|
||||
class ETH3D(StereoDataset):
|
||||
def __init__(self, aug_params=None, root='./datasets/ETH3D', split='training', is_eval=False, is_test=False):
|
||||
super(ETH3D, self).__init__(aug_params, sparse=True, is_eval=is_eval, is_test=is_test)
|
||||
|
||||
image1_list = sorted(glob(osp.join(root, f'two_view_{split}/*/im0.png')))
|
||||
image2_list = sorted(glob(osp.join(root, f'two_view_{split}/*/im1.png')))
|
||||
disp_list = sorted(glob(osp.join(root, 'two_view_training_gt/*/disp0GT.pfm'))) if split == 'training'\
|
||||
else [osp.join(root, 'two_view_training_gt/playground_1l/disp0GT.pfm')]*len(image1_list)
|
||||
|
||||
for img1, img2, disp in zip(image1_list, image2_list, disp_list):
|
||||
self.image_list += [[img1, img2]]
|
||||
self.disparity_list += [[disp]]
|
||||
|
||||
|
||||
class KITTI(StereoDataset):
|
||||
def __init__(self, aug_params=None, root='./datasets/KITTI', split='15', image_set='training', is_eval=False, is_test=False):
|
||||
super(KITTI, self).__init__(aug_params, sparse=True, reader=frame_utils.readDispKITTI, is_eval=is_eval, is_test=is_test)
|
||||
assert split in ["12", "15"]
|
||||
root = root + split
|
||||
assert os.path.exists(root)
|
||||
|
||||
if split == '15':
|
||||
image1_list = sorted(glob(os.path.join(root, image_set, 'image_2/*_10.png')))
|
||||
image2_list = sorted(glob(os.path.join(root, image_set, 'image_3/*_10.png')))
|
||||
disp_list = sorted(
|
||||
glob(os.path.join(root, 'training', 'disp_occ_0/*_10.png'))) if image_set == 'training' else [osp.join(
|
||||
root, 'training/disp_occ_0/000085_10.png')]*len(image1_list)
|
||||
else:
|
||||
image1_list = sorted(glob(os.path.join(root, image_set, 'colored_0/*_10.png')))
|
||||
image2_list = sorted(glob(os.path.join(root, image_set, 'colored_1/*_10.png')))
|
||||
disp_list = sorted(
|
||||
glob(os.path.join(root, 'training', 'disp_occ/*_10.png'))) if image_set == 'training' else [osp.join(
|
||||
root, 'training/disp_occ/000085_10.png')] * len(image1_list)
|
||||
|
||||
for idx, (img1, img2, disp) in enumerate(zip(image1_list, image2_list, disp_list)):
|
||||
self.image_list += [[img1, img2]]
|
||||
self.disparity_list += [[disp]]
|
||||
|
||||
|
||||
class Middlebury(StereoDataset):
|
||||
def __init__(self, aug_params=None, root='./datasets/Middlebury', split='F', image_set='training', is_eval=False, is_test=False):
|
||||
super(Middlebury, self).__init__(aug_params, sparse=True, reader=frame_utils.readDispMiddlebury, is_eval=is_eval, is_test=is_test)
|
||||
assert os.path.exists(root)
|
||||
assert split in ["F", "H", "Q", "2005", "2006", "2014", "2021"]
|
||||
assert image_set in ["training", "test"]
|
||||
|
||||
if split == "2005":
|
||||
scenes = list((Path(root) / "2005").glob("*"))
|
||||
for scene in scenes:
|
||||
self.image_list += [[str(scene / "view1.png"), str(scene / "view5.png")]]
|
||||
self.disparity_list += [[str(scene / "disp1.png"), str(scene / "disp5.png")]]
|
||||
for illum in ["1", "2", "3"]:
|
||||
for exp in ["0", "1", "2"]:
|
||||
self.image_list += [[str(scene / f"Illum{illum}/Exp{exp}/view1.png"), str(scene / f"Illum{illum}/Exp{exp}/view5.png")]]
|
||||
self.disparity_list += [[str(scene / "disp1.png"), str(scene / "disp5.png")]]
|
||||
elif split == "2006":
|
||||
scenes = list((Path(root) / "2006").glob("*"))
|
||||
for scene in scenes:
|
||||
self.image_list += [[str(scene / "view1.png"), str(scene / "view5.png")]]
|
||||
self.disparity_list += [[str(scene / "disp1.png"), str(scene / "disp5.png")]]
|
||||
for illum in ["1", "2", "3"]:
|
||||
for exp in ["0", "1", "2"]:
|
||||
self.image_list += [[str(scene / f"Illum{illum}/Exp{exp}/view1.png"), str(scene / f"Illum{illum}/Exp{exp}/view5.png")]]
|
||||
self.disparity_list += [[str(scene / "disp1.png"), str(scene / "disp5.png")]]
|
||||
elif split == "2014": # datasets/Middlebury/2014/Pipes-perfect/im0.png
|
||||
scenes = list((Path(root) / "2014").glob("*"))
|
||||
for scene in scenes:
|
||||
for s in ["E", "L", ""]:
|
||||
self.image_list += [[str(scene / "im0.png"), str(scene / f"im1{s}.png")]]
|
||||
self.disparity_list += [[str(scene / "disp0.pfm"), str(scene / "disp1.pfm")]]
|
||||
elif split == "2021":
|
||||
scenes = list((Path(root) / "2021/data").glob("*"))
|
||||
for scene in scenes:
|
||||
self.image_list += [[str(scene / "im0.png"), str(scene / "im1.png")]]
|
||||
self.disparity_list += [[str(scene / "disp0.pfm"), str(scene / "disp1.pfm")]]
|
||||
for s in ["0", "1", "2", "3"]:
|
||||
if os.path.exists(str(scene / f"ambient/L0/im0e{s}.png")):
|
||||
self.image_list += [[str(scene / f"ambient/L0/im0e{s}.png"), str(scene / f"ambient/L0/im1e{s}.png")]]
|
||||
self.disparity_list += [[str(scene / "disp0.pfm"), str(scene / "disp1.pfm")]]
|
||||
else:
|
||||
if image_set == 'training':
|
||||
lines = list(map(osp.basename, glob(os.path.join(root, "MiddEval3/trainingF/*"))))
|
||||
if is_eval:
|
||||
lines = list(filter(lambda p: any(s in p.split('/') for s in Path(os.path.join(root, "MiddEval3/official_train.txt")).read_text().splitlines()), lines))
|
||||
else:
|
||||
lines = list(map(osp.basename, glob(os.path.join(root, "MiddEval3/testF/*"))))
|
||||
|
||||
image1_list = sorted([os.path.join(root, "MiddEval3", f'{image_set}{split}', f'{name}/im0.png') for name in lines])
|
||||
image2_list = sorted([os.path.join(root, "MiddEval3", f'{image_set}{split}', f'{name}/im1.png') for name in lines])
|
||||
|
||||
disp_list = sorted([os.path.join(root, "MiddEval3", f'training{split}', f'{name}/disp0GT.pfm') for name in lines]) \
|
||||
if image_set == 'training' else [os.path.join(root, "MiddEval3", f'training{split}', 'Adirondack/disp0GT.pfm')]*len(image1_list)
|
||||
|
||||
assert len(image1_list) == len(image2_list) == len(disp_list) > 0, [image1_list, split]
|
||||
for img1, img2, disp in zip(image1_list, image2_list, disp_list):
|
||||
self.image_list += [[img1, img2]]
|
||||
self.disparity_list += [[disp]]
|
||||
|
||||
|
||||
class SintelStereo(StereoDataset):
|
||||
def __init__(self, aug_params=None, root='./datasets/SintelStereo'):
|
||||
super().__init__(aug_params, reader=frame_utils.readDispSintelStereo)
|
||||
|
||||
image1_list = sorted(glob(osp.join(root, 'training/*_left/*/frame_*.png')))
|
||||
image2_list = sorted(glob(osp.join(root, 'training/*_right/*/frame_*.png')))
|
||||
disp_list = sorted(glob(osp.join(root, 'training/disparities/*/frame_*.png'))) * 2
|
||||
|
||||
for img1, img2, disp in zip(image1_list, image2_list, disp_list):
|
||||
assert img1.split('/')[-2:] == disp.split('/')[-2:]
|
||||
self.image_list += [[img1, img2]]
|
||||
self.disparity_list += [[disp]]
|
||||
|
||||
|
||||
class FallingThings(StereoDataset):
|
||||
def __init__(self, aug_params=None, root='./datasets/FallingThings'):
|
||||
super().__init__(aug_params, reader=frame_utils.readDispFallingThings)
|
||||
assert os.path.exists(root)
|
||||
|
||||
image1_list = sorted(glob(osp.join(root, 'fat/single/*/*/*.left.jpg'))) + \
|
||||
sorted(glob(osp.join(root, 'fat/mixed/*/*.left.jpg')))
|
||||
image2_list = [e.replace('left.jpg', 'right.jpg') for e in image1_list]
|
||||
disp_list = [e.replace('left.jpg', 'left.depth.png') for e in image1_list]
|
||||
|
||||
for img1, img2, disp in zip(image1_list, image2_list, disp_list):
|
||||
self.image_list += [[img1, img2]]
|
||||
self.disparity_list += [[disp, disp.replace('left', 'right')]]
|
||||
|
||||
|
||||
class TartanAir(StereoDataset):
|
||||
def __init__(self, aug_params=None, root='./datasets/TartanAir'):
|
||||
super().__init__(aug_params, reader=frame_utils.readDispTartanAir)
|
||||
assert os.path.exists(root)
|
||||
|
||||
image1_list = sorted(glob(osp.join(root, '*/*/*/*/image_left/*_left.png')))
|
||||
image2_list = [e.replace('_left', '_right') for e in image1_list]
|
||||
disp_list = [e.replace('image_left', 'depth_left').replace('left.png', 'left_depth.npy') for e in image1_list]
|
||||
|
||||
for img1, img2, disp in zip(image1_list, image2_list, disp_list):
|
||||
self.image_list += [[img1, img2]]
|
||||
self.disparity_list += [[disp, disp.replace('left', 'right')]]
|
||||
|
||||
|
||||
class CarlaHighres(StereoDataset):
|
||||
def __init__(self, aug_params=None, root='./datasets/HRVS/carla-highres'):
|
||||
super().__init__(aug_params)
|
||||
assert os.path.exists(root)
|
||||
|
||||
image1_list = sorted(glob(osp.join(root, 'trainingF/*/im0.png')))
|
||||
image2_list = [e.replace('im0', 'im1') for e in image1_list]
|
||||
disp1_list = [e.replace('im0.png', 'disp0GT.pfm') for e in image1_list]
|
||||
disp2_list = [e.replace('im1.png', 'disp1GT.pfm') for e in image2_list]
|
||||
|
||||
for img1, img2, disp1, disp2 in zip(image1_list, image2_list, disp1_list, disp2_list):
|
||||
self.image_list += [[img1, img2]]
|
||||
self.disparity_list += [[disp1, disp2]]
|
||||
|
||||
|
||||
class InStereo2K(StereoDataset):
|
||||
def __init__(self, aug_params=None, root='./datasets/InStereo2K', split='training'):
|
||||
super(InStereo2K, self).__init__(aug_params, sparse=True, reader=frame_utils.readDispInStereo2K, is_eval=split!="training")
|
||||
if split == "training":
|
||||
image1_list = sorted(glob(osp.join(root, 'part*/*/left.png')))
|
||||
else:
|
||||
image1_list = sorted(glob(osp.join(root, 'test/*/left.png')))
|
||||
|
||||
image2_list = [e.replace('left', 'right') for e in image1_list]
|
||||
disp_list = [e.replace('left', 'left_disp') for e in image1_list]
|
||||
|
||||
for img1, img2, disp in zip(image1_list, image2_list, disp_list):
|
||||
self.image_list += [[img1, img2]]
|
||||
self.disparity_list += [[disp, disp.replace('left', 'right')]]
|
||||
|
||||
|
||||
class CreStereo(StereoDataset):
|
||||
def __init__(self, aug_params=None, root='./datasets/CreStereo'):
|
||||
super(CreStereo, self).__init__(aug_params, reader=frame_utils.readDispCreStereo)
|
||||
|
||||
image1_list = sorted(glob(osp.join(root, '*/*_left.jpg')))
|
||||
image2_list = [e.replace('left', 'right') for e in image1_list]
|
||||
disp_list = [e.replace('_left.jpg', '_left.disp.png') for e in image1_list]
|
||||
|
||||
for img1, img2, disp in zip(image1_list, image2_list, disp_list):
|
||||
self.image_list += [[img1, img2]]
|
||||
self.disparity_list += [[disp, disp.replace('left', 'right')]]
|
||||
|
||||
|
||||
class IRS(StereoDataset):
|
||||
def __init__(self, aug_params=None, root='./datasets/IRSDataset'):
|
||||
super().__init__(aug_params)
|
||||
image1_list = sorted(glob(osp.join(root, '*/*/l_*.png')))
|
||||
image2_list = sorted(glob(osp.join(root, '*/*/r_*.png')))
|
||||
disp_list = sorted(glob(osp.join(root, '*/*/d_*.pfm')))
|
||||
for img1, img2, disp in zip(image1_list, image2_list, disp_list):
|
||||
assert img1.split('/')[-2] == disp.split('/')[-2]
|
||||
assert img1.split('.')[0].split('_')[-1] == disp.split('.')[0].split('_')[-1]
|
||||
if 'QAOfficeAndSecurityRoom2_Night' in img1: # bad scenes
|
||||
continue
|
||||
self.image_list += [[img1, img2]]
|
||||
self.disparity_list += [[disp]]
|
||||
|
||||
|
||||
class Booster(StereoDataset):
|
||||
def __init__(self, aug_params=None, root='./datasets/Booster_Dataset', split='train', is_eval=False, is_test=False):
|
||||
super().__init__(aug_params, sparse=True, reader=frame_utils.readDispBooster, is_eval=is_eval, is_test=is_test)
|
||||
assert os.path.exists(root)
|
||||
|
||||
folder_list = sorted(glob(osp.join(root, split+'/balanced/*')))
|
||||
for folder in folder_list:
|
||||
image1_list = sorted(glob(osp.join(folder, 'camera_00/im*.png')))
|
||||
image2_list = sorted(glob(osp.join(folder, 'camera_02/im*.png')))
|
||||
if split=="train":
|
||||
for img1 in image1_list:
|
||||
for img2 in image2_list:
|
||||
self.image_list += [[img1, img2]]
|
||||
self.disparity_list += [[osp.join(folder, 'disp_00.npy'), osp.join(folder, 'disp_02.npy')]]
|
||||
else:
|
||||
for img1, img2 in zip(image1_list, image2_list):
|
||||
self.image_list += [[img1, img2]]
|
||||
self.disparity_list += [[osp.join(folder, 'disp_00.npy'), osp.join(folder, 'disp_02.npy')]]
|
||||
|
||||
|
||||
class ThreeDKenBurns(StereoDataset):
|
||||
def __init__(self, aug_params=None, root='./datasets/3dkenburns'):
|
||||
super().__init__(aug_params, reader=frame_utils.readDisp3DKenBurns)
|
||||
|
||||
image1_list = sorted(glob(osp.join(root, '*/*l-image.png')))
|
||||
image2_list = sorted(glob(osp.join(root, '*/*r-image.png')))
|
||||
|
||||
disp1_list = sorted(glob(osp.join(root, '*/*l-depth.exr')))
|
||||
disp2_list = sorted(glob(osp.join(root, '*/*r-depth.exr')))
|
||||
|
||||
for img1, img2, disp1, disp2 in zip(image1_list, image2_list, disp1_list, disp2_list):
|
||||
self.image_list += [[img1, img2]]
|
||||
self.disparity_list += [[disp1, disp2]]
|
||||
|
||||
|
||||
class VKITTI2(StereoDataset):
|
||||
def __init__(self, aug_params=None, root='./datasets/VKITTI2'):
|
||||
super().__init__(aug_params, reader=frame_utils.readDispVKITTI2)
|
||||
|
||||
image1_list = sorted(glob(osp.join(root, 'Scene*/*/frames/rgb/Camera_0/rgb_*.jpg')))
|
||||
image2_list = sorted(glob(osp.join(root, 'Scene*/*/frames/rgb/Camera_1/rgb_*.jpg')))
|
||||
|
||||
disp1_list = sorted(glob(osp.join(root, 'Scene*/*/frames/depth/Camera_0/depth_*.png')))
|
||||
disp2_list = sorted(glob(osp.join(root, 'Scene*/*/frames/depth/Camera_1/depth_*.png')))
|
||||
|
||||
for img1, img2, disp1, disp2 in zip(image1_list, image2_list, disp1_list, disp2_list):
|
||||
self.image_list += [[img1, img2]]
|
||||
self.disparity_list += [[disp1, disp2]]
|
||||
|
||||
|
||||
def fetch_dataloader(args):
|
||||
""" Create the data loader for the corresponding trainign set """
|
||||
|
||||
aug_params = {'crop_size': args.image_size, 'min_scale': args.spatial_scale[0],
|
||||
'max_scale': args.spatial_scale[1], 'do_flip': False, 'yjitter': not args.noyjitter}
|
||||
if hasattr(args, "saturation_range") and args.saturation_range is not None:
|
||||
aug_params["saturation_range"] = args.saturation_range
|
||||
if hasattr(args, "img_gamma") and args.img_gamma is not None:
|
||||
aug_params["gamma"] = args.img_gamma
|
||||
if hasattr(args, "do_flip") and args.do_flip is not None:
|
||||
aug_params["do_flip"] = args.do_flip
|
||||
|
||||
assert len(args.train_datasets) == len(args.train_folds)
|
||||
|
||||
train_dataset = None
|
||||
for fold, dataset_name in zip(args.train_folds, args.train_datasets):
|
||||
if dataset_name.startswith("middlebury_"):
|
||||
new_dataset = Middlebury(aug_params, split=dataset_name.replace('middlebury_','')) * fold
|
||||
elif dataset_name == 'sceneflow':
|
||||
clean_dataset = SceneFlowDatasets(aug_params, dstype='frames_cleanpass')
|
||||
final_dataset = SceneFlowDatasets(aug_params, dstype='frames_finalpass')
|
||||
new_dataset = clean_dataset*fold+final_dataset*fold
|
||||
logging.info(f"Adding {len(new_dataset)} samples from SceneFlow")
|
||||
elif 'kitti1' in dataset_name:
|
||||
new_dataset = KITTI(aug_params, split=dataset_name[-2:], image_set='training') * fold
|
||||
logging.info(f"Adding {len(new_dataset)} samples from KITTI"+dataset_name[-2:])
|
||||
elif 'eth3d' in dataset_name:
|
||||
new_dataset = ETH3D(aug_params, split='training') * fold
|
||||
logging.info(f"Adding {len(new_dataset)} samples from ETH3D")
|
||||
elif dataset_name == 'sintel_stereo':
|
||||
new_dataset = SintelStereo(aug_params) * fold
|
||||
logging.info(f"Adding {len(new_dataset)} samples from Sintel Stereo")
|
||||
elif dataset_name == 'falling_things':
|
||||
new_dataset = FallingThings(aug_params)*fold
|
||||
logging.info(f"Adding {len(new_dataset)} samples from FallingThings")
|
||||
elif dataset_name.startswith('tartan_air'):
|
||||
new_dataset = TartanAir(aug_params) * fold
|
||||
logging.info(f"Adding {len(new_dataset)} samples from Tartain Air")
|
||||
elif dataset_name.startswith('carla_highres'):
|
||||
new_dataset = CarlaHighres(aug_params) * fold
|
||||
logging.info(f"Adding {len(new_dataset)} samples from Carla Highres")
|
||||
elif dataset_name.startswith('irs'):
|
||||
new_dataset = IRS(aug_params) * fold
|
||||
logging.info(f"Adding {len(new_dataset)} samples from IRS")
|
||||
elif dataset_name.startswith('crestereo'):
|
||||
new_dataset = CreStereo(aug_params) * fold
|
||||
logging.info(f"Adding {len(new_dataset)} samples from CreStereo")
|
||||
elif dataset_name.startswith('instereo2k'):
|
||||
new_dataset = InStereo2K(aug_params) * fold
|
||||
logging.info(f"Adding {len(new_dataset)} samples from InStereo2K")
|
||||
elif dataset_name.startswith('booster'):
|
||||
new_dataset = Booster(aug_params) * fold
|
||||
logging.info(f"Adding {len(new_dataset)} samples from Booster")
|
||||
elif dataset_name.startswith('3dkenburns'):
|
||||
new_dataset = ThreeDKenBurns(aug_params) * fold
|
||||
logging.info(f"Adding {len(new_dataset)} samples from 3D Ken Burns")
|
||||
elif dataset_name.startswith('vkitti2'):
|
||||
new_dataset = VKITTI2(aug_params) * fold
|
||||
logging.info(f"Adding {len(new_dataset)} samples from VKITTI2")
|
||||
|
||||
train_dataset = new_dataset if train_dataset is None else train_dataset + new_dataset
|
||||
|
||||
train_loader = data.DataLoader(train_dataset, batch_size=args.batch_size,
|
||||
pin_memory=True, shuffle=True, num_workers=int(os.environ.get('SLURM_CPUS_PER_TASK', 6))-2, drop_last=True)
|
||||
|
||||
logging.info('Training with %d image pairs' % len(train_dataset))
|
||||
return train_loader
|
||||
|
||||
|
||||
def fetch_dataset(args):
|
||||
""" Create the dataset for the corresponding training set """
|
||||
|
||||
aug_params = {'crop_size': args.image_size, 'min_scale': args.spatial_scale[0],
|
||||
'max_scale': args.spatial_scale[1], 'do_flip': False, 'yjitter': not args.noyjitter}
|
||||
if hasattr(args, "saturation_range") and args.saturation_range is not None:
|
||||
aug_params["saturation_range"] = args.saturation_range
|
||||
if hasattr(args, "img_gamma") and args.img_gamma is not None:
|
||||
aug_params["gamma"] = args.img_gamma
|
||||
if hasattr(args, "do_flip") and args.do_flip is not None:
|
||||
aug_params["do_flip"] = args.do_flip
|
||||
|
||||
assert len(args.train_datasets) == len(args.train_folds)
|
||||
|
||||
train_dataset = None
|
||||
for fold, dataset_name in zip(args.train_folds, args.train_datasets):
|
||||
if dataset_name.startswith("middlebury_"):
|
||||
new_dataset = Middlebury(aug_params, split=dataset_name.replace('middlebury_', '')) * fold
|
||||
logging.info(f"Adding {len(new_dataset)} samples from {dataset_name}")
|
||||
elif 'eth3d' in dataset_name:
|
||||
new_dataset = ETH3D(aug_params, split='training') * fold
|
||||
logging.info(f"Adding {len(new_dataset)} samples from ETH3D")
|
||||
elif 'kitti1' in dataset_name:
|
||||
new_dataset = KITTI(aug_params, split=dataset_name[-2:], image_set='training') * fold
|
||||
logging.info(f"Adding {len(new_dataset)} samples from KITTI"+dataset_name[-2:])
|
||||
elif dataset_name == 'sceneflow':
|
||||
clean_dataset = SceneFlowDatasets(aug_params, dstype='frames_cleanpass')
|
||||
final_dataset = SceneFlowDatasets(aug_params, dstype='frames_finalpass')
|
||||
new_dataset = clean_dataset*fold+final_dataset*fold
|
||||
logging.info(f"Adding {len(new_dataset)} samples from SceneFlow")
|
||||
elif dataset_name == 'sintel_stereo':
|
||||
new_dataset = SintelStereo(aug_params) * fold
|
||||
logging.info(f"Adding {len(new_dataset)} samples from Sintel Stereo")
|
||||
elif dataset_name == 'falling_things':
|
||||
new_dataset = FallingThings(aug_params) * fold
|
||||
logging.info(f"Adding {len(new_dataset)} samples from FallingThings")
|
||||
elif dataset_name.startswith('tartan_air'):
|
||||
new_dataset = TartanAir(aug_params) * fold
|
||||
logging.info(f"Adding {len(new_dataset)} samples from Tartain Air")
|
||||
elif dataset_name.startswith('carla_highres'):
|
||||
new_dataset = CarlaHighres(aug_params) * fold
|
||||
logging.info(f"Adding {len(new_dataset)} samples from Carla Highres")
|
||||
elif dataset_name.startswith('irs'):
|
||||
new_dataset = IRS(aug_params) * fold
|
||||
logging.info(f"Adding {len(new_dataset)} samples from IRS")
|
||||
elif dataset_name.startswith('crestereo'):
|
||||
new_dataset = CreStereo(aug_params) * fold
|
||||
logging.info(f"Adding {len(new_dataset)} samples from CreStereo")
|
||||
elif dataset_name.startswith('instereo2k'):
|
||||
new_dataset = InStereo2K(aug_params) * fold
|
||||
logging.info(f"Adding {len(new_dataset)} samples from InStereo2K")
|
||||
elif dataset_name.startswith('booster'):
|
||||
new_dataset = Booster(aug_params) * fold
|
||||
logging.info(f"Adding {len(new_dataset)} samples from Booster")
|
||||
elif dataset_name.startswith('3dkenburns'):
|
||||
new_dataset = ThreeDKenBurns(aug_params) * fold
|
||||
logging.info(f"Adding {len(new_dataset)} samples from 3D Ken Burns")
|
||||
elif dataset_name.startswith('vkitti2'):
|
||||
new_dataset = VKITTI2(aug_params) * fold
|
||||
logging.info(f"Adding {len(new_dataset)} samples from VKITTI2")
|
||||
|
||||
train_dataset = new_dataset if train_dataset is None else train_dataset + new_dataset
|
||||
|
||||
logging.info('Training with %d image pairs' % len(train_dataset))
|
||||
return train_dataset
|
||||
|
||||
@@ -0,0 +1,195 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from opt_einsum import contract
|
||||
|
||||
|
||||
class DispHead(nn.Module):
|
||||
def __init__(self, input_dim=128, hidden_dim=256, output_dim=1):
|
||||
super(DispHead, self).__init__()
|
||||
self.conv1 = nn.Conv2d(input_dim, hidden_dim, 3, padding=1)
|
||||
self.conv2 = nn.Conv2d(hidden_dim, output_dim, 3, padding=1)
|
||||
self.relu = nn.ReLU(inplace=True)
|
||||
|
||||
def forward(self, x):
|
||||
return self.conv2(self.relu(self.conv1(x)))
|
||||
|
||||
|
||||
class ConvGRU(nn.Module):
|
||||
def __init__(self, hidden_dim, input_dim, kernel_size=3):
|
||||
super(ConvGRU, self).__init__()
|
||||
self.convz = nn.Conv2d(hidden_dim+input_dim, hidden_dim, kernel_size,
|
||||
padding=kernel_size//2)
|
||||
self.convr = nn.Conv2d(hidden_dim+input_dim, hidden_dim, kernel_size,
|
||||
padding=kernel_size//2)
|
||||
self.convq = nn.Conv2d(hidden_dim+input_dim, hidden_dim, kernel_size,
|
||||
padding=kernel_size//2)
|
||||
|
||||
def forward(self, h, cz, cr, cq, *x_list):
|
||||
x = torch.cat(x_list, dim=1)
|
||||
hx = torch.cat([h, x], dim=1)
|
||||
|
||||
z = torch.sigmoid(self.convz(hx) + cz)
|
||||
r = torch.sigmoid(self.convr(hx) + cr)
|
||||
q = torch.tanh(self.convq(torch.cat([r*h, x], dim=1)) + cq)
|
||||
|
||||
h = (1-z) * h + z * q
|
||||
return h
|
||||
|
||||
|
||||
class SepConvGRU(nn.Module):
|
||||
def __init__(self, hidden_dim=128, input_dim=192+128):
|
||||
super(SepConvGRU, self).__init__()
|
||||
self.convz1 = nn.Conv2d(hidden_dim+input_dim, hidden_dim, (1,5), padding=(0,2))
|
||||
self.convr1 = nn.Conv2d(hidden_dim+input_dim, hidden_dim, (1,5), padding=(0,2))
|
||||
self.convq1 = nn.Conv2d(hidden_dim+input_dim, hidden_dim, (1,5), padding=(0,2))
|
||||
|
||||
self.convz2 = nn.Conv2d(hidden_dim+input_dim, hidden_dim, (5,1), padding=(2,0))
|
||||
self.convr2 = nn.Conv2d(hidden_dim+input_dim, hidden_dim, (5,1), padding=(2,0))
|
||||
self.convq2 = nn.Conv2d(hidden_dim+input_dim, hidden_dim, (5,1), padding=(2,0))
|
||||
|
||||
def forward(self, h, *x):
|
||||
# horizontal
|
||||
x = torch.cat(x, dim=1)
|
||||
hx = torch.cat([h, x], dim=1)
|
||||
z = torch.sigmoid(self.convz1(hx))
|
||||
r = torch.sigmoid(self.convr1(hx))
|
||||
q = torch.tanh(self.convq1(torch.cat([r*h, x], dim=1)))
|
||||
h = (1-z) * h + z * q
|
||||
|
||||
# vertical
|
||||
hx = torch.cat([h, x], dim=1)
|
||||
z = torch.sigmoid(self.convz2(hx))
|
||||
r = torch.sigmoid(self.convr2(hx))
|
||||
q = torch.tanh(self.convq2(torch.cat([r*h, x], dim=1)))
|
||||
h = (1-z) * h + z * q
|
||||
|
||||
return h
|
||||
|
||||
|
||||
class BasicMotionEncoder(nn.Module):
|
||||
def __init__(self, cor_planes, c1_planes=64, c2_planes=64, f1_planes=64, f2_planes=64, out_planes=128):
|
||||
super(BasicMotionEncoder, self).__init__()
|
||||
|
||||
self.convc1 = nn.Conv2d(cor_planes, c1_planes, 1, padding=0)
|
||||
self.convc2 = nn.Conv2d(c1_planes, c2_planes, 3, padding=1)
|
||||
self.convd1 = nn.Conv2d(1, f1_planes, 7, padding=3)
|
||||
self.convd2 = nn.Conv2d(f1_planes, f2_planes, 3, padding=1)
|
||||
self.conv = nn.Conv2d(c2_planes+f2_planes, out_planes-1, 3, padding=1)
|
||||
|
||||
def forward(self, disp, corr):
|
||||
cor = F.relu(self.convc1(corr))
|
||||
cor = F.relu(self.convc2(cor))
|
||||
dis = F.relu(self.convd1(disp))
|
||||
dis = F.relu(self.convd2(dis))
|
||||
|
||||
cor_dis = torch.cat([cor, dis], dim=1)
|
||||
out = F.relu(self.conv(cor_dis))
|
||||
return torch.cat([out, disp], dim=1)
|
||||
|
||||
|
||||
def pool2x(x):
|
||||
return F.avg_pool2d(x, 3, stride=2, padding=1)
|
||||
|
||||
|
||||
def pool4x(x):
|
||||
return F.avg_pool2d(x, 5, stride=4, padding=1)
|
||||
|
||||
|
||||
def interp(x, dest):
|
||||
interp_args = {'mode': 'bilinear', 'align_corners': True}
|
||||
return F.interpolate(x, dest.shape[2:], **interp_args)
|
||||
|
||||
|
||||
# for RAFT-Stereo
|
||||
class BasicMultiUpdateBlock(nn.Module):
|
||||
def __init__(self, args, hidden_dims=[128, 128, 128]):
|
||||
super().__init__()
|
||||
self.args = args
|
||||
encoder_output_dim = 128
|
||||
cor_planes = args.corr_levels * (2*args.corr_radius + 1)
|
||||
self.encoder = BasicMotionEncoder(cor_planes, out_planes=encoder_output_dim)
|
||||
|
||||
self.gru08 = ConvGRU(hidden_dims[2], encoder_output_dim + hidden_dims[1] * (args.n_gru_layers > 1))
|
||||
self.gru16 = ConvGRU(hidden_dims[1], hidden_dims[0] * (args.n_gru_layers == 3) + hidden_dims[2])
|
||||
self.gru32 = ConvGRU(hidden_dims[0], hidden_dims[1])
|
||||
self.disp_head = DispHead(hidden_dims[2], hidden_dim=256, output_dim=1)
|
||||
|
||||
factor = 2**self.args.n_downsample
|
||||
|
||||
self.mask = nn.Sequential(
|
||||
nn.Conv2d(hidden_dims[2], 256, 3, padding=1),
|
||||
nn.ReLU(inplace=True),
|
||||
nn.Conv2d(256, (factor**2)*9, 1, padding=0))
|
||||
|
||||
def forward(self, net, inp, corr=None, disp=None, iter08=True, iter16=True, iter32=True, update=True):
|
||||
|
||||
if iter32:
|
||||
net[2] = self.gru32(net[2], *(inp[2]), pool2x(net[1]))
|
||||
if iter16:
|
||||
if self.args.n_gru_layers > 2:
|
||||
net[1] = self.gru16(net[1], *(inp[1]), pool2x(net[0]), interp(net[2], net[1]))
|
||||
else:
|
||||
net[1] = self.gru16(net[1], *(inp[1]), pool2x(net[0]))
|
||||
if iter08:
|
||||
motion_features = self.encoder(disp, corr)
|
||||
if self.args.n_gru_layers > 1:
|
||||
net[0] = self.gru08(net[0], *(inp[0]), motion_features, interp(net[1], net[0]))
|
||||
else:
|
||||
net[0] = self.gru08(net[0], *(inp[0]), motion_features)
|
||||
|
||||
if not update:
|
||||
return net
|
||||
|
||||
delta_disp = self.disp_head(net[0])
|
||||
|
||||
# scale mask to balence gradients
|
||||
mask = .25 * self.mask(net[0])
|
||||
return net, mask, delta_disp
|
||||
|
||||
|
||||
class ScaleBasicMultiUpdateBlock(nn.Module):
|
||||
def __init__(self, args, hidden_dims=[128, 128, 128]):
|
||||
super().__init__()
|
||||
self.args = args
|
||||
encoder_output_dim = 128
|
||||
cor_planes = len(args.scale_list) * (2*args.scale_corr_radius + 1)
|
||||
self.encoder = BasicMotionEncoder(cor_planes, out_planes=encoder_output_dim)
|
||||
|
||||
self.gru08 = ConvGRU(hidden_dims[2], encoder_output_dim + hidden_dims[1] * (args.n_gru_layers > 1))
|
||||
self.gru16 = ConvGRU(hidden_dims[1], hidden_dims[0] * (args.n_gru_layers == 3) + hidden_dims[2])
|
||||
self.gru32 = ConvGRU(hidden_dims[0], hidden_dims[1])
|
||||
self.disp_head = DispHead(hidden_dims[2], hidden_dim=256, output_dim=1)
|
||||
|
||||
factor = 2**self.args.n_downsample
|
||||
|
||||
self.mask = nn.Sequential(
|
||||
nn.Conv2d(hidden_dims[2], 256, 3, padding=1),
|
||||
nn.ReLU(inplace=True),
|
||||
nn.Conv2d(256, (factor**2)*9, 1, padding=0))
|
||||
|
||||
def forward(self, net, inp, corr=None, disp=None, iter08=True, iter16=True, iter32=True, update=True):
|
||||
|
||||
if iter32:
|
||||
net[2] = self.gru32(net[2], *(inp[2]), pool2x(net[1]))
|
||||
if iter16:
|
||||
if self.args.n_gru_layers > 2:
|
||||
net[1] = self.gru16(net[1], *(inp[1]), pool2x(net[0]), interp(net[2], net[1]))
|
||||
else:
|
||||
net[1] = self.gru16(net[1], *(inp[1]), pool2x(net[0]))
|
||||
if iter08:
|
||||
motion_features = self.encoder(disp, corr)
|
||||
if self.args.n_gru_layers > 1:
|
||||
net[0] = self.gru08(net[0], *(inp[0]), motion_features, interp(net[1], net[0]))
|
||||
else:
|
||||
net[0] = self.gru08(net[0], *(inp[0]), motion_features)
|
||||
|
||||
if not update:
|
||||
return net
|
||||
|
||||
x_disp = self.disp_head(net[0])
|
||||
scale_disp = F.relu6(torch.exp(.25*x_disp))
|
||||
|
||||
# scale mask to balence gradients
|
||||
mask = .25 * self.mask(net[0])
|
||||
return net, mask, scale_disp
|
||||
@@ -0,0 +1,307 @@
|
||||
import numpy as np
|
||||
import random
|
||||
import warnings
|
||||
import os
|
||||
import time
|
||||
from glob import glob
|
||||
from skimage import color, io
|
||||
from PIL import Image
|
||||
|
||||
import cv2
|
||||
cv2.setNumThreads(0)
|
||||
cv2.ocl.setUseOpenCL(False)
|
||||
|
||||
import torch
|
||||
from torchvision.transforms import ColorJitter, functional, Compose
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
def get_middlebury_images():
|
||||
root = "../datasets/Middlebury/MiddEval3"
|
||||
with open(os.path.join(root, "official_train.txt"), 'r') as f:
|
||||
lines = f.read().splitlines()
|
||||
return sorted([os.path.join(root, 'trainingQ', f'{name}/im0.png') for name in lines])
|
||||
|
||||
|
||||
def get_eth3d_images():
|
||||
return sorted(glob('../datasets/ETH3D/two_view_training/*/im0.png'))
|
||||
|
||||
|
||||
def get_kitti_images():
|
||||
return sorted(glob('..datasets/KITTI/training/image_2/*_10.png'))
|
||||
|
||||
|
||||
def transfer_color(image, style_mean, style_stddev):
|
||||
reference_image_lab = color.rgb2lab(image)
|
||||
reference_stddev = np.std(reference_image_lab, axis=(0, 1), keepdims=True)# + 1
|
||||
reference_mean = np.mean(reference_image_lab, axis=(0, 1), keepdims=True)
|
||||
|
||||
reference_image_lab = reference_image_lab - reference_mean
|
||||
lamb = style_stddev/reference_stddev
|
||||
style_image_lab = lamb * reference_image_lab
|
||||
output_image_lab = style_image_lab + style_mean
|
||||
l, a, b = np.split(output_image_lab, 3, axis=2)
|
||||
l = l.clip(0, 100)
|
||||
output_image_lab = np.concatenate((l, a, b), axis=2)
|
||||
with warnings.catch_warnings():
|
||||
warnings.simplefilter("ignore", category=UserWarning)
|
||||
output_image_rgb = color.lab2rgb(output_image_lab) * 255
|
||||
return output_image_rgb
|
||||
|
||||
|
||||
class AdjustGamma(object):
|
||||
|
||||
def __init__(self, gamma_min, gamma_max, gain_min=1.0, gain_max=1.0):
|
||||
self.gamma_min, self.gamma_max, self.gain_min, self.gain_max = gamma_min, gamma_max, gain_min, gain_max
|
||||
|
||||
def __call__(self, sample):
|
||||
gain = random.uniform(self.gain_min, self.gain_max)
|
||||
gamma = random.uniform(self.gamma_min, self.gamma_max)
|
||||
return functional.adjust_gamma(sample, gamma, gain)
|
||||
|
||||
def __repr__(self):
|
||||
return f"Adjust Gamma {self.gamma_min}, ({self.gamma_max}) and Gain ({self.gain_min}, {self.gain_max})"
|
||||
|
||||
|
||||
class DispAugmentor:
|
||||
def __init__(self, crop_size, min_scale=-0.2, max_scale=0.5, do_flip=True, yjitter=False,
|
||||
saturation_range=[0.6, 1.4], gamma=[1, 1, 1, 1]):
|
||||
|
||||
# spatial augmentation params
|
||||
self.crop_size = crop_size
|
||||
self.min_scale = min_scale
|
||||
self.max_scale = max_scale
|
||||
self.spatial_aug_prob = 1.0
|
||||
self.stretch_prob = 0.8
|
||||
self.max_stretch = 0.2
|
||||
|
||||
# flip augmentation params
|
||||
self.yjitter = yjitter
|
||||
self.do_flip = do_flip
|
||||
self.v_flip_prob = 0.1
|
||||
|
||||
# photometric augmentation params
|
||||
self.photo_aug = Compose([ColorJitter(brightness=0.4, contrast=0.4, saturation=saturation_range, hue=0.5/3.14), AdjustGamma(*gamma)])
|
||||
self.asymmetric_color_aug_prob = 0.2
|
||||
self.eraser_aug_prob = 0.5
|
||||
|
||||
def color_transform(self, img1, img2):
|
||||
""" Photometric augmentation """
|
||||
|
||||
# asymmetric
|
||||
if np.random.rand() < self.asymmetric_color_aug_prob:
|
||||
img1 = np.array(self.photo_aug(Image.fromarray(img1)), dtype=np.uint8)
|
||||
img2 = np.array(self.photo_aug(Image.fromarray(img2)), dtype=np.uint8)
|
||||
|
||||
# symmetric
|
||||
else:
|
||||
image_stack = np.concatenate([img1, img2], axis=0)
|
||||
image_stack = np.array(self.photo_aug(Image.fromarray(image_stack)), dtype=np.uint8)
|
||||
img1, img2 = np.split(image_stack, 2, axis=0)
|
||||
|
||||
return img1, img2
|
||||
|
||||
def eraser_transform(self, img1, img2, bounds=[50, 100]):
|
||||
""" Occlusion augmentation """
|
||||
|
||||
ht, wd = img1.shape[:2]
|
||||
if np.random.rand() < self.eraser_aug_prob:
|
||||
mean_color = np.mean(img2.reshape(-1, 3), axis=0)
|
||||
for _ in range(np.random.randint(1, 3)):
|
||||
x0 = np.random.randint(0, wd)
|
||||
y0 = np.random.randint(0, ht)
|
||||
dx = np.random.randint(bounds[0], bounds[1])
|
||||
dy = np.random.randint(bounds[0], bounds[1])
|
||||
img2[y0:y0 + dy, x0:x0 + dx, :] = mean_color
|
||||
|
||||
return img1, img2
|
||||
|
||||
def spatial_transform(self, img1, img2, disp):
|
||||
# randomly sample scale
|
||||
ht, wd = img1.shape[:2]
|
||||
min_scale = np.maximum(
|
||||
(self.crop_size[0] + 8) / float(ht),
|
||||
(self.crop_size[1] + 8) / float(wd))
|
||||
|
||||
scale = 2 ** np.random.uniform(self.min_scale, self.max_scale)
|
||||
if scale>min_scale:
|
||||
scale = np.random.uniform(min_scale, scale)
|
||||
scale_x = scale
|
||||
scale_y = scale
|
||||
if np.random.rand() < self.stretch_prob:
|
||||
scale_x *= 2 ** np.random.uniform(-self.max_stretch, self.max_stretch)
|
||||
scale_y *= 2 ** np.random.uniform(-self.max_stretch, self.max_stretch)
|
||||
|
||||
scale_x = np.clip(scale_x, min_scale, 2*min_scale)
|
||||
scale_y = np.clip(scale_y, min_scale, 2*min_scale)
|
||||
|
||||
if np.random.rand() < self.spatial_aug_prob or min_scale >= 1.0:
|
||||
# rescale the images
|
||||
img1 = cv2.resize(img1, None, fx=scale_x, fy=scale_y, interpolation=cv2.INTER_LINEAR)
|
||||
img2 = cv2.resize(img2, None, fx=scale_x, fy=scale_y, interpolation=cv2.INTER_LINEAR)
|
||||
disp = cv2.resize(disp, None, fx=scale_x, fy=scale_y, interpolation=cv2.INTER_LINEAR)
|
||||
disp = disp * scale_x
|
||||
|
||||
if self.do_flip:
|
||||
if np.random.rand() < self.v_flip_prob and self.do_flip == 'v': # v-flip
|
||||
img1 = img1[::-1, :]
|
||||
img2 = img2[::-1, :]
|
||||
disp = disp[::-1, :]
|
||||
|
||||
if self.yjitter:
|
||||
y0 = np.random.randint(2, img1.shape[0] - self.crop_size[0] - 2)
|
||||
x0 = np.random.randint(0, img1.shape[1] - self.crop_size[1] - 0)
|
||||
|
||||
y1 = y0 + np.random.randint(-2, 2 + 1)
|
||||
y1 = np.clip(y1, 0, img1.shape[0] - self.crop_size[0])
|
||||
img1 = img1[y0:y0 + self.crop_size[0], x0:x0 + self.crop_size[1]]
|
||||
img2 = img2[y1:y1 + self.crop_size[0], x0:x0 + self.crop_size[1]]
|
||||
disp = disp[y0:y0 + self.crop_size[0], x0:x0 + self.crop_size[1]]
|
||||
|
||||
else:
|
||||
y0 = np.random.randint(0, img1.shape[0] - self.crop_size[0])
|
||||
x0 = np.random.randint(0, img1.shape[1] - self.crop_size[1])
|
||||
|
||||
img1 = img1[y0:y0 + self.crop_size[0], x0:x0 + self.crop_size[1]]
|
||||
img2 = img2[y0:y0 + self.crop_size[0], x0:x0 + self.crop_size[1]]
|
||||
disp = disp[y0:y0 + self.crop_size[0], x0:x0 + self.crop_size[1]]
|
||||
|
||||
return img1, img2, disp
|
||||
|
||||
def __call__(self, img1, img2, disp):
|
||||
img1, img2 = self.color_transform(img1, img2)
|
||||
img1, img2 = self.eraser_transform(img1, img2)
|
||||
img1, img2, disp = self.spatial_transform(img1, img2, disp)
|
||||
|
||||
img1 = np.ascontiguousarray(img1)
|
||||
img2 = np.ascontiguousarray(img2)
|
||||
disp = np.ascontiguousarray(disp)
|
||||
|
||||
return img1, img2, disp
|
||||
|
||||
|
||||
class SparseDispAugmentor:
|
||||
def __init__(self, crop_size, min_scale=-0.2, max_scale=0.5, do_flip=False, yjitter=False,
|
||||
saturation_range=[0.7, 1.3], gamma=[1, 1, 1, 1]):
|
||||
# spatial augmentation params
|
||||
self.crop_size = crop_size
|
||||
self.min_scale = min_scale
|
||||
self.max_scale = max_scale
|
||||
self.spatial_aug_prob = 0.8
|
||||
self.stretch_prob = 0.8
|
||||
self.max_stretch = 0.2
|
||||
|
||||
# flip augmentation params
|
||||
self.do_flip = do_flip
|
||||
self.v_flip_prob = 0.1
|
||||
|
||||
# photometric augmentation params
|
||||
self.photo_aug = Compose(
|
||||
[ColorJitter(brightness=0.3, contrast=0.3, saturation=saturation_range, hue=0.3/3.14),
|
||||
AdjustGamma(*gamma)])
|
||||
self.asymmetric_color_aug_prob = 0.2
|
||||
self.eraser_aug_prob = 0.5
|
||||
|
||||
def color_transform(self, img1, img2):
|
||||
image_stack = np.concatenate([img1, img2], axis=0)
|
||||
image_stack = np.array(self.photo_aug(Image.fromarray(image_stack)), dtype=np.uint8)
|
||||
img1, img2 = np.split(image_stack, 2, axis=0)
|
||||
return img1, img2
|
||||
|
||||
def eraser_transform(self, img1, img2):
|
||||
ht, wd = img1.shape[:2]
|
||||
if np.random.rand() < self.eraser_aug_prob:
|
||||
mean_color = np.mean(img2.reshape(-1, 3), axis=0)
|
||||
for _ in range(np.random.randint(1, 3)):
|
||||
x0 = np.random.randint(0, wd)
|
||||
y0 = np.random.randint(0, ht)
|
||||
dx = np.random.randint(50, 100)
|
||||
dy = np.random.randint(50, 100)
|
||||
img2[y0:y0 + dy, x0:x0 + dx, :] = mean_color
|
||||
|
||||
return img1, img2
|
||||
|
||||
def resize_sparse_flow_map(self, disp, valid, fx=1.0, fy=1.0):
|
||||
ht, wd = disp.shape[:2]
|
||||
coords = np.meshgrid(np.arange(wd), np.arange(ht))
|
||||
coords = np.stack(coords, axis=-1)
|
||||
|
||||
coords = coords.reshape(-1, 2).astype(np.float32)
|
||||
disp = disp.reshape(-1).astype(np.float32)
|
||||
valid = valid.reshape(-1).astype(np.float32)
|
||||
|
||||
coords0 = coords[valid >= 1]
|
||||
disp0 = disp[valid >= 1]
|
||||
|
||||
ht1 = int(round(ht * fy))
|
||||
wd1 = int(round(wd * fx))
|
||||
|
||||
coords1 = coords0 * [fx, fy]
|
||||
disp1 = disp0 * fx
|
||||
|
||||
xx = np.round(coords1[:, 0]).astype(np.int32)
|
||||
yy = np.round(coords1[:, 1]).astype(np.int32)
|
||||
|
||||
v = (xx > 0) & (xx < wd1) & (yy > 0) & (yy < ht1)
|
||||
xx = xx[v]
|
||||
yy = yy[v]
|
||||
disp1 = disp1[v]
|
||||
|
||||
disp_img = np.zeros([ht1, wd1], dtype=np.float32)
|
||||
valid_img = np.zeros([ht1, wd1], dtype=np.int32)
|
||||
|
||||
disp_img[yy, xx] = disp1
|
||||
valid_img[yy, xx] = 1
|
||||
|
||||
return disp_img, valid_img
|
||||
|
||||
def spatial_transform(self, img1, img2, disp, valid):
|
||||
# randomly sample scale
|
||||
|
||||
ht, wd = img1.shape[:2]
|
||||
min_scale = np.maximum(
|
||||
(self.crop_size[0] + 1) / float(ht),
|
||||
(self.crop_size[1] + 1) / float(wd))
|
||||
|
||||
scale = 2 ** np.random.uniform(self.min_scale, self.max_scale)
|
||||
if scale>min_scale:
|
||||
scale = np.random.uniform(min_scale, 2*min_scale)
|
||||
scale_x = scale
|
||||
scale_y = scale
|
||||
|
||||
scale_x = np.clip(scale_x, min_scale, 2*min_scale)
|
||||
scale_y = np.clip(scale_y, min_scale, 2*min_scale)
|
||||
|
||||
if np.random.rand() < self.spatial_aug_prob or min_scale >= 1.0:
|
||||
# rescale the images
|
||||
img1 = cv2.resize(img1, None, fx=scale_x, fy=scale_y, interpolation=cv2.INTER_LINEAR)
|
||||
img2 = cv2.resize(img2, None, fx=scale_x, fy=scale_y, interpolation=cv2.INTER_LINEAR)
|
||||
disp, valid = self.resize_sparse_flow_map(disp, valid, fx=scale_x, fy=scale_y)
|
||||
|
||||
if self.do_flip:
|
||||
if np.random.rand() < self.v_flip_prob and self.do_flip == 'v': # v-flip
|
||||
img1 = img1[::-1, :]
|
||||
img2 = img2[::-1, :]
|
||||
disp = disp[::-1, :]
|
||||
valid = valid[::-1, :]
|
||||
|
||||
y0 = np.random.randint(0, img1.shape[0] - self.crop_size[0])
|
||||
x0 = np.random.randint(0, img1.shape[1] - self.crop_size[1])
|
||||
|
||||
img1 = img1[y0:y0 + self.crop_size[0], x0:x0 + self.crop_size[1]]
|
||||
img2 = img2[y0:y0 + self.crop_size[0], x0:x0 + self.crop_size[1]]
|
||||
disp = disp[y0:y0 + self.crop_size[0], x0:x0 + self.crop_size[1]]
|
||||
valid = valid[y0:y0 + self.crop_size[0], x0:x0 + self.crop_size[1]]
|
||||
return img1, img2, disp, valid
|
||||
|
||||
def __call__(self, img1, img2, disp, valid):
|
||||
img1, img2 = self.color_transform(img1, img2)
|
||||
img1, img2 = self.eraser_transform(img1, img2)
|
||||
img1, img2, disp, valid = self.spatial_transform(img1, img2, disp, valid)
|
||||
|
||||
img1 = np.ascontiguousarray(img1)
|
||||
img2 = np.ascontiguousarray(img2)
|
||||
disp = np.ascontiguousarray(disp)
|
||||
valid = np.ascontiguousarray(valid)
|
||||
|
||||
return img1, img2, disp, valid
|
||||
@@ -0,0 +1,105 @@
|
||||
# Copyright (c) OpenMMLab. All rights reserved.
|
||||
# https://github.com/open-mmlab/mmcv/blob/7540cf73ac7e5d1e14d0ffbd9b6759e83929ecfc/mmcv/runner/dist_utils.py
|
||||
|
||||
import os
|
||||
import subprocess
|
||||
|
||||
import torch
|
||||
import torch.multiprocessing as mp
|
||||
from torch import distributed as dist
|
||||
|
||||
|
||||
def init_dist(launcher, backend='nccl', **kwargs):
|
||||
if mp.get_start_method(allow_none=True) is None:
|
||||
mp.set_start_method('spawn')
|
||||
if launcher == 'pytorch':
|
||||
_init_dist_pytorch(backend, **kwargs)
|
||||
elif launcher == 'mpi':
|
||||
_init_dist_mpi(backend, **kwargs)
|
||||
elif launcher == 'slurm':
|
||||
_init_dist_slurm(backend, **kwargs)
|
||||
else:
|
||||
raise ValueError(f'Invalid launcher type: {launcher}')
|
||||
|
||||
|
||||
def _init_dist_pytorch(backend, **kwargs):
|
||||
# TODO: use local_rank instead of rank % num_gpus
|
||||
rank = int(os.environ['RANK'])
|
||||
num_gpus = torch.cuda.device_count()
|
||||
torch.cuda.set_device(rank % num_gpus)
|
||||
dist.init_process_group(backend=backend, **kwargs)
|
||||
|
||||
|
||||
def _init_dist_mpi(backend, **kwargs):
|
||||
# TODO: use local_rank instead of rank % num_gpus
|
||||
rank = int(os.environ['OMPI_COMM_WORLD_RANK'])
|
||||
num_gpus = torch.cuda.device_count()
|
||||
torch.cuda.set_device(rank % num_gpus)
|
||||
dist.init_process_group(backend=backend, **kwargs)
|
||||
|
||||
|
||||
def _init_dist_slurm(backend, port=None):
|
||||
"""Initialize slurm distributed training environment.
|
||||
If argument ``port`` is not specified, then the master port will be system
|
||||
environment variable ``MASTER_PORT``. If ``MASTER_PORT`` is not in system
|
||||
environment variable, then a default port ``29500`` will be used.
|
||||
Args:
|
||||
backend (str): Backend of torch.distributed.
|
||||
port (int, optional): Master port. Defaults to None.
|
||||
"""
|
||||
proc_id = int(os.environ['SLURM_PROCID'])
|
||||
ntasks = int(os.environ['SLURM_NTASKS'])
|
||||
node_list = os.environ['SLURM_NODELIST']
|
||||
num_gpus = torch.cuda.device_count()
|
||||
torch.cuda.set_device(proc_id % num_gpus)
|
||||
addr = subprocess.getoutput(
|
||||
f'scontrol show hostname {node_list} | head -n1')
|
||||
# specify master port
|
||||
if port is not None:
|
||||
os.environ['MASTER_PORT'] = str(port)
|
||||
elif 'MASTER_PORT' in os.environ:
|
||||
pass # use MASTER_PORT in the environment variable
|
||||
else:
|
||||
# 29500 is torch.distributed default port
|
||||
os.environ['MASTER_PORT'] = '29500'
|
||||
# use MASTER_ADDR in the environment variable if it already exists
|
||||
if 'MASTER_ADDR' not in os.environ:
|
||||
os.environ['MASTER_ADDR'] = addr
|
||||
os.environ['WORLD_SIZE'] = str(ntasks)
|
||||
os.environ['LOCAL_RANK'] = str(proc_id % num_gpus)
|
||||
os.environ['RANK'] = str(proc_id)
|
||||
dist.init_process_group(backend=backend)
|
||||
|
||||
|
||||
def get_dist_info():
|
||||
# if (TORCH_VERSION != 'parrots'
|
||||
# and digit_version(TORCH_VERSION) < digit_version('1.0')):
|
||||
# initialized = dist._initialized
|
||||
# else:
|
||||
if dist.is_available():
|
||||
initialized = dist.is_initialized()
|
||||
else:
|
||||
initialized = False
|
||||
if initialized:
|
||||
rank = dist.get_rank()
|
||||
world_size = dist.get_world_size()
|
||||
else:
|
||||
rank = 0
|
||||
world_size = 1
|
||||
return rank, world_size
|
||||
|
||||
|
||||
# from DETR repo
|
||||
def setup_for_distributed(is_master):
|
||||
"""
|
||||
This function disables printing when not in master process
|
||||
"""
|
||||
import builtins as __builtin__
|
||||
builtin_print = __builtin__.print
|
||||
|
||||
def print(*args, **kwargs):
|
||||
force = kwargs.pop('force', False)
|
||||
if is_master or force:
|
||||
builtin_print(*args, **kwargs)
|
||||
|
||||
__builtin__.print = print
|
||||
@@ -0,0 +1,286 @@
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
from os.path import *
|
||||
import re
|
||||
import json
|
||||
import imageio
|
||||
import os
|
||||
import math
|
||||
|
||||
os.environ["OPENCV_IO_ENABLE_OPENEXR"]="1"
|
||||
import cv2
|
||||
cv2.setNumThreads(0)
|
||||
cv2.ocl.setUseOpenCL(False)
|
||||
|
||||
TAG_CHAR = np.array([202021.25], np.float32)
|
||||
|
||||
|
||||
def readFlow(fn):
|
||||
""" Read .flo file in Middlebury format"""
|
||||
# Code adapted from:
|
||||
# http://stackoverflow.com/questions/28013200/reading-middlebury-flow-files-with-python-bytes-array-numpy
|
||||
|
||||
# WARNING: this will work on little-endian architectures (eg Intel x86) only!
|
||||
# print 'fn = %s'%(fn)
|
||||
with open(fn, 'rb') as f:
|
||||
magic = np.fromfile(f, np.float32, count=1)
|
||||
if 202021.25 != magic:
|
||||
print('Magic number incorrect. Invalid .flo file')
|
||||
return None
|
||||
else:
|
||||
w = np.fromfile(f, np.int32, count=1)
|
||||
h = np.fromfile(f, np.int32, count=1)
|
||||
# print 'Reading %d x %d flo file\n' % (w, h)
|
||||
data = np.fromfile(f, np.float32, count=2*int(w)*int(h))
|
||||
# Reshape data into 3D array (columns, rows, bands)
|
||||
# The reshape here is for visualization, the original code is (w,h,2)
|
||||
return np.resize(data, (int(h), int(w), 2))
|
||||
|
||||
|
||||
def readPFM(file):
|
||||
file = open(file, 'rb')
|
||||
|
||||
color = None
|
||||
width = None
|
||||
height = None
|
||||
scale = None
|
||||
endian = None
|
||||
|
||||
header = file.readline().rstrip()
|
||||
if header == b'PF':
|
||||
color = True
|
||||
elif header == b'Pf':
|
||||
color = False
|
||||
else:
|
||||
raise Exception('Not a PFM file.')
|
||||
|
||||
dim_match = re.match(rb'^(\d+)\s(\d+)\s$', file.readline())
|
||||
if dim_match:
|
||||
width, height = map(int, dim_match.groups())
|
||||
else:
|
||||
raise Exception('Malformed PFM header.')
|
||||
|
||||
scale = float(file.readline().rstrip())
|
||||
if scale < 0: # little-endian
|
||||
endian = '<'
|
||||
scale = -scale
|
||||
else:
|
||||
endian = '>' # big-endian
|
||||
|
||||
data = np.fromfile(file, endian + 'f')
|
||||
shape = (height, width, 3) if color else (height, width)
|
||||
|
||||
data = np.reshape(data, shape)
|
||||
data = np.flipud(data)
|
||||
return data
|
||||
|
||||
|
||||
def writePFM(file, array):
|
||||
import os
|
||||
assert type(file) is str and type(array) is np.ndarray and \
|
||||
os.path.splitext(file)[1] == ".pfm"
|
||||
with open(file, 'wb') as f:
|
||||
H, W = array.shape
|
||||
headers = ["Pf\n", f"{W} {H}\n", "-1\n"]
|
||||
for header in headers:
|
||||
f.write(str.encode(header))
|
||||
array = np.flip(array, axis=0).astype(np.float32)
|
||||
f.write(array.tobytes())
|
||||
|
||||
|
||||
def writeFlow(filename,uv,v=None):
|
||||
""" Write optical flow to file.
|
||||
|
||||
If v is None, uv is assumed to contain both u and v channels,
|
||||
stacked in depth.
|
||||
Original code by Deqing Sun, adapted from Daniel Scharstein.
|
||||
"""
|
||||
nBands = 2
|
||||
|
||||
if v is None:
|
||||
assert(uv.ndim == 3)
|
||||
assert(uv.shape[2] == 2)
|
||||
u = uv[:,:,0]
|
||||
v = uv[:,:,1]
|
||||
else:
|
||||
u = uv
|
||||
|
||||
assert(u.shape == v.shape)
|
||||
height,width = u.shape
|
||||
f = open(filename,'wb')
|
||||
# write the header
|
||||
f.write(TAG_CHAR)
|
||||
np.array(width).astype(np.int32).tofile(f)
|
||||
np.array(height).astype(np.int32).tofile(f)
|
||||
# arrange into matrix form
|
||||
tmp = np.zeros((height, width*nBands))
|
||||
tmp[:,np.arange(width)*2] = u
|
||||
tmp[:,np.arange(width)*2 + 1] = v
|
||||
tmp.astype(np.float32).tofile(f)
|
||||
f.close()
|
||||
|
||||
|
||||
def readFlowKITTI(filename):
|
||||
flow = cv2.imread(filename, cv2.IMREAD_ANYDEPTH|cv2.IMREAD_COLOR)
|
||||
flow = flow[:,:,::-1].astype(np.float32)
|
||||
flow, valid = flow[:, :, :2], flow[:, :, 2]
|
||||
flow = (flow - 2**15) / 64.0
|
||||
return flow, valid
|
||||
|
||||
|
||||
def readDispKITTI(filename):
|
||||
disp = cv2.imread(filename, cv2.IMREAD_ANYDEPTH) / 256.0
|
||||
valid = disp > 0.0
|
||||
return disp, valid
|
||||
|
||||
|
||||
def readDispInStereo2K(filename):
|
||||
disp = cv2.imread(filename, cv2.IMREAD_ANYDEPTH) / 100.0
|
||||
valid = disp > 0.0
|
||||
return disp, valid
|
||||
|
||||
|
||||
def readDispVKITTI2(filename):
|
||||
depth = cv2.imread(filename, cv2.IMREAD_ANYCOLOR | cv2.IMREAD_ANYDEPTH).astype(np.float32) / 100.0
|
||||
valid = depth > 0.0
|
||||
baseline = 0.532725
|
||||
focus_length = 725.0087
|
||||
disp = baseline*focus_length/(depth+1e-8)
|
||||
return disp, valid
|
||||
|
||||
|
||||
def readDispCreStereo(filename):
|
||||
disp = cv2.imread(filename, cv2.IMREAD_ANYDEPTH) / 32
|
||||
valid = disp > -1e-8
|
||||
return disp, valid
|
||||
|
||||
|
||||
# Method taken from /n/fs/raft-depth/RAFT-Stereo/datasets/SintelStereo/sdk/python/sintel_io.py
|
||||
def readDispSintelStereo(file_name):
|
||||
a = np.array(Image.open(file_name))
|
||||
d_r, d_g, d_b = np.split(a, axis=2, indices_or_sections=3)
|
||||
disp = (d_r.astype('float64') * 4 + d_g.astype('float64') / (2**6) + d_b.astype('float64') / (2**14))[..., 0]
|
||||
mask = np.array(Image.open(file_name.replace('disparities', 'occlusions')))
|
||||
valid = ((mask == 0) & (disp > -1e-8))
|
||||
return disp, valid
|
||||
|
||||
|
||||
# Method taken from https://research.nvidia.com/sites/default/files/pubs/2018-06_Falling-Things/readme_0.txt
|
||||
def readDispFallingThings(file_name):
|
||||
a = np.array(Image.open(file_name))
|
||||
with open('/'.join(file_name.split('/')[:-1] + ['_camera_settings.json']), 'r') as f:
|
||||
intrinsics = json.load(f)
|
||||
if 'left' in file_name:
|
||||
idx = 0
|
||||
else:
|
||||
idx = 1
|
||||
fx = intrinsics['camera_settings'][idx]['intrinsic_settings']['fx']
|
||||
disp = (fx * 6.0 * 100) / a.astype(np.float32)
|
||||
valid = disp > -1e-8
|
||||
return disp, valid
|
||||
|
||||
|
||||
# Method taken from https://github.com/castacks/tartanair_tools/blob/master/data_type.md
|
||||
def readDispTartanAir(file_name):
|
||||
depth = np.load(file_name)
|
||||
disp = 80.0 / depth
|
||||
valid = disp > -1e-8
|
||||
return disp, valid
|
||||
|
||||
|
||||
def readDispBooster(file_name):
|
||||
disp = np.load(file_name)
|
||||
valid = disp > 0
|
||||
return disp, valid
|
||||
|
||||
|
||||
def readDisp3DKenBurns(file_name):
|
||||
depth = cv2.imread(file_name, cv2.IMREAD_ANYCOLOR | cv2.IMREAD_ANYDEPTH)
|
||||
meta_file_name = file_name.replace('-depth', '')[:-7]+'-meta.json'
|
||||
fltFov = json.loads(open(meta_file_name, 'r').read())['fltFov']
|
||||
fltFocal = 0.5 * 512 * math.tan(math.radians(90.0) - (0.5 * math.radians(fltFov)))
|
||||
fltBaseline = 40.0
|
||||
disp = (fltFocal * fltBaseline) / depth
|
||||
valid = disp > 0
|
||||
return disp, valid
|
||||
|
||||
|
||||
def readDispMiddlebury0(file_name):
|
||||
if basename(file_name) == 'disp0GT.pfm':
|
||||
disp = readPFM(file_name).astype(np.float32)
|
||||
assert len(disp.shape) == 2
|
||||
nocc_pix = file_name.replace('disp0GT.pfm', 'mask0nocc.png')
|
||||
assert exists(nocc_pix)
|
||||
nocc_pix = imageio.imread(nocc_pix) == 255
|
||||
assert np.any(nocc_pix)
|
||||
return disp, nocc_pix
|
||||
elif basename(file_name) == 'disp1GT.pfm':
|
||||
disp = readPFM(file_name).astype(np.float32)
|
||||
assert len(disp.shape) == 2
|
||||
nocc_pix = file_name.replace('disp1GT.pfm', 'mask1nocc.png')
|
||||
assert exists(nocc_pix)
|
||||
nocc_pix = imageio.imread(nocc_pix) == 255
|
||||
assert np.any(nocc_pix)
|
||||
return disp, nocc_pix
|
||||
elif basename(file_name) == 'disp0.pfm':
|
||||
disp = readPFM(file_name).astype(np.float32)
|
||||
valid = disp < 1e3
|
||||
return disp, valid
|
||||
elif basename(file_name) == 'disp1.pfm':
|
||||
disp = readPFM(file_name).astype(np.float32)
|
||||
valid = disp < 1e3
|
||||
return disp, valid
|
||||
elif splitext(file_name)[-1] == '.png':
|
||||
disp = np.array(Image.open(file_name)).astype(np.float32)
|
||||
valid = disp > 0.0
|
||||
return disp, valid
|
||||
|
||||
|
||||
def readDispMiddlebury(file_name):
|
||||
if basename(file_name) == 'disp0GT.pfm':
|
||||
disp = readPFM(file_name).astype(np.float32)
|
||||
return disp, disp<1e3
|
||||
elif basename(file_name) == 'disp1GT.pfm':
|
||||
disp = readPFM(file_name).astype(np.float32)
|
||||
return disp, disp<1e3
|
||||
elif basename(file_name) == 'disp0.pfm':
|
||||
disp = readPFM(file_name).astype(np.float32)
|
||||
valid = disp < 1e3
|
||||
return disp, valid
|
||||
elif basename(file_name) == 'disp1.pfm':
|
||||
disp = readPFM(file_name).astype(np.float32)
|
||||
valid = disp < 1e3
|
||||
return disp, valid
|
||||
elif splitext(file_name)[-1] == '.png':
|
||||
disp = np.array(Image.open(file_name)).astype(np.float32)
|
||||
valid = disp > 0.0
|
||||
return disp, valid
|
||||
|
||||
|
||||
def writeFlowKITTI(filename, uv):
|
||||
uv = 64.0 * uv + 2**15
|
||||
valid = np.ones([uv.shape[0], uv.shape[1], 1])
|
||||
uv = np.concatenate([uv, valid], axis=-1).astype(np.uint16)
|
||||
cv2.imwrite(filename, uv[..., ::-1])
|
||||
|
||||
|
||||
def read_gen(file_name, pil=False):
|
||||
ext = splitext(file_name)[-1]
|
||||
if ext == '.png' or ext == '.jpeg' or ext == '.ppm' or ext == '.jpg':
|
||||
return Image.open(file_name)
|
||||
elif ext == '.bin' or ext == '.raw':
|
||||
return np.load(file_name)
|
||||
elif ext == '.flo':
|
||||
return readFlow(file_name).astype(np.float32)
|
||||
elif ext == '.pfm':
|
||||
flow = readPFM(file_name).astype(np.float32)
|
||||
if len(flow.shape) == 2:
|
||||
return flow
|
||||
else:
|
||||
return flow[:, :, :-1]
|
||||
elif ext == '.exr':
|
||||
disp = cv2.imread(file_name, cv2.IMREAD_ANYCOLOR | cv2.IMREAD_ANYDEPTH)
|
||||
if len(disp.shape) > 2:
|
||||
disp = disp[..., 0]
|
||||
return disp
|
||||
return []
|
||||
@@ -0,0 +1,242 @@
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import numpy as np
|
||||
from scipy import interpolate
|
||||
import glob
|
||||
import os.path as osp
|
||||
|
||||
|
||||
def get_danv2_io_size(h, w, nds, max_i_size=2688, multiple_of=14):
|
||||
"""compute the input and output sizes of danv2 network"""
|
||||
danv2_oh, danv2_ow = h//2**nds, w//2**nds
|
||||
danv2_io_factor = 3.5 # more precise, 14/8=3.5
|
||||
ih, iw = danv2_io_factor*danv2_oh, danv2_io_factor*danv2_ow
|
||||
ih = int(np.ceil(ih / multiple_of) * multiple_of)
|
||||
iw = int(np.ceil(iw / multiple_of) * multiple_of)
|
||||
|
||||
max_i_size = int(np.floor(max_i_size / multiple_of) * multiple_of)
|
||||
|
||||
if ih <= max_i_size and iw <= max_i_size:
|
||||
danv2_ih, danv2_iw = ih, iw
|
||||
else:
|
||||
factor_h = max_i_size/ih
|
||||
factor_w = max_i_size/iw
|
||||
|
||||
if factor_w > factor_h:
|
||||
danv2_ih = max_i_size
|
||||
danv2_iw = int(np.ceil(factor_h * iw / multiple_of) * multiple_of)
|
||||
else:
|
||||
danv2_iw = max_i_size
|
||||
danv2_ih = int(np.ceil(factor_w * ih / multiple_of) * multiple_of)
|
||||
|
||||
return danv2_ih, danv2_iw, danv2_oh, danv2_ow
|
||||
|
||||
|
||||
class InputPadder:
|
||||
""" Pads images such that dimensions are divisible by 8 """
|
||||
def __init__(self, dims, mode='sintel', divis_by=8):
|
||||
self.ht, self.wd = dims[-2:]
|
||||
pad_ht = (((self.ht // divis_by) + 1) * divis_by - self.ht) % divis_by
|
||||
pad_wd = (((self.wd // divis_by) + 1) * divis_by - self.wd) % divis_by
|
||||
if mode == 'sintel':
|
||||
self._pad = [pad_wd//2, pad_wd - pad_wd//2, pad_ht//2, pad_ht - pad_ht//2]
|
||||
else:
|
||||
self._pad = [pad_wd//2, pad_wd - pad_wd//2, 0, pad_ht]
|
||||
|
||||
def pad(self, *inputs):
|
||||
assert all((x.ndim == 4) for x in inputs)
|
||||
return [F.pad(x, self._pad, mode='replicate') for x in inputs]
|
||||
|
||||
def unpad(self, x):
|
||||
assert x.ndim == 4
|
||||
ht, wd = x.shape[-2:]
|
||||
c = [self._pad[2], ht-self._pad[3], self._pad[0], wd-self._pad[1]]
|
||||
return x[..., c[0]:c[1], c[2]:c[3]]
|
||||
|
||||
|
||||
def forward_interpolate(flow):
|
||||
flow = flow.detach().cpu().numpy()
|
||||
dx, dy = flow[0], flow[1]
|
||||
|
||||
ht, wd = dx.shape
|
||||
x0, y0 = np.meshgrid(np.arange(wd), np.arange(ht))
|
||||
|
||||
x1 = x0 + dx
|
||||
y1 = y0 + dy
|
||||
|
||||
x1 = x1.reshape(-1)
|
||||
y1 = y1.reshape(-1)
|
||||
dx = dx.reshape(-1)
|
||||
dy = dy.reshape(-1)
|
||||
|
||||
valid = (x1 > 0) & (x1 < wd) & (y1 > 0) & (y1 < ht)
|
||||
x1 = x1[valid]
|
||||
y1 = y1[valid]
|
||||
dx = dx[valid]
|
||||
dy = dy[valid]
|
||||
|
||||
flow_x = interpolate.griddata(
|
||||
(x1, y1), dx, (x0, y0), method='nearest', fill_value=0)
|
||||
|
||||
flow_y = interpolate.griddata(
|
||||
(x1, y1), dy, (x0, y0), method='nearest', fill_value=0)
|
||||
|
||||
flow = np.stack([flow_x, flow_y], axis=0)
|
||||
return torch.from_numpy(flow).float()
|
||||
|
||||
|
||||
def bilinear_sampler(img, coords, mode='bilinear', mask=False):
|
||||
""" Wrapper for grid_sample, uses pixel coordinates """
|
||||
H, W = img.shape[-2:]
|
||||
xgrid, ygrid = coords.split([1, 1], dim=-1)
|
||||
xgrid = 2*xgrid/(W-1) - 1
|
||||
if H > 1:
|
||||
ygrid = 2*ygrid/(H-1) - 1
|
||||
|
||||
grid = torch.cat([xgrid, ygrid], dim=-1)
|
||||
img = F.grid_sample(img, grid, align_corners=True)
|
||||
# img = bilinear_grid_sample(img, grid, align_corners=True)
|
||||
|
||||
if mask:
|
||||
mask = (xgrid > -1) & (ygrid > -1) & (xgrid < 1) & (ygrid < 1)
|
||||
return img, mask.float()
|
||||
|
||||
return img
|
||||
|
||||
|
||||
def coords_grid(batch, ht, wd):
|
||||
coords = torch.meshgrid(torch.arange(ht), torch.arange(wd))
|
||||
coords = torch.stack(coords[::-1], dim=0).float()
|
||||
return coords[None].repeat(batch, 1, 1, 1)
|
||||
|
||||
|
||||
def upflow(flow, factor=8, mode='bilinear', sacle=True):
|
||||
new_size = (factor * flow.shape[2], factor * flow.shape[3])
|
||||
if sacle:
|
||||
return factor * F.interpolate(flow, size=new_size, mode=mode, align_corners=True)
|
||||
else:
|
||||
return F.interpolate(flow, size=new_size, mode=mode, align_corners=True)
|
||||
|
||||
|
||||
def gauss_blur(input, N=5, std=1):
|
||||
B, D, H, W = input.shape
|
||||
x, y = torch.meshgrid(torch.arange(N).float() - N//2, torch.arange(N).float() - N//2)
|
||||
unnormalized_gaussian = torch.exp(-(x.pow(2) + y.pow(2)) / (2 * std ** 2))
|
||||
weights = unnormalized_gaussian / unnormalized_gaussian.sum().clamp(min=1e-4)
|
||||
weights = weights.view(1, 1, N, N).to(input)
|
||||
output = F.conv2d(input.reshape(B*D, 1, H, W), weights, padding=N//2)
|
||||
return output.view(B, D, H, W)
|
||||
|
||||
|
||||
# Ref: https://zenn.dev/pinto0309/scraps/7d4032067d0160
|
||||
def bilinear_grid_sample(im, grid, align_corners=False):
|
||||
"""Given an input and a flow-field grid, computes the output using input
|
||||
values and pixel locations from grid. Supported only bilinear interpolation
|
||||
method to sample the input pixels.
|
||||
|
||||
Args:
|
||||
im (torch.Tensor): Input feature map, shape (N, C, H, W)
|
||||
grid (torch.Tensor): Point coordinates, shape (N, Hg, Wg, 2)
|
||||
align_corners {bool}: If set to True, the extrema (-1 and 1) are
|
||||
considered as referring to the center points of the input’s
|
||||
corner pixels. If set to False, they are instead considered as
|
||||
referring to the corner points of the input’s corner pixels,
|
||||
making the sampling more resolution agnostic.
|
||||
|
||||
Returns:
|
||||
torch.Tensor: A tensor with sampled points, shape (N, C, Hg, Wg)
|
||||
"""
|
||||
n, c, h, w = im.shape
|
||||
gn, gh, gw, _ = grid.shape
|
||||
assert n == gn
|
||||
|
||||
x = grid[:, :, :, 0]
|
||||
y = grid[:, :, :, 1]
|
||||
|
||||
if align_corners:
|
||||
x = ((x + 1) / 2) * (w - 1)
|
||||
y = ((y + 1) / 2) * (h - 1)
|
||||
else:
|
||||
x = ((x + 1) * w - 1) / 2
|
||||
y = ((y + 1) * h - 1) / 2
|
||||
|
||||
x = x.view(n, -1)
|
||||
y = y.view(n, -1)
|
||||
|
||||
x0 = torch.floor(x).long()
|
||||
y0 = torch.floor(y).long()
|
||||
x1 = x0 + 1
|
||||
y1 = y0 + 1
|
||||
|
||||
wa = ((x1 - x) * (y1 - y)).unsqueeze(1)
|
||||
wb = ((x1 - x) * (y - y0)).unsqueeze(1)
|
||||
wc = ((x - x0) * (y1 - y)).unsqueeze(1)
|
||||
wd = ((x - x0) * (y - y0)).unsqueeze(1)
|
||||
|
||||
# Apply default for grid_sample function zero padding
|
||||
im_padded = torch.nn.functional.pad(im, pad=[1, 1, 1, 1], mode='constant', value=0)
|
||||
padded_h = h + 2
|
||||
padded_w = w + 2
|
||||
# save points positions after padding
|
||||
x0, x1, y0, y1 = x0 + 1, x1 + 1, y0 + 1, y1 + 1
|
||||
|
||||
# Clip coordinates to padded image size
|
||||
x0 = torch.where(x0 < 0, torch.tensor(0, device=im.device), x0)
|
||||
x0 = torch.where(x0 > padded_w - 1, torch.tensor(padded_w - 1, device=im.device), x0)
|
||||
x1 = torch.where(x1 < 0, torch.tensor(0, device=im.device), x1)
|
||||
x1 = torch.where(x1 > padded_w - 1, torch.tensor(padded_w - 1, device=im.device), x1)
|
||||
y0 = torch.where(y0 < 0, torch.tensor(0, device=im.device), y0)
|
||||
y0 = torch.where(y0 > padded_h - 1, torch.tensor(padded_h - 1, device=im.device), y0)
|
||||
y1 = torch.where(y1 < 0, torch.tensor(0, device=im.device), y1)
|
||||
y1 = torch.where(y1 > padded_h - 1, torch.tensor(padded_h - 1, device=im.device), y1)
|
||||
|
||||
im_padded = im_padded.view(n, c, -1)
|
||||
|
||||
x0_y0 = (x0 + y0 * padded_w).unsqueeze(1).expand(-1, c, -1)
|
||||
x0_y1 = (x0 + y1 * padded_w).unsqueeze(1).expand(-1, c, -1)
|
||||
x1_y0 = (x1 + y0 * padded_w).unsqueeze(1).expand(-1, c, -1)
|
||||
x1_y1 = (x1 + y1 * padded_w).unsqueeze(1).expand(-1, c, -1)
|
||||
|
||||
Ia = torch.gather(im_padded, 2, x0_y0)
|
||||
Ib = torch.gather(im_padded, 2, x0_y1)
|
||||
Ic = torch.gather(im_padded, 2, x1_y0)
|
||||
Id = torch.gather(im_padded, 2, x1_y1)
|
||||
|
||||
return (Ia * wa + Ib * wb + Ic * wc + Id * wd).reshape(n, c, gh, gw)
|
||||
|
||||
|
||||
def read_kitti_calib_file(path):
|
||||
"""Read KITTI calibration file
|
||||
(from https://github.com/hunse/kitti)
|
||||
"""
|
||||
float_chars = set("0123456789.e+- ")
|
||||
data = {}
|
||||
with open(path, 'r') as f:
|
||||
for line in f.readlines():
|
||||
key, value = line.split(':', 1)
|
||||
value = value.strip()
|
||||
data[key] = value
|
||||
if float_chars.issuperset(value):
|
||||
# try to cast to float array
|
||||
try:
|
||||
data[key] = np.array(list(map(float, value.split(' '))))
|
||||
except ValueError:
|
||||
# casting error: data[key] already eq. value, so pass
|
||||
pass
|
||||
|
||||
return data
|
||||
|
||||
|
||||
# from https://github.com/ozendelait/rvc_devkit/blob/master/stereo/stereo_devkit.py
|
||||
def ReadMiddlebury2014CalibFile(path):
|
||||
result = dict()
|
||||
with open(path, 'rb') as calib_file:
|
||||
for line in calib_file.readlines():
|
||||
line = line.decode('UTF-8').rstrip('\n')
|
||||
if len(line) == 0:
|
||||
continue
|
||||
eq_pos = line.find('=')
|
||||
if eq_pos < 0:
|
||||
raise Exception('Cannot parse Middlebury 2014 calib file: ' + path)
|
||||
result[line[:eq_pos]] = line[eq_pos + 1:]
|
||||
return result
|
||||
@@ -0,0 +1,89 @@
|
||||
import sys
|
||||
sys.path.append('core')
|
||||
|
||||
import argparse
|
||||
import glob
|
||||
import numpy as np
|
||||
import torch
|
||||
from tqdm import tqdm
|
||||
from pathlib import Path
|
||||
from core.defom_stereo import DEFOMStereo
|
||||
from utils.utils import InputPadder
|
||||
from PIL import Image
|
||||
from matplotlib import pyplot as plt
|
||||
|
||||
|
||||
DEVICE = 'cuda'
|
||||
|
||||
def load_image(imfile):
|
||||
img = np.array(Image.open(imfile)).astype(np.uint8)
|
||||
img = torch.from_numpy(img).permute(2, 0, 1).float()
|
||||
return img[None].to(DEVICE)
|
||||
|
||||
def demo(args):
|
||||
model = DEFOMStereo(args)
|
||||
checkpoint = torch.load(args.restore_ckpt, map_location='cuda')
|
||||
if 'model' in checkpoint:
|
||||
model.load_state_dict(checkpoint['model'])
|
||||
else:
|
||||
model.load_state_dict(checkpoint)
|
||||
|
||||
model.to(DEVICE)
|
||||
model.eval()
|
||||
|
||||
output_directory = Path(args.output_directory)
|
||||
output_directory.mkdir(exist_ok=True)
|
||||
|
||||
with torch.no_grad():
|
||||
left_images = sorted(glob.glob(args.left_imgs, recursive=True))
|
||||
right_images = sorted(glob.glob(args.right_imgs, recursive=True))
|
||||
print(f"Found {len(left_images)} images. Saving files to {output_directory}/")
|
||||
|
||||
for (imfile1, imfile2) in tqdm(list(zip(left_images, right_images))):
|
||||
image1 = load_image(imfile1)
|
||||
image2 = load_image(imfile2)
|
||||
|
||||
padder = InputPadder(image1.shape, divis_by=32)
|
||||
image1, image2 = padder.pad(image1, image2)
|
||||
|
||||
with torch.no_grad():
|
||||
disp_pr = model(image1, image2, iters=args.valid_iters, scale_iters=args.scale_iters, test_mode=True)
|
||||
disp_pr = padder.unpad(disp_pr).cpu().squeeze().numpy()
|
||||
|
||||
file_stem = imfile1.split('/')[-1].split('_')[0]+'_'+args.restore_ckpt.split('/')[-1][:-4]
|
||||
if args.save_numpy:
|
||||
np.save(output_directory / f"{file_stem}.npy", disp_pr)
|
||||
plt.imsave(output_directory / f"{file_stem}.png", disp_pr, cmap='jet')
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--restore_ckpt', help="restore checkpoint", required=True)
|
||||
parser.add_argument('--save_numpy', action='store_true', help='save output as numpy arrays')
|
||||
parser.add_argument('-l', '--left_imgs', help="path to all first (left) frames", default="demo/*_left.png")
|
||||
parser.add_argument('-r', '--right_imgs', help="path to all second (right) frames", default="demo/*_right.png")
|
||||
parser.add_argument('--output_directory', help="directory to save output", default="demo")
|
||||
parser.add_argument('--mixed_precision', action='store_true', help='use mixed precision')
|
||||
parser.add_argument('--valid_iters', type=int, default=32, help='number of flow-field updates during forward pass')
|
||||
parser.add_argument('--scale_iters', type=int, default=8, help="number of scaling updates to the disparity field in each forward pass.")
|
||||
|
||||
# Architecture choices
|
||||
parser.add_argument('--dinov2_encoder', type=str, default='vitl', choices=['vits', 'vitb', 'vitl', 'vitg'])
|
||||
parser.add_argument('--idepth_scale', type=float, default=0.5, help="the scale of inverse depth to initialize disparity")
|
||||
parser.add_argument('--hidden_dims', nargs='+', type=int, default=[128]*3, help="hidden state and context dimensions")
|
||||
parser.add_argument('--corr_implementation', choices=["reg", "alt", "reg_cuda", "alt_cuda"], default="reg", help="correlation volume implementation")
|
||||
parser.add_argument('--shared_backbone', action='store_true', help="use a single backbone for the context and feature encoders")
|
||||
parser.add_argument('--corr_levels', type=int, default=2, help="number of levels in the correlation pyramid")
|
||||
parser.add_argument('--corr_radius', type=int, default=4, help="width of the correlation pyramid")
|
||||
parser.add_argument('--scale_list', type=float, nargs='+', default=[0.125, 0.25, 0.5, 0.75, 1.0, 1.25, 1.5, 2.0],
|
||||
help='the list of scaling factors of disparity')
|
||||
parser.add_argument('--scale_corr_radius', type=int, default=2,
|
||||
help="width of the correlation pyramid for scaled disparity")
|
||||
|
||||
parser.add_argument('--n_downsample', type=int, default=2, choices=[2, 3], help="resolution of the disparity field (1/2^K)")
|
||||
parser.add_argument('--context_norm', type=str, default="batch", choices=['group', 'batch', 'instance', 'none'], help="normalization of context encoder")
|
||||
parser.add_argument('--n_gru_layers', type=int, default=3, help="number of hidden GRU levels")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
demo(args)
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 1.1 MiB |
Binary file not shown.
|
After Width: | Height: | Size: 1.2 MiB |
Binary file not shown.
|
After Width: | Height: | Size: 975 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 994 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 1.1 MiB |
Binary file not shown.
|
After Width: | Height: | Size: 1.1 MiB |
Binary file not shown.
@@ -0,0 +1,415 @@
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
||||
#
|
||||
# This source code is licensed under the Apache License, Version 2.0
|
||||
# found in the LICENSE file in the root directory of this source tree.
|
||||
|
||||
# References:
|
||||
# https://github.com/facebookresearch/dino/blob/main/vision_transformer.py
|
||||
# https://github.com/rwightman/pytorch-image-models/tree/master/timm/models/vision_transformer.py
|
||||
|
||||
from functools import partial
|
||||
import math
|
||||
import logging
|
||||
from typing import Sequence, Tuple, Union, Callable
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.utils.checkpoint
|
||||
from torch.nn.init import trunc_normal_
|
||||
|
||||
from .dinov2_layers import Mlp, PatchEmbed, SwiGLUFFNFused, MemEffAttention, NestedTensorBlock as Block
|
||||
|
||||
|
||||
logger = logging.getLogger("dinov2")
|
||||
|
||||
|
||||
def named_apply(fn: Callable, module: nn.Module, name="", depth_first=True, include_root=False) -> nn.Module:
|
||||
if not depth_first and include_root:
|
||||
fn(module=module, name=name)
|
||||
for child_name, child_module in module.named_children():
|
||||
child_name = ".".join((name, child_name)) if name else child_name
|
||||
named_apply(fn=fn, module=child_module, name=child_name, depth_first=depth_first, include_root=True)
|
||||
if depth_first and include_root:
|
||||
fn(module=module, name=name)
|
||||
return module
|
||||
|
||||
|
||||
class BlockChunk(nn.ModuleList):
|
||||
def forward(self, x):
|
||||
for b in self:
|
||||
x = b(x)
|
||||
return x
|
||||
|
||||
|
||||
class DinoVisionTransformer(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
img_size=224,
|
||||
patch_size=16,
|
||||
in_chans=3,
|
||||
embed_dim=768,
|
||||
depth=12,
|
||||
num_heads=12,
|
||||
mlp_ratio=4.0,
|
||||
qkv_bias=True,
|
||||
ffn_bias=True,
|
||||
proj_bias=True,
|
||||
drop_path_rate=0.0,
|
||||
drop_path_uniform=False,
|
||||
init_values=None, # for layerscale: None or 0 => no layerscale
|
||||
embed_layer=PatchEmbed,
|
||||
act_layer=nn.GELU,
|
||||
block_fn=Block,
|
||||
ffn_layer="mlp",
|
||||
block_chunks=1,
|
||||
num_register_tokens=0,
|
||||
interpolate_antialias=False,
|
||||
interpolate_offset=0.1,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
img_size (int, tuple): input image size
|
||||
patch_size (int, tuple): patch size
|
||||
in_chans (int): number of input channels
|
||||
embed_dim (int): embedding dimension
|
||||
depth (int): depth of transformer
|
||||
num_heads (int): number of attention heads
|
||||
mlp_ratio (int): ratio of mlp hidden dim to embedding dim
|
||||
qkv_bias (bool): enable bias for qkv if True
|
||||
proj_bias (bool): enable bias for proj in attn if True
|
||||
ffn_bias (bool): enable bias for ffn if True
|
||||
drop_path_rate (float): stochastic depth rate
|
||||
drop_path_uniform (bool): apply uniform drop rate across blocks
|
||||
weight_init (str): weight init scheme
|
||||
init_values (float): layer-scale init values
|
||||
embed_layer (nn.Module): patch embedding layer
|
||||
act_layer (nn.Module): MLP activation layer
|
||||
block_fn (nn.Module): transformer block class
|
||||
ffn_layer (str): "mlp", "swiglu", "swiglufused" or "identity"
|
||||
block_chunks: (int) split block sequence into block_chunks units for FSDP wrap
|
||||
num_register_tokens: (int) number of extra cls tokens (so-called "registers")
|
||||
interpolate_antialias: (str) flag to apply anti-aliasing when interpolating positional embeddings
|
||||
interpolate_offset: (float) work-around offset to apply when interpolating positional embeddings
|
||||
"""
|
||||
super().__init__()
|
||||
norm_layer = partial(nn.LayerNorm, eps=1e-6)
|
||||
|
||||
self.num_features = self.embed_dim = embed_dim # num_features for consistency with other models
|
||||
self.num_tokens = 1
|
||||
self.n_blocks = depth
|
||||
self.num_heads = num_heads
|
||||
self.patch_size = patch_size
|
||||
self.num_register_tokens = num_register_tokens
|
||||
self.interpolate_antialias = interpolate_antialias
|
||||
self.interpolate_offset = interpolate_offset
|
||||
|
||||
self.patch_embed = embed_layer(img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim)
|
||||
num_patches = self.patch_embed.num_patches
|
||||
|
||||
self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
|
||||
self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + self.num_tokens, embed_dim))
|
||||
assert num_register_tokens >= 0
|
||||
self.register_tokens = (
|
||||
nn.Parameter(torch.zeros(1, num_register_tokens, embed_dim)) if num_register_tokens else None
|
||||
)
|
||||
|
||||
if drop_path_uniform is True:
|
||||
dpr = [drop_path_rate] * depth
|
||||
else:
|
||||
dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)] # stochastic depth decay rule
|
||||
|
||||
if ffn_layer == "mlp":
|
||||
logger.info("using MLP layer as FFN")
|
||||
ffn_layer = Mlp
|
||||
elif ffn_layer == "swiglufused" or ffn_layer == "swiglu":
|
||||
logger.info("using SwiGLU layer as FFN")
|
||||
ffn_layer = SwiGLUFFNFused
|
||||
elif ffn_layer == "identity":
|
||||
logger.info("using Identity layer as FFN")
|
||||
|
||||
def f(*args, **kwargs):
|
||||
return nn.Identity()
|
||||
|
||||
ffn_layer = f
|
||||
else:
|
||||
raise NotImplementedError
|
||||
|
||||
blocks_list = [
|
||||
block_fn(
|
||||
dim=embed_dim,
|
||||
num_heads=num_heads,
|
||||
mlp_ratio=mlp_ratio,
|
||||
qkv_bias=qkv_bias,
|
||||
proj_bias=proj_bias,
|
||||
ffn_bias=ffn_bias,
|
||||
drop_path=dpr[i],
|
||||
norm_layer=norm_layer,
|
||||
act_layer=act_layer,
|
||||
ffn_layer=ffn_layer,
|
||||
init_values=init_values,
|
||||
)
|
||||
for i in range(depth)
|
||||
]
|
||||
if block_chunks > 0:
|
||||
self.chunked_blocks = True
|
||||
chunked_blocks = []
|
||||
chunksize = depth // block_chunks
|
||||
for i in range(0, depth, chunksize):
|
||||
# this is to keep the block index consistent if we chunk the block list
|
||||
chunked_blocks.append([nn.Identity()] * i + blocks_list[i : i + chunksize])
|
||||
self.blocks = nn.ModuleList([BlockChunk(p) for p in chunked_blocks])
|
||||
else:
|
||||
self.chunked_blocks = False
|
||||
self.blocks = nn.ModuleList(blocks_list)
|
||||
|
||||
self.norm = norm_layer(embed_dim)
|
||||
self.head = nn.Identity()
|
||||
|
||||
self.mask_token = nn.Parameter(torch.zeros(1, embed_dim))
|
||||
|
||||
self.init_weights()
|
||||
|
||||
def init_weights(self):
|
||||
trunc_normal_(self.pos_embed, std=0.02)
|
||||
nn.init.normal_(self.cls_token, std=1e-6)
|
||||
if self.register_tokens is not None:
|
||||
nn.init.normal_(self.register_tokens, std=1e-6)
|
||||
named_apply(init_weights_vit_timm, self)
|
||||
|
||||
def interpolate_pos_encoding(self, x, w, h):
|
||||
previous_dtype = x.dtype
|
||||
npatch = x.shape[1] - 1
|
||||
N = self.pos_embed.shape[1] - 1
|
||||
if npatch == N and w == h:
|
||||
return self.pos_embed
|
||||
pos_embed = self.pos_embed.float()
|
||||
class_pos_embed = pos_embed[:, 0]
|
||||
patch_pos_embed = pos_embed[:, 1:]
|
||||
dim = x.shape[-1]
|
||||
w0 = w // self.patch_size
|
||||
h0 = h // self.patch_size
|
||||
# we add a small number to avoid floating point error in the interpolation
|
||||
# see discussion at https://github.com/facebookresearch/dino/issues/8
|
||||
# DINOv2 with register modify the interpolate_offset from 0.1 to 0.0
|
||||
w0, h0 = w0 + self.interpolate_offset, h0 + self.interpolate_offset
|
||||
# w0, h0 = w0 + 0.1, h0 + 0.1
|
||||
|
||||
sqrt_N = math.sqrt(N)
|
||||
sx, sy = float(w0) / sqrt_N, float(h0) / sqrt_N
|
||||
patch_pos_embed = nn.functional.interpolate(
|
||||
patch_pos_embed.reshape(1, int(sqrt_N), int(sqrt_N), dim).permute(0, 3, 1, 2),
|
||||
scale_factor=(sx, sy),
|
||||
# (int(w0), int(h0)), # to solve the upsampling shape issue
|
||||
mode="bicubic",
|
||||
antialias=self.interpolate_antialias
|
||||
)
|
||||
|
||||
assert int(w0) == patch_pos_embed.shape[-2]
|
||||
assert int(h0) == patch_pos_embed.shape[-1]
|
||||
patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
|
||||
return torch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1).to(previous_dtype)
|
||||
|
||||
def prepare_tokens_with_masks(self, x, masks=None):
|
||||
B, nc, w, h = x.shape
|
||||
x = self.patch_embed(x)
|
||||
if masks is not None:
|
||||
x = torch.where(masks.unsqueeze(-1), self.mask_token.to(x.dtype).unsqueeze(0), x)
|
||||
|
||||
x = torch.cat((self.cls_token.expand(x.shape[0], -1, -1), x), dim=1)
|
||||
x = x + self.interpolate_pos_encoding(x, w, h)
|
||||
|
||||
if self.register_tokens is not None:
|
||||
x = torch.cat(
|
||||
(
|
||||
x[:, :1],
|
||||
self.register_tokens.expand(x.shape[0], -1, -1),
|
||||
x[:, 1:],
|
||||
),
|
||||
dim=1,
|
||||
)
|
||||
|
||||
return x
|
||||
|
||||
def forward_features_list(self, x_list, masks_list):
|
||||
x = [self.prepare_tokens_with_masks(x, masks) for x, masks in zip(x_list, masks_list)]
|
||||
for blk in self.blocks:
|
||||
x = blk(x)
|
||||
|
||||
all_x = x
|
||||
output = []
|
||||
for x, masks in zip(all_x, masks_list):
|
||||
x_norm = self.norm(x)
|
||||
output.append(
|
||||
{
|
||||
"x_norm_clstoken": x_norm[:, 0],
|
||||
"x_norm_regtokens": x_norm[:, 1 : self.num_register_tokens + 1],
|
||||
"x_norm_patchtokens": x_norm[:, self.num_register_tokens + 1 :],
|
||||
"x_prenorm": x,
|
||||
"masks": masks,
|
||||
}
|
||||
)
|
||||
return output
|
||||
|
||||
def forward_features(self, x, masks=None):
|
||||
if isinstance(x, list):
|
||||
return self.forward_features_list(x, masks)
|
||||
|
||||
x = self.prepare_tokens_with_masks(x, masks)
|
||||
|
||||
for blk in self.blocks:
|
||||
x = blk(x)
|
||||
|
||||
x_norm = self.norm(x)
|
||||
return {
|
||||
"x_norm_clstoken": x_norm[:, 0],
|
||||
"x_norm_regtokens": x_norm[:, 1 : self.num_register_tokens + 1],
|
||||
"x_norm_patchtokens": x_norm[:, self.num_register_tokens + 1 :],
|
||||
"x_prenorm": x,
|
||||
"masks": masks,
|
||||
}
|
||||
|
||||
def _get_intermediate_layers_not_chunked(self, x, n=1):
|
||||
x = self.prepare_tokens_with_masks(x)
|
||||
# If n is an int, take the n last blocks. If it's a list, take them
|
||||
output, total_block_len = [], len(self.blocks)
|
||||
blocks_to_take = range(total_block_len - n, total_block_len) if isinstance(n, int) else n
|
||||
for i, blk in enumerate(self.blocks):
|
||||
x = blk(x)
|
||||
if i in blocks_to_take:
|
||||
output.append(x)
|
||||
assert len(output) == len(blocks_to_take), f"only {len(output)} / {len(blocks_to_take)} blocks found"
|
||||
return output
|
||||
|
||||
def _get_intermediate_layers_chunked(self, x, n=1):
|
||||
x = self.prepare_tokens_with_masks(x)
|
||||
output, i, total_block_len = [], 0, len(self.blocks[-1])
|
||||
# If n is an int, take the n last blocks. If it's a list, take them
|
||||
blocks_to_take = range(total_block_len - n, total_block_len) if isinstance(n, int) else n
|
||||
for block_chunk in self.blocks:
|
||||
for blk in block_chunk[i:]: # Passing the nn.Identity()
|
||||
x = blk(x)
|
||||
if i in blocks_to_take:
|
||||
output.append(x)
|
||||
i += 1
|
||||
assert len(output) == len(blocks_to_take), f"only {len(output)} / {len(blocks_to_take)} blocks found"
|
||||
return output
|
||||
|
||||
def get_intermediate_layers(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
n: Union[int, Sequence] = 1, # Layers or n last layers to take
|
||||
reshape: bool = False,
|
||||
return_class_token: bool = False,
|
||||
norm=True
|
||||
) -> Tuple[Union[torch.Tensor, Tuple[torch.Tensor]]]:
|
||||
if self.chunked_blocks:
|
||||
outputs = self._get_intermediate_layers_chunked(x, n)
|
||||
else:
|
||||
outputs = self._get_intermediate_layers_not_chunked(x, n)
|
||||
if norm:
|
||||
outputs = [self.norm(out) for out in outputs]
|
||||
class_tokens = [out[:, 0] for out in outputs]
|
||||
outputs = [out[:, 1 + self.num_register_tokens:] for out in outputs]
|
||||
if reshape:
|
||||
B, _, w, h = x.shape
|
||||
outputs = [
|
||||
out.reshape(B, w // self.patch_size, h // self.patch_size, -1).permute(0, 3, 1, 2).contiguous()
|
||||
for out in outputs
|
||||
]
|
||||
if return_class_token:
|
||||
return tuple(zip(outputs, class_tokens))
|
||||
return tuple(outputs)
|
||||
|
||||
def forward(self, *args, is_training=False, **kwargs):
|
||||
ret = self.forward_features(*args, **kwargs)
|
||||
if is_training:
|
||||
return ret
|
||||
else:
|
||||
return self.head(ret["x_norm_clstoken"])
|
||||
|
||||
|
||||
def init_weights_vit_timm(module: nn.Module, name: str = ""):
|
||||
"""ViT weight initialization, original timm impl (for reproducibility)"""
|
||||
if isinstance(module, nn.Linear):
|
||||
trunc_normal_(module.weight, std=0.02)
|
||||
if module.bias is not None:
|
||||
nn.init.zeros_(module.bias)
|
||||
|
||||
|
||||
def vit_small(patch_size=16, num_register_tokens=0, **kwargs):
|
||||
model = DinoVisionTransformer(
|
||||
patch_size=patch_size,
|
||||
embed_dim=384,
|
||||
depth=12,
|
||||
num_heads=6,
|
||||
mlp_ratio=4,
|
||||
block_fn=partial(Block, attn_class=MemEffAttention),
|
||||
num_register_tokens=num_register_tokens,
|
||||
**kwargs,
|
||||
)
|
||||
return model
|
||||
|
||||
|
||||
def vit_base(patch_size=16, num_register_tokens=0, **kwargs):
|
||||
model = DinoVisionTransformer(
|
||||
patch_size=patch_size,
|
||||
embed_dim=768,
|
||||
depth=12,
|
||||
num_heads=12,
|
||||
mlp_ratio=4,
|
||||
block_fn=partial(Block, attn_class=MemEffAttention),
|
||||
num_register_tokens=num_register_tokens,
|
||||
**kwargs,
|
||||
)
|
||||
return model
|
||||
|
||||
|
||||
def vit_large(patch_size=16, num_register_tokens=0, **kwargs):
|
||||
model = DinoVisionTransformer(
|
||||
patch_size=patch_size,
|
||||
embed_dim=1024,
|
||||
depth=24,
|
||||
num_heads=16,
|
||||
mlp_ratio=4,
|
||||
block_fn=partial(Block, attn_class=MemEffAttention),
|
||||
num_register_tokens=num_register_tokens,
|
||||
**kwargs,
|
||||
)
|
||||
return model
|
||||
|
||||
|
||||
def vit_giant2(patch_size=16, num_register_tokens=0, **kwargs):
|
||||
"""
|
||||
Close to ViT-giant, with embed-dim 1536 and 24 heads => embed-dim per head 64
|
||||
"""
|
||||
model = DinoVisionTransformer(
|
||||
patch_size=patch_size,
|
||||
embed_dim=1536,
|
||||
depth=40,
|
||||
num_heads=24,
|
||||
mlp_ratio=4,
|
||||
block_fn=partial(Block, attn_class=MemEffAttention),
|
||||
num_register_tokens=num_register_tokens,
|
||||
**kwargs,
|
||||
)
|
||||
return model
|
||||
|
||||
|
||||
def DINOv2(model_name):
|
||||
model_zoo = {
|
||||
"vits": vit_small,
|
||||
"vitb": vit_base,
|
||||
"vitl": vit_large,
|
||||
"vitg": vit_giant2
|
||||
}
|
||||
|
||||
return model_zoo[model_name](
|
||||
img_size=518,
|
||||
patch_size=14,
|
||||
init_values=1.0,
|
||||
ffn_layer="mlp" if model_name != "vitg" else "swiglufused",
|
||||
block_chunks=0,
|
||||
num_register_tokens=0,
|
||||
interpolate_antialias=False,
|
||||
interpolate_offset=0.1
|
||||
)
|
||||
@@ -0,0 +1,11 @@
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
||||
# All rights reserved.
|
||||
#
|
||||
# This source code is licensed under the license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
from .mlp import Mlp
|
||||
from .patch_embed import PatchEmbed
|
||||
from .swiglu_ffn import SwiGLUFFN, SwiGLUFFNFused
|
||||
from .block import NestedTensorBlock
|
||||
from .attention import MemEffAttention
|
||||
@@ -0,0 +1,82 @@
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
||||
# All rights reserved.
|
||||
#
|
||||
# This source code is licensed under the license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
# References:
|
||||
# https://github.com/facebookresearch/dino/blob/master/vision_transformer.py
|
||||
# https://github.com/rwightman/pytorch-image-models/tree/master/timm/models/vision_transformer.py
|
||||
|
||||
import logging
|
||||
|
||||
from torch import Tensor
|
||||
from torch import nn
|
||||
|
||||
|
||||
logger = logging.getLogger("dinov2")
|
||||
|
||||
|
||||
try:
|
||||
from xformers.ops import memory_efficient_attention, unbind, fmha
|
||||
|
||||
XFORMERS_AVAILABLE = True
|
||||
except ImportError:
|
||||
logger.warning("xFormers not available")
|
||||
XFORMERS_AVAILABLE = False
|
||||
|
||||
|
||||
class Attention(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
dim: int,
|
||||
num_heads: int = 8,
|
||||
qkv_bias: bool = False,
|
||||
proj_bias: bool = True,
|
||||
attn_drop: float = 0.0,
|
||||
proj_drop: float = 0.0,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.num_heads = num_heads
|
||||
head_dim = dim // num_heads
|
||||
self.scale = head_dim**-0.5
|
||||
|
||||
self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
|
||||
self.attn_drop = nn.Dropout(attn_drop)
|
||||
self.proj = nn.Linear(dim, dim, bias=proj_bias)
|
||||
self.proj_drop = nn.Dropout(proj_drop)
|
||||
|
||||
def forward(self, x: Tensor) -> Tensor:
|
||||
B, N, C = x.shape
|
||||
qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
|
||||
|
||||
q, k, v = qkv[0] * self.scale, qkv[1], qkv[2]
|
||||
attn = q @ k.transpose(-2, -1)
|
||||
|
||||
attn = attn.softmax(dim=-1)
|
||||
attn = self.attn_drop(attn)
|
||||
|
||||
x = (attn @ v).transpose(1, 2).reshape(B, N, C)
|
||||
x = self.proj(x)
|
||||
x = self.proj_drop(x)
|
||||
return x
|
||||
|
||||
|
||||
class MemEffAttention(Attention):
|
||||
def forward(self, x: Tensor, attn_bias=None) -> Tensor:
|
||||
if not XFORMERS_AVAILABLE:
|
||||
assert attn_bias is None, "xFormers is required for nested tensors usage"
|
||||
return super().forward(x)
|
||||
|
||||
B, N, C = x.shape
|
||||
qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads)
|
||||
|
||||
q, k, v = unbind(qkv, 2)
|
||||
|
||||
x = memory_efficient_attention(q, k, v, attn_bias=attn_bias)
|
||||
x = x.reshape([B, N, C])
|
||||
|
||||
x = self.proj(x)
|
||||
x = self.proj_drop(x)
|
||||
return x
|
||||
|
||||
@@ -0,0 +1,252 @@
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
||||
# All rights reserved.
|
||||
#
|
||||
# This source code is licensed under the license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
# References:
|
||||
# https://github.com/facebookresearch/dino/blob/master/vision_transformer.py
|
||||
# https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/patch_embed.py
|
||||
|
||||
import logging
|
||||
from typing import Callable, List, Any, Tuple, Dict
|
||||
|
||||
import torch
|
||||
from torch import nn, Tensor
|
||||
|
||||
from .attention import Attention, MemEffAttention
|
||||
from .drop_path import DropPath
|
||||
from .layer_scale import LayerScale
|
||||
from .mlp import Mlp
|
||||
|
||||
|
||||
logger = logging.getLogger("dinov2")
|
||||
|
||||
|
||||
try:
|
||||
from xformers.ops import fmha
|
||||
from xformers.ops import scaled_index_add, index_select_cat
|
||||
|
||||
XFORMERS_AVAILABLE = True
|
||||
except ImportError:
|
||||
logger.warning("xFormers not available")
|
||||
XFORMERS_AVAILABLE = False
|
||||
|
||||
|
||||
class Block(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
dim: int,
|
||||
num_heads: int,
|
||||
mlp_ratio: float = 4.0,
|
||||
qkv_bias: bool = False,
|
||||
proj_bias: bool = True,
|
||||
ffn_bias: bool = True,
|
||||
drop: float = 0.0,
|
||||
attn_drop: float = 0.0,
|
||||
init_values=None,
|
||||
drop_path: float = 0.0,
|
||||
act_layer: Callable[..., nn.Module] = nn.GELU,
|
||||
norm_layer: Callable[..., nn.Module] = nn.LayerNorm,
|
||||
attn_class: Callable[..., nn.Module] = Attention,
|
||||
ffn_layer: Callable[..., nn.Module] = Mlp,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
# print(f"biases: qkv: {qkv_bias}, proj: {proj_bias}, ffn: {ffn_bias}")
|
||||
self.norm1 = norm_layer(dim)
|
||||
self.attn = attn_class(
|
||||
dim,
|
||||
num_heads=num_heads,
|
||||
qkv_bias=qkv_bias,
|
||||
proj_bias=proj_bias,
|
||||
attn_drop=attn_drop,
|
||||
proj_drop=drop,
|
||||
)
|
||||
self.ls1 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity()
|
||||
self.drop_path1 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
|
||||
|
||||
self.norm2 = norm_layer(dim)
|
||||
mlp_hidden_dim = int(dim * mlp_ratio)
|
||||
self.mlp = ffn_layer(
|
||||
in_features=dim,
|
||||
hidden_features=mlp_hidden_dim,
|
||||
act_layer=act_layer,
|
||||
drop=drop,
|
||||
bias=ffn_bias,
|
||||
)
|
||||
self.ls2 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity()
|
||||
self.drop_path2 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
|
||||
|
||||
self.sample_drop_ratio = drop_path
|
||||
|
||||
def forward(self, x: Tensor) -> Tensor:
|
||||
def attn_residual_func(x: Tensor) -> Tensor:
|
||||
return self.ls1(self.attn(self.norm1(x)))
|
||||
|
||||
def ffn_residual_func(x: Tensor) -> Tensor:
|
||||
return self.ls2(self.mlp(self.norm2(x)))
|
||||
|
||||
if self.training and self.sample_drop_ratio > 0.1:
|
||||
# the overhead is compensated only for a drop path rate larger than 0.1
|
||||
x = drop_add_residual_stochastic_depth(
|
||||
x,
|
||||
residual_func=attn_residual_func,
|
||||
sample_drop_ratio=self.sample_drop_ratio,
|
||||
)
|
||||
x = drop_add_residual_stochastic_depth(
|
||||
x,
|
||||
residual_func=ffn_residual_func,
|
||||
sample_drop_ratio=self.sample_drop_ratio,
|
||||
)
|
||||
elif self.training and self.sample_drop_ratio > 0.0:
|
||||
x = x + self.drop_path1(attn_residual_func(x))
|
||||
x = x + self.drop_path1(ffn_residual_func(x)) # FIXME: drop_path2
|
||||
else:
|
||||
x = x + attn_residual_func(x)
|
||||
x = x + ffn_residual_func(x)
|
||||
return x
|
||||
|
||||
|
||||
def drop_add_residual_stochastic_depth(
|
||||
x: Tensor,
|
||||
residual_func: Callable[[Tensor], Tensor],
|
||||
sample_drop_ratio: float = 0.0,
|
||||
) -> Tensor:
|
||||
# 1) extract subset using permutation
|
||||
b, n, d = x.shape
|
||||
sample_subset_size = max(int(b * (1 - sample_drop_ratio)), 1)
|
||||
brange = (torch.randperm(b, device=x.device))[:sample_subset_size]
|
||||
x_subset = x[brange]
|
||||
|
||||
# 2) apply residual_func to get residual
|
||||
residual = residual_func(x_subset)
|
||||
|
||||
x_flat = x.flatten(1)
|
||||
residual = residual.flatten(1)
|
||||
|
||||
residual_scale_factor = b / sample_subset_size
|
||||
|
||||
# 3) add the residual
|
||||
x_plus_residual = torch.index_add(x_flat, 0, brange, residual.to(dtype=x.dtype), alpha=residual_scale_factor)
|
||||
return x_plus_residual.view_as(x)
|
||||
|
||||
|
||||
def get_branges_scales(x, sample_drop_ratio=0.0):
|
||||
b, n, d = x.shape
|
||||
sample_subset_size = max(int(b * (1 - sample_drop_ratio)), 1)
|
||||
brange = (torch.randperm(b, device=x.device))[:sample_subset_size]
|
||||
residual_scale_factor = b / sample_subset_size
|
||||
return brange, residual_scale_factor
|
||||
|
||||
|
||||
def add_residual(x, brange, residual, residual_scale_factor, scaling_vector=None):
|
||||
if scaling_vector is None:
|
||||
x_flat = x.flatten(1)
|
||||
residual = residual.flatten(1)
|
||||
x_plus_residual = torch.index_add(x_flat, 0, brange, residual.to(dtype=x.dtype), alpha=residual_scale_factor)
|
||||
else:
|
||||
x_plus_residual = scaled_index_add(
|
||||
x, brange, residual.to(dtype=x.dtype), scaling=scaling_vector, alpha=residual_scale_factor
|
||||
)
|
||||
return x_plus_residual
|
||||
|
||||
|
||||
attn_bias_cache: Dict[Tuple, Any] = {}
|
||||
|
||||
|
||||
def get_attn_bias_and_cat(x_list, branges=None):
|
||||
"""
|
||||
this will perform the index select, cat the tensors, and provide the attn_bias from cache
|
||||
"""
|
||||
batch_sizes = [b.shape[0] for b in branges] if branges is not None else [x.shape[0] for x in x_list]
|
||||
all_shapes = tuple((b, x.shape[1]) for b, x in zip(batch_sizes, x_list))
|
||||
if all_shapes not in attn_bias_cache.keys():
|
||||
seqlens = []
|
||||
for b, x in zip(batch_sizes, x_list):
|
||||
for _ in range(b):
|
||||
seqlens.append(x.shape[1])
|
||||
attn_bias = fmha.BlockDiagonalMask.from_seqlens(seqlens)
|
||||
attn_bias._batch_sizes = batch_sizes
|
||||
attn_bias_cache[all_shapes] = attn_bias
|
||||
|
||||
if branges is not None:
|
||||
cat_tensors = index_select_cat([x.flatten(1) for x in x_list], branges).view(1, -1, x_list[0].shape[-1])
|
||||
else:
|
||||
tensors_bs1 = tuple(x.reshape([1, -1, *x.shape[2:]]) for x in x_list)
|
||||
cat_tensors = torch.cat(tensors_bs1, dim=1)
|
||||
|
||||
return attn_bias_cache[all_shapes], cat_tensors
|
||||
|
||||
|
||||
def drop_add_residual_stochastic_depth_list(
|
||||
x_list: List[Tensor],
|
||||
residual_func: Callable[[Tensor, Any], Tensor],
|
||||
sample_drop_ratio: float = 0.0,
|
||||
scaling_vector=None,
|
||||
) -> Tensor:
|
||||
# 1) generate random set of indices for dropping samples in the batch
|
||||
branges_scales = [get_branges_scales(x, sample_drop_ratio=sample_drop_ratio) for x in x_list]
|
||||
branges = [s[0] for s in branges_scales]
|
||||
residual_scale_factors = [s[1] for s in branges_scales]
|
||||
|
||||
# 2) get attention bias and index+concat the tensors
|
||||
attn_bias, x_cat = get_attn_bias_and_cat(x_list, branges)
|
||||
|
||||
# 3) apply residual_func to get residual, and split the result
|
||||
residual_list = attn_bias.split(residual_func(x_cat, attn_bias=attn_bias)) # type: ignore
|
||||
|
||||
outputs = []
|
||||
for x, brange, residual, residual_scale_factor in zip(x_list, branges, residual_list, residual_scale_factors):
|
||||
outputs.append(add_residual(x, brange, residual, residual_scale_factor, scaling_vector).view_as(x))
|
||||
return outputs
|
||||
|
||||
|
||||
class NestedTensorBlock(Block):
|
||||
def forward_nested(self, x_list: List[Tensor]) -> List[Tensor]:
|
||||
"""
|
||||
x_list contains a list of tensors to nest together and run
|
||||
"""
|
||||
assert isinstance(self.attn, MemEffAttention)
|
||||
|
||||
if self.training and self.sample_drop_ratio > 0.0:
|
||||
|
||||
def attn_residual_func(x: Tensor, attn_bias=None) -> Tensor:
|
||||
return self.attn(self.norm1(x), attn_bias=attn_bias)
|
||||
|
||||
def ffn_residual_func(x: Tensor, attn_bias=None) -> Tensor:
|
||||
return self.mlp(self.norm2(x))
|
||||
|
||||
x_list = drop_add_residual_stochastic_depth_list(
|
||||
x_list,
|
||||
residual_func=attn_residual_func,
|
||||
sample_drop_ratio=self.sample_drop_ratio,
|
||||
scaling_vector=self.ls1.gamma if isinstance(self.ls1, LayerScale) else None,
|
||||
)
|
||||
x_list = drop_add_residual_stochastic_depth_list(
|
||||
x_list,
|
||||
residual_func=ffn_residual_func,
|
||||
sample_drop_ratio=self.sample_drop_ratio,
|
||||
scaling_vector=self.ls2.gamma if isinstance(self.ls1, LayerScale) else None,
|
||||
)
|
||||
return x_list
|
||||
else:
|
||||
|
||||
def attn_residual_func(x: Tensor, attn_bias=None) -> Tensor:
|
||||
return self.ls1(self.attn(self.norm1(x), attn_bias=attn_bias))
|
||||
|
||||
def ffn_residual_func(x: Tensor, attn_bias=None) -> Tensor:
|
||||
return self.ls2(self.mlp(self.norm2(x)))
|
||||
|
||||
attn_bias, x = get_attn_bias_and_cat(x_list)
|
||||
x = x + attn_residual_func(x, attn_bias=attn_bias)
|
||||
x = x + ffn_residual_func(x)
|
||||
return attn_bias.split(x)
|
||||
|
||||
def forward(self, x_or_x_list):
|
||||
if isinstance(x_or_x_list, Tensor):
|
||||
return super().forward(x_or_x_list)
|
||||
elif isinstance(x_or_x_list, list):
|
||||
assert XFORMERS_AVAILABLE, "Please install xFormers for nested tensors usage"
|
||||
return self.forward_nested(x_or_x_list)
|
||||
else:
|
||||
raise AssertionError
|
||||
@@ -0,0 +1,35 @@
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
||||
# All rights reserved.
|
||||
#
|
||||
# This source code is licensed under the license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
# References:
|
||||
# https://github.com/facebookresearch/dino/blob/master/vision_transformer.py
|
||||
# https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/drop.py
|
||||
|
||||
|
||||
from torch import nn
|
||||
|
||||
|
||||
def drop_path(x, drop_prob: float = 0.0, training: bool = False):
|
||||
if drop_prob == 0.0 or not training:
|
||||
return x
|
||||
keep_prob = 1 - drop_prob
|
||||
shape = (x.shape[0],) + (1,) * (x.ndim - 1) # work with diff dim tensors, not just 2D ConvNets
|
||||
random_tensor = x.new_empty(shape).bernoulli_(keep_prob)
|
||||
if keep_prob > 0.0:
|
||||
random_tensor.div_(keep_prob)
|
||||
output = x * random_tensor
|
||||
return output
|
||||
|
||||
|
||||
class DropPath(nn.Module):
|
||||
"""Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks)."""
|
||||
|
||||
def __init__(self, drop_prob=None):
|
||||
super(DropPath, self).__init__()
|
||||
self.drop_prob = drop_prob
|
||||
|
||||
def forward(self, x):
|
||||
return drop_path(x, self.drop_prob, self.training)
|
||||
@@ -0,0 +1,28 @@
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
||||
# All rights reserved.
|
||||
#
|
||||
# This source code is licensed under the license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
# Modified from: https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/vision_transformer.py#L103-L110
|
||||
|
||||
from typing import Union
|
||||
|
||||
import torch
|
||||
from torch import Tensor
|
||||
from torch import nn
|
||||
|
||||
|
||||
class LayerScale(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
dim: int,
|
||||
init_values: Union[float, Tensor] = 1e-5,
|
||||
inplace: bool = False,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.inplace = inplace
|
||||
self.gamma = nn.Parameter(init_values * torch.ones(dim))
|
||||
|
||||
def forward(self, x: Tensor) -> Tensor:
|
||||
return x.mul_(self.gamma) if self.inplace else x * self.gamma
|
||||
@@ -0,0 +1,41 @@
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
||||
# All rights reserved.
|
||||
#
|
||||
# This source code is licensed under the license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
# References:
|
||||
# https://github.com/facebookresearch/dino/blob/master/vision_transformer.py
|
||||
# https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/mlp.py
|
||||
|
||||
|
||||
from typing import Callable, Optional
|
||||
|
||||
from torch import Tensor, nn
|
||||
|
||||
|
||||
class Mlp(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
in_features: int,
|
||||
hidden_features: Optional[int] = None,
|
||||
out_features: Optional[int] = None,
|
||||
act_layer: Callable[..., nn.Module] = nn.GELU,
|
||||
drop: float = 0.0,
|
||||
bias: bool = True,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
out_features = out_features or in_features
|
||||
hidden_features = hidden_features or in_features
|
||||
self.fc1 = nn.Linear(in_features, hidden_features, bias=bias)
|
||||
self.act = act_layer()
|
||||
self.fc2 = nn.Linear(hidden_features, out_features, bias=bias)
|
||||
self.drop = nn.Dropout(drop)
|
||||
|
||||
def forward(self, x: Tensor) -> Tensor:
|
||||
x = self.fc1(x)
|
||||
x = self.act(x)
|
||||
x = self.drop(x)
|
||||
x = self.fc2(x)
|
||||
x = self.drop(x)
|
||||
return x
|
||||
@@ -0,0 +1,89 @@
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
||||
# All rights reserved.
|
||||
#
|
||||
# This source code is licensed under the license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
# References:
|
||||
# https://github.com/facebookresearch/dino/blob/master/vision_transformer.py
|
||||
# https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/patch_embed.py
|
||||
|
||||
from typing import Callable, Optional, Tuple, Union
|
||||
|
||||
from torch import Tensor
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
def make_2tuple(x):
|
||||
if isinstance(x, tuple):
|
||||
assert len(x) == 2
|
||||
return x
|
||||
|
||||
assert isinstance(x, int)
|
||||
return (x, x)
|
||||
|
||||
|
||||
class PatchEmbed(nn.Module):
|
||||
"""
|
||||
2D image to patch embedding: (B,C,H,W) -> (B,N,D)
|
||||
|
||||
Args:
|
||||
img_size: Image size.
|
||||
patch_size: Patch token size.
|
||||
in_chans: Number of input image channels.
|
||||
embed_dim: Number of linear projection output channels.
|
||||
norm_layer: Normalization layer.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
img_size: Union[int, Tuple[int, int]] = 224,
|
||||
patch_size: Union[int, Tuple[int, int]] = 16,
|
||||
in_chans: int = 3,
|
||||
embed_dim: int = 768,
|
||||
norm_layer: Optional[Callable] = None,
|
||||
flatten_embedding: bool = True,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
|
||||
image_HW = make_2tuple(img_size)
|
||||
patch_HW = make_2tuple(patch_size)
|
||||
patch_grid_size = (
|
||||
image_HW[0] // patch_HW[0],
|
||||
image_HW[1] // patch_HW[1],
|
||||
)
|
||||
|
||||
self.img_size = image_HW
|
||||
self.patch_size = patch_HW
|
||||
self.patches_resolution = patch_grid_size
|
||||
self.num_patches = patch_grid_size[0] * patch_grid_size[1]
|
||||
|
||||
self.in_chans = in_chans
|
||||
self.embed_dim = embed_dim
|
||||
|
||||
self.flatten_embedding = flatten_embedding
|
||||
|
||||
self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_HW, stride=patch_HW)
|
||||
self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()
|
||||
|
||||
def forward(self, x: Tensor) -> Tensor:
|
||||
_, _, H, W = x.shape
|
||||
patch_H, patch_W = self.patch_size
|
||||
|
||||
assert H % patch_H == 0, f"Input image height {H} is not a multiple of patch height {patch_H}"
|
||||
assert W % patch_W == 0, f"Input image width {W} is not a multiple of patch width: {patch_W}"
|
||||
|
||||
x = self.proj(x) # B C H W
|
||||
H, W = x.size(2), x.size(3)
|
||||
x = x.flatten(2).transpose(1, 2) # B HW C
|
||||
x = self.norm(x)
|
||||
if not self.flatten_embedding:
|
||||
x = x.reshape(-1, H, W, self.embed_dim) # B H W C
|
||||
return x
|
||||
|
||||
def flops(self) -> float:
|
||||
Ho, Wo = self.patches_resolution
|
||||
flops = Ho * Wo * self.embed_dim * self.in_chans * (self.patch_size[0] * self.patch_size[1])
|
||||
if self.norm is not None:
|
||||
flops += Ho * Wo * self.embed_dim
|
||||
return flops
|
||||
@@ -0,0 +1,63 @@
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
||||
# All rights reserved.
|
||||
#
|
||||
# This source code is licensed under the license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
from typing import Callable, Optional
|
||||
|
||||
from torch import Tensor, nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
class SwiGLUFFN(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
in_features: int,
|
||||
hidden_features: Optional[int] = None,
|
||||
out_features: Optional[int] = None,
|
||||
act_layer: Callable[..., nn.Module] = None,
|
||||
drop: float = 0.0,
|
||||
bias: bool = True,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
out_features = out_features or in_features
|
||||
hidden_features = hidden_features or in_features
|
||||
self.w12 = nn.Linear(in_features, 2 * hidden_features, bias=bias)
|
||||
self.w3 = nn.Linear(hidden_features, out_features, bias=bias)
|
||||
|
||||
def forward(self, x: Tensor) -> Tensor:
|
||||
x12 = self.w12(x)
|
||||
x1, x2 = x12.chunk(2, dim=-1)
|
||||
hidden = F.silu(x1) * x2
|
||||
return self.w3(hidden)
|
||||
|
||||
|
||||
try:
|
||||
from xformers.ops import SwiGLU
|
||||
|
||||
XFORMERS_AVAILABLE = True
|
||||
except ImportError:
|
||||
SwiGLU = SwiGLUFFN
|
||||
XFORMERS_AVAILABLE = False
|
||||
|
||||
|
||||
class SwiGLUFFNFused(SwiGLU):
|
||||
def __init__(
|
||||
self,
|
||||
in_features: int,
|
||||
hidden_features: Optional[int] = None,
|
||||
out_features: Optional[int] = None,
|
||||
act_layer: Callable[..., nn.Module] = None,
|
||||
drop: float = 0.0,
|
||||
bias: bool = True,
|
||||
) -> None:
|
||||
out_features = out_features or in_features
|
||||
hidden_features = hidden_features or in_features
|
||||
hidden_features = (int(hidden_features * 2 / 3) + 7) // 8 * 8
|
||||
super().__init__(
|
||||
in_features=in_features,
|
||||
hidden_features=hidden_features,
|
||||
out_features=out_features,
|
||||
bias=bias,
|
||||
)
|
||||
@@ -0,0 +1,309 @@
|
||||
import cv2
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from torchvision.transforms import Compose
|
||||
|
||||
from .dinov2 import DINOv2
|
||||
from .util.blocks import FeatureFusionBlock, _make_scratch
|
||||
from .util.transform import Resize, NormalizeImage, PrepareForNet
|
||||
|
||||
|
||||
def _make_fusion_block(features, use_bn, size=None):
|
||||
return FeatureFusionBlock(
|
||||
features,
|
||||
nn.ReLU(False),
|
||||
deconv=False,
|
||||
bn=use_bn,
|
||||
expand=False,
|
||||
align_corners=True,
|
||||
size=size,
|
||||
)
|
||||
|
||||
|
||||
class ConvBlock(nn.Module):
|
||||
def __init__(self, in_feature, out_feature):
|
||||
super().__init__()
|
||||
|
||||
self.conv_block = nn.Sequential(
|
||||
nn.Conv2d(in_feature, out_feature, kernel_size=3, stride=1, padding=1),
|
||||
nn.BatchNorm2d(out_feature),
|
||||
nn.ReLU(True)
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
return self.conv_block(x)
|
||||
|
||||
|
||||
class DPTHead(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
in_channels,
|
||||
features=256,
|
||||
use_bn=False,
|
||||
out_channels=[256, 512, 1024, 1024],
|
||||
use_clstoken=False,
|
||||
):
|
||||
super(DPTHead, self).__init__()
|
||||
|
||||
self.use_clstoken = use_clstoken
|
||||
|
||||
self.projects = nn.ModuleList([
|
||||
nn.Conv2d(
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channel,
|
||||
kernel_size=1,
|
||||
stride=1,
|
||||
padding=0,
|
||||
) for out_channel in out_channels
|
||||
])
|
||||
|
||||
self.resize_layers = nn.ModuleList([
|
||||
nn.ConvTranspose2d(
|
||||
in_channels=out_channels[0],
|
||||
out_channels=out_channels[0],
|
||||
kernel_size=4,
|
||||
stride=4,
|
||||
padding=0),
|
||||
nn.ConvTranspose2d(
|
||||
in_channels=out_channels[1],
|
||||
out_channels=out_channels[1],
|
||||
kernel_size=2,
|
||||
stride=2,
|
||||
padding=0),
|
||||
nn.Identity(),
|
||||
nn.Conv2d(
|
||||
in_channels=out_channels[3],
|
||||
out_channels=out_channels[3],
|
||||
kernel_size=3,
|
||||
stride=2,
|
||||
padding=1)
|
||||
])
|
||||
|
||||
if use_clstoken:
|
||||
self.readout_projects = nn.ModuleList()
|
||||
for _ in range(len(self.projects)):
|
||||
self.readout_projects.append(
|
||||
nn.Sequential(
|
||||
nn.Linear(2 * in_channels, in_channels),
|
||||
nn.GELU()))
|
||||
|
||||
self.scratch = _make_scratch(
|
||||
out_channels,
|
||||
features,
|
||||
groups=1,
|
||||
expand=False,
|
||||
)
|
||||
|
||||
self.scratch.stem_transpose = None
|
||||
|
||||
self.scratch.refinenet1 = _make_fusion_block(features, use_bn)
|
||||
self.scratch.refinenet2 = _make_fusion_block(features, use_bn)
|
||||
self.scratch.refinenet3 = _make_fusion_block(features, use_bn)
|
||||
self.scratch.refinenet4 = _make_fusion_block(features, use_bn)
|
||||
|
||||
head_features_1 = features
|
||||
head_features_2 = 32
|
||||
|
||||
self.scratch.output_conv1 = nn.Conv2d(head_features_1, head_features_1 // 2, kernel_size=3, stride=1, padding=1)
|
||||
|
||||
self.scratch.output_conv2 = nn.Sequential(
|
||||
nn.Conv2d(head_features_1 // 2, head_features_2, kernel_size=3, stride=1, padding=1),
|
||||
nn.ReLU(True),
|
||||
nn.Conv2d(head_features_2, 1, kernel_size=1, stride=1, padding=0),
|
||||
nn.ReLU(True),
|
||||
nn.Identity(),
|
||||
)
|
||||
|
||||
def forward(self, out_features, patch_h, patch_w, out_h, out_w):
|
||||
bs = out_features[0][0].shape[0]
|
||||
out = []
|
||||
for i, x in enumerate(out_features):
|
||||
if self.use_clstoken:
|
||||
x, cls_token = x[0][:bs//2], x[1][:bs//2]
|
||||
readout = cls_token.unsqueeze(1).expand_as(x)
|
||||
x = self.readout_projects[i](torch.cat((x, readout), -1))
|
||||
else:
|
||||
x = x[0][:bs//2]
|
||||
|
||||
x = x.permute(0, 2, 1).reshape((x.shape[0], x.shape[-1], patch_h, patch_w))
|
||||
|
||||
x = self.projects[i](x)
|
||||
x = self.resize_layers[i](x)
|
||||
|
||||
out.append(x)
|
||||
|
||||
layer_1, layer_2, layer_3, layer_4 = out
|
||||
|
||||
layer_1_rn = self.scratch.layer1_rn(layer_1)
|
||||
layer_2_rn = self.scratch.layer2_rn(layer_2)
|
||||
layer_3_rn = self.scratch.layer3_rn(layer_3)
|
||||
layer_4_rn = self.scratch.layer4_rn(layer_4)
|
||||
|
||||
path_4 = self.scratch.refinenet4(layer_4_rn, size=layer_3_rn.shape[2:])
|
||||
path_3 = self.scratch.refinenet3(path_4, layer_3_rn, size=layer_2_rn.shape[2:])
|
||||
path_2 = self.scratch.refinenet2(path_3, layer_2_rn, size=layer_1_rn.shape[2:])
|
||||
path_1 = self.scratch.refinenet1(path_2, layer_1_rn)
|
||||
|
||||
out = self.scratch.output_conv1(path_1)
|
||||
out = F.interpolate(out, (out_h, out_w), mode="bilinear", align_corners=True)
|
||||
|
||||
idepth = self.scratch.output_conv2(out)
|
||||
return idepth
|
||||
|
||||
|
||||
class DPTFeat(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
in_channels,
|
||||
features=256,
|
||||
use_bn=False,
|
||||
out_channels=[256, 512, 1024, 1024],
|
||||
use_clstoken=False,
|
||||
):
|
||||
super(DPTFeat, self).__init__()
|
||||
|
||||
self.use_clstoken = use_clstoken
|
||||
|
||||
self.projects = nn.ModuleList([
|
||||
nn.Conv2d(
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channel,
|
||||
kernel_size=1,
|
||||
stride=1,
|
||||
padding=0,
|
||||
) for out_channel in out_channels
|
||||
])
|
||||
|
||||
self.resize_layers = nn.ModuleList([
|
||||
nn.ConvTranspose2d(
|
||||
in_channels=out_channels[0],
|
||||
out_channels=out_channels[0],
|
||||
kernel_size=4,
|
||||
stride=4,
|
||||
padding=0),
|
||||
nn.ConvTranspose2d(
|
||||
in_channels=out_channels[1],
|
||||
out_channels=out_channels[1],
|
||||
kernel_size=2,
|
||||
stride=2,
|
||||
padding=0),
|
||||
nn.Identity(),
|
||||
nn.Conv2d(
|
||||
in_channels=out_channels[3],
|
||||
out_channels=out_channels[3],
|
||||
kernel_size=3,
|
||||
stride=2,
|
||||
padding=1)
|
||||
])
|
||||
|
||||
if use_clstoken:
|
||||
self.readout_projects = nn.ModuleList()
|
||||
for _ in range(len(self.projects)):
|
||||
self.readout_projects.append(
|
||||
nn.Sequential(
|
||||
nn.Linear(2 * in_channels, in_channels),
|
||||
nn.GELU()))
|
||||
|
||||
self.scratch = _make_scratch(
|
||||
out_channels,
|
||||
features,
|
||||
groups=1,
|
||||
expand=False,
|
||||
)
|
||||
|
||||
self.scratch.stem_transpose = None
|
||||
|
||||
self.scratch.refinenet1 = _make_fusion_block(features, use_bn)
|
||||
self.scratch.refinenet2 = _make_fusion_block(features, use_bn)
|
||||
self.scratch.refinenet3 = _make_fusion_block(features, use_bn)
|
||||
self.scratch.refinenet4 = _make_fusion_block(features, use_bn)
|
||||
|
||||
def forward(self, out_features, patch_h, patch_w, out_h, out_w):
|
||||
bs = out_features[0][0].shape[0]
|
||||
out = []
|
||||
for i, x in enumerate(out_features):
|
||||
if self.use_clstoken:
|
||||
x, cls_token = x[0], x[1]
|
||||
readout = cls_token.unsqueeze(1).expand_as(x)
|
||||
x = self.readout_projects[i](torch.cat((x, readout), -1))
|
||||
else:
|
||||
x = x[0]
|
||||
|
||||
x = x.permute(0, 2, 1).reshape((x.shape[0], x.shape[-1], patch_h, patch_w))
|
||||
|
||||
x = self.projects[i](x)
|
||||
x = self.resize_layers[i](x)
|
||||
|
||||
out.append(x)
|
||||
|
||||
layer_1, layer_2, layer_3, layer_4 = out
|
||||
|
||||
layer_1_rn = self.scratch.layer1_rn(layer_1)
|
||||
layer_2_rn = self.scratch.layer2_rn(layer_2)
|
||||
layer_3_rn = self.scratch.layer3_rn(layer_3)
|
||||
layer_4_rn = self.scratch.layer4_rn(layer_4)
|
||||
|
||||
layer_1_rn = F.interpolate(layer_1_rn, (out_h, out_w), mode="bilinear", align_corners=True)
|
||||
layer_2_rn = F.interpolate(layer_2_rn, (out_h // 2, out_w // 2), mode="bilinear", align_corners=True)
|
||||
layer_3_rn = F.interpolate(layer_3_rn, (out_h // 4, out_w // 4), mode="bilinear", align_corners=True)
|
||||
layer_4_rn = F.interpolate(layer_4_rn, (out_h//8, out_w//8), mode="bilinear", align_corners=True)
|
||||
|
||||
out_features = [layer_1_rn[:bs//2], layer_2_rn[:bs//2], layer_3_rn[:bs//2]]
|
||||
|
||||
path_4 = self.scratch.refinenet4(layer_4_rn, size=layer_3_rn.shape[2:])
|
||||
path_3 = self.scratch.refinenet3(path_4, layer_3_rn, size=layer_2_rn.shape[2:])
|
||||
path_2 = self.scratch.refinenet2(path_3, layer_2_rn, size=layer_1_rn.shape[2:])
|
||||
path_1 = self.scratch.refinenet1(path_2, layer_1_rn)
|
||||
|
||||
return out_features, path_1[:bs//2], path_1[bs//2:]
|
||||
|
||||
|
||||
class DepthAnythingV2(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
encoder='vitl',
|
||||
features=256,
|
||||
out_channels=[256, 512, 1024, 1024],
|
||||
use_bn=False,
|
||||
use_clstoken=False,
|
||||
):
|
||||
super(DepthAnythingV2, self).__init__()
|
||||
|
||||
self.intermediate_layer_idx = {
|
||||
'vits': [2, 5, 8, 11],
|
||||
'vitb': [2, 5, 8, 11],
|
||||
'vitl': [4, 11, 17, 23],
|
||||
'vitg': [9, 19, 29, 39]
|
||||
}
|
||||
self.encoder = encoder
|
||||
self.pretrained = DINOv2(model_name=encoder)
|
||||
|
||||
self.depth_head = DPTHead(self.pretrained.embed_dim, features, use_bn,
|
||||
out_channels=out_channels, use_clstoken=use_clstoken)
|
||||
self.depth_feat = DPTFeat(self.pretrained.embed_dim, features, use_bn,
|
||||
out_channels=out_channels, use_clstoken=use_clstoken)
|
||||
|
||||
|
||||
def forward(self, x, out_h, out_w):
|
||||
patch_h, patch_w = x.shape[-2] // 14, x.shape[-1] // 14
|
||||
|
||||
features = self.pretrained.get_intermediate_layers(x, self.intermediate_layer_idx[self.encoder], return_class_token=True)
|
||||
|
||||
d_features, left_feat, right_feat = self.depth_feat(features, patch_h, patch_w, out_h, out_w)
|
||||
idepth = self.depth_head(features, patch_h, patch_w, out_h, out_w)
|
||||
|
||||
return d_features, left_feat, right_feat, idepth
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def forward_test(self, x, out_h, out_w):
|
||||
patch_h, patch_w = x.shape[-2] // 14, x.shape[-1] // 14
|
||||
|
||||
features = self.pretrained.get_intermediate_layers(x, self.intermediate_layer_idx[self.encoder],
|
||||
return_class_token=True)
|
||||
|
||||
d_features, left_feat, right_feat = self.depth_feat(features, patch_h, patch_w, out_h, out_w)
|
||||
idepth = self.depth_head(features, patch_h, patch_w, out_h, out_w)
|
||||
|
||||
return d_features, left_feat, right_feat, idepth
|
||||
@@ -0,0 +1,149 @@
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
def _make_scratch(in_shape, out_shape, groups=1, expand=False):
|
||||
scratch = nn.Module()
|
||||
|
||||
out_shape1 = out_shape
|
||||
out_shape2 = out_shape
|
||||
out_shape3 = out_shape
|
||||
if len(in_shape) >= 4:
|
||||
out_shape4 = out_shape
|
||||
|
||||
if expand:
|
||||
out_shape1 = out_shape
|
||||
out_shape2 = out_shape * 2
|
||||
out_shape3 = out_shape * 4
|
||||
if len(in_shape) >= 4:
|
||||
out_shape4 = out_shape * 8
|
||||
|
||||
scratch.layer1_rn = nn.Conv2d(in_shape[0], out_shape1, kernel_size=3, stride=1, padding=1, bias=False, groups=groups)
|
||||
scratch.layer2_rn = nn.Conv2d(in_shape[1], out_shape2, kernel_size=3, stride=1, padding=1, bias=False, groups=groups)
|
||||
scratch.layer3_rn = nn.Conv2d(in_shape[2], out_shape3, kernel_size=3, stride=1, padding=1, bias=False, groups=groups)
|
||||
if len(in_shape) >= 4:
|
||||
scratch.layer4_rn = nn.Conv2d(in_shape[3], out_shape4, kernel_size=3, stride=1, padding=1, bias=False, groups=groups)
|
||||
|
||||
return scratch
|
||||
|
||||
|
||||
class ResidualConvUnit(nn.Module):
|
||||
"""Residual convolution module.
|
||||
"""
|
||||
|
||||
def __init__(self, features, activation, bn):
|
||||
"""Init.
|
||||
|
||||
Args:
|
||||
features (int): number of features
|
||||
"""
|
||||
super().__init__()
|
||||
|
||||
self.bn = bn
|
||||
|
||||
self.groups=1
|
||||
|
||||
self.conv1 = nn.Conv2d(features, features, kernel_size=3, stride=1, padding=1, bias=True, groups=self.groups)
|
||||
|
||||
self.conv2 = nn.Conv2d(features, features, kernel_size=3, stride=1, padding=1, bias=True, groups=self.groups)
|
||||
|
||||
if self.bn == True:
|
||||
self.bn1 = nn.BatchNorm2d(features)
|
||||
self.bn2 = nn.BatchNorm2d(features)
|
||||
|
||||
self.activation = activation
|
||||
|
||||
self.skip_add = nn.quantized.FloatFunctional()
|
||||
|
||||
def forward(self, x):
|
||||
"""Forward pass.
|
||||
|
||||
Args:
|
||||
x (tensor): input
|
||||
|
||||
Returns:
|
||||
tensor: output
|
||||
"""
|
||||
|
||||
out = self.activation(x)
|
||||
out = self.conv1(out)
|
||||
if self.bn == True:
|
||||
out = self.bn1(out)
|
||||
|
||||
out = self.activation(out)
|
||||
out = self.conv2(out)
|
||||
if self.bn == True:
|
||||
out = self.bn2(out)
|
||||
|
||||
if self.groups > 1:
|
||||
out = self.conv_merge(out)
|
||||
|
||||
return self.skip_add.add(out, x)
|
||||
|
||||
|
||||
class FeatureFusionBlock(nn.Module):
|
||||
"""Feature fusion block.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
features,
|
||||
activation,
|
||||
deconv=False,
|
||||
bn=False,
|
||||
expand=False,
|
||||
align_corners=True,
|
||||
size=None
|
||||
):
|
||||
"""Init.
|
||||
|
||||
Args:
|
||||
features (int): number of features
|
||||
"""
|
||||
super(FeatureFusionBlock, self).__init__()
|
||||
|
||||
self.deconv = deconv
|
||||
self.align_corners = align_corners
|
||||
|
||||
self.groups=1
|
||||
|
||||
self.expand = expand
|
||||
out_features = features
|
||||
if self.expand == True:
|
||||
out_features = features // 2
|
||||
|
||||
self.out_conv = nn.Conv2d(features, out_features, kernel_size=1, stride=1, padding=0, bias=True, groups=1)
|
||||
|
||||
self.resConfUnit1 = ResidualConvUnit(features, activation, bn)
|
||||
self.resConfUnit2 = ResidualConvUnit(features, activation, bn)
|
||||
|
||||
self.skip_add = nn.quantized.FloatFunctional()
|
||||
|
||||
self.size=size
|
||||
|
||||
def forward(self, *xs, size=None):
|
||||
"""Forward pass.
|
||||
|
||||
Returns:
|
||||
tensor: output
|
||||
"""
|
||||
output = xs[0]
|
||||
|
||||
if len(xs) == 2:
|
||||
res = self.resConfUnit1(xs[1])
|
||||
output = self.skip_add.add(output, res)
|
||||
|
||||
output = self.resConfUnit2(output)
|
||||
|
||||
if (size is None) and (self.size is None):
|
||||
modifier = None #{"scale_factor": 2}
|
||||
elif size is None:
|
||||
modifier = {"size": self.size}
|
||||
else:
|
||||
modifier = {"size": size}
|
||||
|
||||
if modifier:
|
||||
output = nn.functional.interpolate(output, **modifier, mode="bilinear", align_corners=self.align_corners)
|
||||
|
||||
output = self.out_conv(output)
|
||||
|
||||
return output
|
||||
@@ -0,0 +1,158 @@
|
||||
import numpy as np
|
||||
import cv2
|
||||
|
||||
|
||||
class Resize(object):
|
||||
"""Resize sample to given size (width, height).
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
width,
|
||||
height,
|
||||
resize_target=True,
|
||||
keep_aspect_ratio=False,
|
||||
ensure_multiple_of=1,
|
||||
resize_method="lower_bound",
|
||||
image_interpolation_method=cv2.INTER_AREA,
|
||||
):
|
||||
"""Init.
|
||||
|
||||
Args:
|
||||
width (int): desired output width
|
||||
height (int): desired output height
|
||||
resize_target (bool, optional):
|
||||
True: Resize the full sample (image, mask, target).
|
||||
False: Resize image only.
|
||||
Defaults to True.
|
||||
keep_aspect_ratio (bool, optional):
|
||||
True: Keep the aspect ratio of the input sample.
|
||||
Output sample might not have the given width and height, and
|
||||
resize behaviour depends on the parameter 'resize_method'.
|
||||
Defaults to False.
|
||||
ensure_multiple_of (int, optional):
|
||||
Output width and height is constrained to be multiple of this parameter.
|
||||
Defaults to 1.
|
||||
resize_method (str, optional):
|
||||
"lower_bound": Output will be at least as large as the given size.
|
||||
"upper_bound": Output will be at max as large as the given size. (Output size might be smaller than given size.)
|
||||
"minimal": Scale as least as possible. (Output size might be smaller than given size.)
|
||||
Defaults to "lower_bound".
|
||||
"""
|
||||
self.__width = width
|
||||
self.__height = height
|
||||
|
||||
self.__resize_target = resize_target
|
||||
self.__keep_aspect_ratio = keep_aspect_ratio
|
||||
self.__multiple_of = ensure_multiple_of
|
||||
self.__resize_method = resize_method
|
||||
self.__image_interpolation_method = image_interpolation_method
|
||||
|
||||
def constrain_to_multiple_of(self, x, min_val=0, max_val=None):
|
||||
y = (np.round(x / self.__multiple_of) * self.__multiple_of).astype(int)
|
||||
|
||||
if max_val is not None and y > max_val:
|
||||
y = (np.floor(x / self.__multiple_of) * self.__multiple_of).astype(int)
|
||||
|
||||
if y < min_val:
|
||||
y = (np.ceil(x / self.__multiple_of) * self.__multiple_of).astype(int)
|
||||
|
||||
return y
|
||||
|
||||
def get_size(self, width, height):
|
||||
# determine new height and width
|
||||
scale_height = self.__height / height
|
||||
scale_width = self.__width / width
|
||||
|
||||
if self.__keep_aspect_ratio:
|
||||
if self.__resize_method == "lower_bound":
|
||||
# scale such that output size is lower bound
|
||||
if scale_width > scale_height:
|
||||
# fit width
|
||||
scale_height = scale_width
|
||||
else:
|
||||
# fit height
|
||||
scale_width = scale_height
|
||||
elif self.__resize_method == "upper_bound":
|
||||
# scale such that output size is upper bound
|
||||
if scale_width < scale_height:
|
||||
# fit width
|
||||
scale_height = scale_width
|
||||
else:
|
||||
# fit height
|
||||
scale_width = scale_height
|
||||
elif self.__resize_method == "minimal":
|
||||
# scale as least as possbile
|
||||
if abs(1 - scale_width) < abs(1 - scale_height):
|
||||
# fit width
|
||||
scale_height = scale_width
|
||||
else:
|
||||
# fit height
|
||||
scale_width = scale_height
|
||||
else:
|
||||
raise ValueError(f"resize_method {self.__resize_method} not implemented")
|
||||
|
||||
if self.__resize_method == "lower_bound":
|
||||
new_height = self.constrain_to_multiple_of(scale_height * height, min_val=self.__height)
|
||||
new_width = self.constrain_to_multiple_of(scale_width * width, min_val=self.__width)
|
||||
elif self.__resize_method == "upper_bound":
|
||||
new_height = self.constrain_to_multiple_of(scale_height * height, max_val=self.__height)
|
||||
new_width = self.constrain_to_multiple_of(scale_width * width, max_val=self.__width)
|
||||
elif self.__resize_method == "minimal":
|
||||
new_height = self.constrain_to_multiple_of(scale_height * height)
|
||||
new_width = self.constrain_to_multiple_of(scale_width * width)
|
||||
else:
|
||||
raise ValueError(f"resize_method {self.__resize_method} not implemented")
|
||||
|
||||
return (new_width, new_height)
|
||||
|
||||
def __call__(self, sample):
|
||||
width, height = self.get_size(sample["image"].shape[1], sample["image"].shape[0])
|
||||
|
||||
# resize sample
|
||||
sample["image"] = cv2.resize(sample["image"], (width, height), interpolation=self.__image_interpolation_method)
|
||||
|
||||
if self.__resize_target:
|
||||
if "depth" in sample:
|
||||
sample["depth"] = cv2.resize(sample["depth"], (width, height), interpolation=cv2.INTER_NEAREST)
|
||||
|
||||
if "mask" in sample:
|
||||
sample["mask"] = cv2.resize(sample["mask"].astype(np.float32), (width, height), interpolation=cv2.INTER_NEAREST)
|
||||
|
||||
return sample
|
||||
|
||||
|
||||
class NormalizeImage(object):
|
||||
"""Normlize image by given mean and std.
|
||||
"""
|
||||
|
||||
def __init__(self, mean, std):
|
||||
self.__mean = mean
|
||||
self.__std = std
|
||||
|
||||
def __call__(self, sample):
|
||||
sample["image"] = (sample["image"] - self.__mean) / self.__std
|
||||
|
||||
return sample
|
||||
|
||||
|
||||
class PrepareForNet(object):
|
||||
"""Prepare sample for usage as network input.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
def __call__(self, sample):
|
||||
image = np.transpose(sample["image"], (2, 0, 1))
|
||||
sample["image"] = np.ascontiguousarray(image).astype(np.float32)
|
||||
|
||||
if "depth" in sample:
|
||||
depth = sample["depth"].astype(np.float32)
|
||||
sample["depth"] = np.ascontiguousarray(depth)
|
||||
|
||||
if "mask" in sample:
|
||||
sample["mask"] = sample["mask"].astype(np.float32)
|
||||
sample["mask"] = np.ascontiguousarray(sample["mask"])
|
||||
|
||||
return sample
|
||||
@@ -0,0 +1,29 @@
|
||||
name: defomstereo
|
||||
channels:
|
||||
- pytorch
|
||||
- nvidia
|
||||
- xformers
|
||||
- conda-forge
|
||||
- bioconda
|
||||
- defaults
|
||||
dependencies:
|
||||
- python=3.9
|
||||
- cudatoolkit=11.8.0
|
||||
- pytorch::pytorch=2.1.1
|
||||
- pytorch::pytorch-cuda=11.8.0
|
||||
- pytorch::torchvision=0.16.1
|
||||
- xformers::xformers=0.0.23
|
||||
- matplotlib
|
||||
- tensorboard
|
||||
- scipy
|
||||
- opencv
|
||||
- tqdm
|
||||
- opt_einsum
|
||||
- imageio
|
||||
- scikit-image
|
||||
- p7zip
|
||||
- pip
|
||||
- pip:
|
||||
- gradio_imageslider
|
||||
- gradio==4.29.0
|
||||
|
||||
@@ -0,0 +1,448 @@
|
||||
from __future__ import print_function, division
|
||||
import sys
|
||||
|
||||
import argparse
|
||||
import time
|
||||
import logging
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
torch.cuda.empty_cache()
|
||||
from PIL import Image
|
||||
|
||||
from tqdm import tqdm
|
||||
from core.defom_stereo import DEFOMStereo, autocast
|
||||
|
||||
import core.stereo_datasets as datasets
|
||||
from core.utils.utils import InputPadder
|
||||
|
||||
|
||||
def count_parameters(model):
|
||||
return sum(p.numel() for p in model.parameters()), sum(p.numel() for p in model.parameters() if p.requires_grad)
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def validate_things(model, iters=32, scale_iters=8, mixed_prec=False, max_disp=192, bad_threshold=1.0):
|
||||
""" Peform validation using the FlyingThings3D (TEST) split """
|
||||
model.eval()
|
||||
val_dataset = datasets.SceneFlowDatasets(dstype='frames_finalpass', things_test=True)
|
||||
|
||||
out_list, epe_list, elapsed_list = [], [], []
|
||||
for val_id in tqdm(range(len(val_dataset))):
|
||||
data_blob = val_dataset[val_id]
|
||||
image1 = data_blob["img1"][None].cuda()
|
||||
image2 = data_blob["img2"][None].cuda()
|
||||
disp_gt = data_blob["disp"]
|
||||
valid = data_blob["valid"]
|
||||
|
||||
padder = InputPadder(image1.shape, divis_by=32)
|
||||
image1, image2 = padder.pad(image1, image2)
|
||||
|
||||
with autocast(enabled=mixed_prec):
|
||||
start = time.time()
|
||||
disp_pr = model(image1, image2, iters=iters, scale_iters=scale_iters, test_mode=True)
|
||||
end = time.time()
|
||||
if val_id > 50:
|
||||
elapsed_list.append(end-start)
|
||||
|
||||
disp_pr = padder.unpad(disp_pr).cpu().squeeze(0)
|
||||
assert disp_pr.shape == disp_gt.shape, (disp_pr.shape, disp_gt.shape)
|
||||
epe = torch.sum(torch.abs(disp_pr - disp_gt), dim=0)
|
||||
|
||||
epe = epe.flatten()
|
||||
val = (valid.flatten() >= 0.5) & (disp_gt.abs().flatten() < max_disp)
|
||||
|
||||
if np.isnan(epe[val].mean().item()):
|
||||
continue
|
||||
out = (epe > bad_threshold)
|
||||
image_out = out[val].float().mean().item()
|
||||
image_epe = epe[val].mean().item()
|
||||
if val_id < 9 or (val_id+1) % 10 == 0:
|
||||
logging.info(f"Fhythings3D Iter {val_id+1} out of {len(val_dataset)}. EPE {round(image_epe,4)} Out{bad_threshold} {round(image_out,4)}. Runtime: {format(end-start, '.3f')}s ({format(1/(end-start), '.2f')}-FPS)")
|
||||
|
||||
epe_list.append(image_epe)
|
||||
out_list.append(out[val].cpu().numpy())
|
||||
|
||||
epe_list = np.array(epe_list)
|
||||
out_list = np.concatenate(out_list)
|
||||
|
||||
epe = np.mean(epe_list)
|
||||
out = 100 * np.mean(out_list)
|
||||
avg_runtime = np.mean(elapsed_list)
|
||||
|
||||
print(f"Validation FlyingThings: EPE {epe}, Out{bad_threshold} {out}, "
|
||||
f"{format(1/avg_runtime, '.2f')}-FPS ({format(avg_runtime, '.3f')}s)")
|
||||
return {'things-epe': epe, 'things-out': out}
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def validate_eth3d(model, iters=32, scale_iters=8, mixed_prec=False):
|
||||
""" Peform validation using the ETH3D (train) split """
|
||||
model.eval()
|
||||
aug_params = {}
|
||||
val_dataset = datasets.ETH3D(aug_params, is_eval=True)
|
||||
|
||||
out_list, epe_list = [], []
|
||||
for val_id in tqdm(range(len(val_dataset))):
|
||||
data_blob = val_dataset[val_id]
|
||||
image1 = data_blob["img1"][None].cuda()
|
||||
image2 = data_blob["img2"][None].cuda()
|
||||
disp_gt = data_blob["disp"]
|
||||
valid = data_blob["valid"]
|
||||
|
||||
padder = InputPadder(image1.shape, divis_by=32)
|
||||
image1, image2 = padder.pad(image1, image2)
|
||||
|
||||
with autocast(enabled=mixed_prec):
|
||||
disp_pr = model(image1, image2, iters=iters, scale_iters=scale_iters, test_mode=True)
|
||||
disp_pr = padder.unpad(disp_pr).cpu().squeeze(0)
|
||||
assert disp_pr.shape == disp_gt.shape, (disp_pr.shape, disp_gt.shape)
|
||||
epe = torch.sum(torch.abs(disp_pr - disp_gt), dim=0)
|
||||
|
||||
epe_flattened = epe.flatten()
|
||||
val = valid.flatten() >= 0.5
|
||||
out = (epe_flattened > 1.0)
|
||||
image_out = out[val].float().mean().item()
|
||||
image_epe = epe_flattened[val].mean().item()
|
||||
logging.info(f"ETH3D {val_id+1} out of {len(val_dataset)}. EPE {round(image_epe,4)} D1 {round(image_out,4)}")
|
||||
epe_list.append(image_epe)
|
||||
out_list.append(image_out)
|
||||
|
||||
epe_list = np.array(epe_list)
|
||||
out_list = np.array(out_list)
|
||||
|
||||
epe = np.mean(epe_list)
|
||||
out1 = 100 * np.mean(out_list)
|
||||
|
||||
print("Validation ETH3D: EPE %f, Out1 %f" % (epe, out1))
|
||||
return {'eth3d-epe': epe, 'eth3d-out1': out1}
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def validate_kitti(model, iters=32, scale_iters=8, split='15', mixed_prec=False):
|
||||
""" Peform validation using the KITTI-2015/2012 (train) split """
|
||||
model.eval()
|
||||
aug_params = {}
|
||||
val_dataset = datasets.KITTI(aug_params, split=split, image_set='training', is_eval=True)
|
||||
torch.backends.cudnn.benchmark = True
|
||||
|
||||
out_list, epe_list, elapsed_list = [], [], []
|
||||
for val_id in range(len(val_dataset)):
|
||||
data_blob = val_dataset[val_id]
|
||||
image1 = data_blob["img1"][None].cuda()
|
||||
image2 = data_blob["img2"][None].cuda()
|
||||
disp_gt = data_blob["disp"]
|
||||
valid = data_blob["valid"]
|
||||
|
||||
padder = InputPadder(image1.shape, divis_by=32)
|
||||
image1, image2 = padder.pad(image1, image2)
|
||||
|
||||
with autocast(enabled=mixed_prec):
|
||||
start = time.time()
|
||||
disp_pr = model(image1, image2, iters=iters, scale_iters=scale_iters, test_mode=True)
|
||||
end = time.time()
|
||||
if val_id > 50:
|
||||
elapsed_list.append(end-start)
|
||||
|
||||
disp_pr = padder.unpad(disp_pr).cpu().squeeze(0)
|
||||
assert disp_pr.shape == disp_gt.shape, (disp_pr.shape, disp_gt.shape)
|
||||
epe = torch.sum(torch.abs(disp_pr - disp_gt), dim=0)
|
||||
|
||||
epe_flattened = epe.flatten()
|
||||
val = valid.flatten() >= 0.5
|
||||
|
||||
out = (epe_flattened > 3.0)
|
||||
image_out = out[val].float().mean().item()
|
||||
image_epe = epe_flattened[val].mean().item()
|
||||
if val_id < 9 or (val_id+1) % 10 == 0:
|
||||
logging.info(f"KITTI{split} Iter {val_id+1} out of {len(val_dataset)}. EPE {round(image_epe,4)} Out3 {round(image_out,4)}. Runtime: {format(end-start, '.3f')}s ({format(1/(end-start), '.2f')}-FPS)")
|
||||
epe_list.append(epe_flattened[val].mean().item())
|
||||
out_list.append(out[val].cpu().numpy())
|
||||
|
||||
epe_list = np.array(epe_list)
|
||||
out_list = np.concatenate(out_list)
|
||||
|
||||
epe = np.mean(epe_list)
|
||||
out3 = 100 * np.mean(out_list)
|
||||
|
||||
avg_runtime = np.mean(elapsed_list)
|
||||
|
||||
print(f"Validation KITTI{split}: EPE {epe}, Out3 {out3}, "
|
||||
f"{format(1/avg_runtime, '.2f')}-FPS ({format(avg_runtime, '.3f')}s)")
|
||||
return {f'kitti{split}-epe': epe, f'kitti{split}-out3': out3}
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def validate_middlebury(model, iters=32, scale_iters=8, split='H', mixed_prec=False):
|
||||
""" Peform validation using the Middlebury-V3 dataset """
|
||||
model.eval()
|
||||
aug_params = {}
|
||||
val_dataset = datasets.Middlebury(aug_params, split=split, is_eval=True)
|
||||
|
||||
out_list, epe_list = [], []
|
||||
for val_id in range(len(val_dataset)):
|
||||
data_blob = val_dataset[val_id]
|
||||
image1 = data_blob["img1"][None].cuda()
|
||||
image2 = data_blob["img2"][None].cuda()
|
||||
disp_gt = data_blob["disp"]
|
||||
valid = data_blob["valid"]
|
||||
|
||||
padder = InputPadder(image1.shape, divis_by=32)
|
||||
image1, image2 = padder.pad(image1, image2)
|
||||
|
||||
with autocast(enabled=mixed_prec):
|
||||
disp_pr = model(image1, image2, iters=iters, scale_iters=scale_iters, test_mode=True)
|
||||
disp_pr = padder.unpad(disp_pr).cpu().squeeze(0)
|
||||
assert disp_pr.shape == disp_gt.shape, (disp_pr.shape, disp_gt.shape)
|
||||
epe = torch.sum(torch.abs(disp_pr - disp_gt), dim=0)
|
||||
|
||||
epe_flattened = epe.flatten()
|
||||
val = (valid.reshape(-1) >= 0.5) & (disp_gt.reshape(-1) < 1000)
|
||||
|
||||
out = (epe_flattened > 2.0)
|
||||
image_out = out[val].float().mean().item()
|
||||
image_epe = epe_flattened[val].mean().item()
|
||||
logging.info(f"Middlebury Iter {val_id+1} out of {len(val_dataset)}. "
|
||||
f"EPE {round(image_epe,4)} Out2 {round(image_out,4)}")
|
||||
epe_list.append(image_epe)
|
||||
out_list.append(image_out)
|
||||
|
||||
epe_list = np.array(epe_list)
|
||||
out_list = np.array(out_list)
|
||||
|
||||
epe = np.mean(epe_list)
|
||||
out2 = 100 * np.mean(out_list)
|
||||
|
||||
print(f"Validation Middlebury{split}: EPE {epe}, Out2 {out2}")
|
||||
return {f'middlebury{split}-epe': epe, f'middlebury{split}-out2': out2}
|
||||
|
||||
|
||||
def compute_nontexture(x, weight=None, c1=0.01**2, c2=0.03**2, weight_epsilon=0.01, window=33, threshold=0.95, split="F"):
|
||||
|
||||
if split=="H":
|
||||
scale = 2
|
||||
threshold += 0.02
|
||||
elif split=="Q":
|
||||
scale = 4
|
||||
threshold += 0.03
|
||||
else:
|
||||
scale = 1
|
||||
|
||||
x = F.interpolate(x, scale_factor=scale, mode='bilinear', align_corners=True)
|
||||
|
||||
if x.max()>1:
|
||||
x = x/x.max()
|
||||
|
||||
y = F.pad(x, (1, 1, 1, 1), mode='replicate')
|
||||
_, _, h, w = y.shape
|
||||
#y = y[..., 0:h-2, 1:w-1] #(y[..., 0:h-2, 1:w-1] + y[..., 2:h, 1:w-1] + y[..., 1:h-1, 0:w-2] + y[..., 1:h-1, 2:w])/4.0
|
||||
|
||||
x = F.pad(x, (window//2, window//2, window//2, window//2), mode='replicate')
|
||||
if c1 == float('inf') and c2 == float('inf'):
|
||||
raise ValueError(
|
||||
'Both c1 and c2 are infinite, SSIM loss is zero. This is '
|
||||
'likely unintended.')
|
||||
_, _, H, W = x.shape
|
||||
|
||||
if weight is None:
|
||||
weight = torch.ones((H, W)).to(x)
|
||||
else:
|
||||
assert weight.shape == (H, W), \
|
||||
f'image shape is {(H, W)}, but weight shape is {weight.shape}'
|
||||
weight = weight[None, None, ...]
|
||||
average_pooled_weight = F.avg_pool2d(weight, (window, window), stride=(1, 1))
|
||||
weight_plus_epsilon = weight + weight_epsilon
|
||||
inverse_average_pooled_weight = 1.0 / (
|
||||
average_pooled_weight + weight_epsilon)
|
||||
|
||||
def weighted_avg_pool(z):
|
||||
weighted_avg = F.avg_pool2d(
|
||||
z * weight_plus_epsilon, (window, window), stride=(1, 1))
|
||||
return weighted_avg * inverse_average_pooled_weight
|
||||
|
||||
mu_x = weighted_avg_pool(x)
|
||||
sigma_x = weighted_avg_pool(x**2) - mu_x**2
|
||||
|
||||
def ssim(x, y):
|
||||
y = F.pad(y, (window//2, window//2, window//2, window//2), mode='replicate')
|
||||
mu_y = weighted_avg_pool(y)
|
||||
sigma_y = weighted_avg_pool(y**2) - mu_y**2
|
||||
sigma_xy = weighted_avg_pool(x * y) - mu_x * mu_y
|
||||
if c1 == float('inf'):
|
||||
ssim_n = (2 * sigma_xy + c2)
|
||||
ssim_d = (sigma_x + sigma_y + c2)
|
||||
elif c2 == float('inf'):
|
||||
ssim_n = 2 * mu_x * mu_y + c1
|
||||
ssim_d = mu_x**2 + mu_y**2 + c1
|
||||
else:
|
||||
ssim_n = (2 * mu_x * mu_y + c1) * (2 * sigma_xy + c2)
|
||||
ssim_d = (mu_x**2 + mu_y**2 + c1) * (sigma_x + sigma_y + c2)
|
||||
|
||||
result = ssim_n / ssim_d
|
||||
|
||||
result = F.avg_pool2d(result, (scale, scale), stride=(scale, scale))
|
||||
|
||||
return result
|
||||
|
||||
mask = (ssim(x, y[..., 0:h-2, 1:w-1])>threshold) & (ssim(x, y[..., 2:h, 1:w-1])>threshold) & (ssim(x, y[..., 1:h-1, 0:w-2])>threshold) & (ssim(x, y[..., 1:h-1, 2:w])>threshold)
|
||||
mask = mask[0, 0] & mask[0, 1] & mask[0, 2]
|
||||
|
||||
return mask.cpu().numpy()
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def validate_middlebury_indetail(model, iters=32, scale_iters=8, split='H', mixed_prec=False):
|
||||
""" Peform validation using the Middlebury-V3 dataset """
|
||||
model.eval()
|
||||
aug_params = {}
|
||||
val_dataset = datasets.Middlebury(aug_params, split=split, is_eval=True)
|
||||
|
||||
out_list, epe_list, portion_list = [[], [], [], []], [[], [], [], []], [[], [], [], []]
|
||||
for val_id in range(len(val_dataset)):
|
||||
data_blob = val_dataset[val_id]
|
||||
image1 = data_blob["img1"][None].cuda()
|
||||
image2 = data_blob["img2"][None].cuda()
|
||||
disp_gt = data_blob["disp"]
|
||||
valid = data_blob["valid"]
|
||||
|
||||
padder = InputPadder(image1.shape, divis_by=32)
|
||||
image1, image2 = padder.pad(image1, image2)
|
||||
|
||||
with autocast(enabled=mixed_prec):
|
||||
disp_pr = model(image1, image2, iters=iters, scale_iters=scale_iters, test_mode=True)
|
||||
disp_pr = padder.unpad(disp_pr).cpu().squeeze(0)
|
||||
assert disp_pr.shape == disp_gt.shape, (disp_pr.shape, disp_gt.shape)
|
||||
epe = torch.sum(torch.abs(disp_pr - disp_gt), dim=0)
|
||||
|
||||
epe_flattened = epe.flatten()
|
||||
|
||||
occ_mask = Image.open(data_blob["imageL_file"].replace('im0.png', 'mask0nocc.png')).convert('L')
|
||||
occ_mask = np.ascontiguousarray(occ_mask, dtype=np.float32).flatten()
|
||||
val_all = (valid.reshape(-1) >= 0.5) & (disp_gt.reshape(-1) < 1000)
|
||||
val_occ = val_all & (occ_mask==128)
|
||||
val_nocc = val_all & (occ_mask==255)
|
||||
|
||||
val_ntt = val_all & compute_nontexture(data_blob["img1"][None].cuda(), split=split).flatten()
|
||||
|
||||
out = (epe_flattened > 2.0)
|
||||
image_out = out[val_all].float().mean().item()
|
||||
image_epe = epe_flattened[val_all].mean().item()
|
||||
|
||||
image_out_occ = out[val_occ].float().mean().item()
|
||||
image_epe_occ = epe_flattened[val_occ].mean().item()
|
||||
|
||||
image_out_nocc = out[val_nocc].float().mean().item()
|
||||
image_epe_nocc = epe_flattened[val_nocc].mean().item()
|
||||
|
||||
image_out_ntt = out[val_ntt].float().mean().item()
|
||||
image_epe_ntt = epe_flattened[val_ntt].mean().item()
|
||||
|
||||
logging.info(f"Middlebury Iter {val_id+1} out of {len(val_dataset)}. "
|
||||
f"All({round((val_all.sum()/val_all.sum()).item(),4)}): EPE {round(image_epe,4)} Out2 {round(image_out,4)}, \n "
|
||||
f"Occ({round((val_occ.sum()/val_all.sum()).item(),4)}): EPE {round(image_epe_occ,4)} Out2 {round(image_out_occ,4)}, "
|
||||
f"NOcc({round((val_nocc.sum()/val_all.sum()).item(),4)}): EPE {round(image_epe_nocc,4)} Out2 {round(image_out_nocc,4)}, "
|
||||
f"NonTexture({round((val_ntt.sum()/val_all.sum()).item(),4)}): EPE {round(image_epe_ntt,4)} Out2 {round(image_out_ntt,4)}")
|
||||
|
||||
epe_list[0].append(image_epe)
|
||||
out_list[0].append(image_out)
|
||||
portion_list[0].append((val_all.sum()/val_all.sum()).item())
|
||||
epe_list[1].append(image_epe_occ)
|
||||
out_list[1].append(image_out_occ)
|
||||
portion_list[1].append((val_occ.sum()/val_all.sum()).item())
|
||||
epe_list[2].append(image_epe_nocc)
|
||||
out_list[2].append(image_out_nocc)
|
||||
portion_list[2].append((val_nocc.sum()/val_all.sum()).item())
|
||||
epe_list[3].append(image_epe_ntt)
|
||||
out_list[3].append(image_out_ntt)
|
||||
portion_list[3].append((val_ntt.sum()/val_all.sum()).item())
|
||||
|
||||
epe_list = np.array(epe_list)
|
||||
out_list = np.array(out_list)
|
||||
portion_list = np.array(portion_list)
|
||||
|
||||
epe = np.mean(epe_list, axis=1)
|
||||
out2 = 100 * np.mean(out_list, axis=1)
|
||||
portion = 100 * np.mean(portion_list, axis=1)
|
||||
|
||||
print(f"Validation Middlebury{split}: All({round(portion[0],8)}%): EPE {round(epe[0],8)} Out2 {round(out2[0],8)}, \n"
|
||||
f"Occ({round(portion[1],8)}%): EPE {round(epe[1],8)} Out2 {round(out2[1],8)}, "
|
||||
f"NOcc({round(portion[2],8)}%): EPE {round(epe[2],8)} Out2 {round(out2[2],8)}, "
|
||||
f"NonTexture({round(portion[3],8)}%): EPE {round(epe[3],8)} Out2 {round(out2[3],8)}")
|
||||
return {f'middlebury{split}-epe': epe[0], f'middlebury{split}-out2': out2[0]}
|
||||
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--restore_ckpt', help="restore checkpoint", default=None)
|
||||
parser.add_argument('--datasets', nargs='+', type=str, help="dataset for evaluation", default=["things"],
|
||||
choices=["things", "eth3d", "kitti12", "kitti15"] + [f"middlebury_{s}" for s in 'FHQ'])
|
||||
parser.add_argument('--indetail', action='store_true', help='evaluate middlebury in detail (for different regions)')
|
||||
|
||||
parser.add_argument('--mixed_precision', action='store_true', help='use mixed precision')
|
||||
parser.add_argument('--valid_iters', type=int, default=32, help='number of disparity field updates during forward pass')
|
||||
parser.add_argument('--scale_iters', type=int, default=8, help="number of scaling updates to the disparity field in each forward pass.")
|
||||
|
||||
# Architecure choices
|
||||
parser.add_argument('--dinov2_encoder', type=str, default='vits', choices=['vits', 'vitb', 'vitl', 'vitg'])
|
||||
parser.add_argument('--idepth_scale', type=float, default=0.5, help="the scale of inverse depth to initialize disparity")
|
||||
parser.add_argument('--hidden_dims', nargs='+', type=int, default=[128]*3, help="hidden state and context dimensions")
|
||||
parser.add_argument('--corr_implementation', choices=["reg", "alt", "reg_cuda", "alt_cuda"], default="reg", help="correlation volume implementation")
|
||||
parser.add_argument('--shared_backbone', action='store_true', help="use a single backbone for the context and feature encoders")
|
||||
parser.add_argument('--corr_levels', type=int, default=2, help="number of levels in the correlation pyramid")
|
||||
parser.add_argument('--corr_radius', type=int, default=4, help="width of the correlation pyramid")
|
||||
parser.add_argument('--scale_list', type=float, nargs='+', default=[0.125, 0.25, 0.5, 0.75, 1.0, 1.25, 1.5, 2.0],
|
||||
help='the list of scaling factors of disparity')
|
||||
parser.add_argument('--scale_corr_radius', type=int, default=2,
|
||||
help="width of the correlation pyramid for scaled disparity")
|
||||
|
||||
parser.add_argument('--n_downsample', type=int, default=2, choices=[2, 3], help="resolution of the disparity field (1/2^K)")
|
||||
parser.add_argument('--context_norm', type=str, default="batch", choices=['group', 'batch', 'instance', 'none'], help="normalization of context encoder")
|
||||
parser.add_argument('--n_gru_layers', type=int, default=3, help="number of hidden GRU levels")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
model = DEFOMStereo(args)
|
||||
|
||||
logging.basicConfig(level=logging.INFO,
|
||||
format='%(asctime)s %(levelname)-8s [%(filename)s:%(lineno)d] %(message)s')
|
||||
|
||||
if args.restore_ckpt is not None:
|
||||
assert args.restore_ckpt.endswith(".pth")
|
||||
logging.info("Loading checkpoint...")
|
||||
checkpoint = torch.load(args.restore_ckpt, map_location='cuda')
|
||||
if 'model' in checkpoint:
|
||||
model.load_state_dict(checkpoint['model'])
|
||||
else:
|
||||
model.load_state_dict(checkpoint)
|
||||
logging.info(f"Done loading checkpoint")
|
||||
|
||||
model.cuda()
|
||||
model.eval()
|
||||
|
||||
print(f"The model has {format(count_parameters(model)[1]/1e6, '.2f')}M learnable parameters.")
|
||||
|
||||
# The CUDA implementations of the correlation volume prevent half-precision
|
||||
# rounding errors in the correlation lookup. This allows us to use mixed precision
|
||||
# in the entire forward pass, not just in the GRUs & feature extractors.
|
||||
use_mixed_precision = args.corr_implementation.endswith("_cuda")
|
||||
|
||||
if 'things' in args.datasets:
|
||||
validate_things(model, iters=args.valid_iters, scale_iters=args.scale_iters, mixed_prec=use_mixed_precision)
|
||||
|
||||
if 'eth3d' in args.datasets:
|
||||
validate_eth3d(model, iters=args.valid_iters, scale_iters=args.scale_iters, mixed_prec=use_mixed_precision)
|
||||
|
||||
if 'kitti12' in args.datasets:
|
||||
validate_kitti(model, iters=args.valid_iters, scale_iters=args.scale_iters, split='12', mixed_prec=use_mixed_precision)
|
||||
|
||||
if 'kitti15' in args.datasets:
|
||||
validate_kitti(model, iters=args.valid_iters, scale_iters=args.scale_iters, split='15', mixed_prec=use_mixed_precision)
|
||||
|
||||
for s in 'FHQ':
|
||||
if f"middlebury_{s}" in args.datasets:
|
||||
if args.indetail:
|
||||
validate_middlebury_indetail(model, iters=args.valid_iters, scale_iters=args.scale_iters, split=s, mixed_prec=use_mixed_precision)
|
||||
else:
|
||||
validate_middlebury(model, iters=args.valid_iters, scale_iters=args.scale_iters, split=s, mixed_prec=use_mixed_precision)
|
||||
|
||||
@@ -0,0 +1,237 @@
|
||||
from __future__ import print_function, division
|
||||
|
||||
import argparse
|
||||
import logging
|
||||
import numpy as np
|
||||
import torch
|
||||
from tqdm import tqdm
|
||||
import time
|
||||
import os
|
||||
import cv2
|
||||
import sys
|
||||
|
||||
|
||||
from core.defom_stereo import DEFOMStereo, autocast
|
||||
|
||||
import core.stereo_datasets as datasets
|
||||
from core.utils.utils import InputPadder
|
||||
from core.utils.frame_utils import writePFM
|
||||
|
||||
|
||||
def makedirs(path):
|
||||
if not os.path.exists(path):
|
||||
os.makedirs(path)
|
||||
|
||||
|
||||
def StrToBytes(text):
|
||||
if sys.version_info[0] == 2:
|
||||
return text
|
||||
else:
|
||||
return bytes(text, 'UTF-8')
|
||||
|
||||
|
||||
def count_parameters(model):
|
||||
return sum(p.numel() for p in model.parameters() if p.requires_grad)
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def test_eth3d(model, save_path, iters=32, scale_iters=3, mixed_prec=False):
|
||||
""" Peform validation using the ETH3D (train) split """
|
||||
model.eval()
|
||||
aug_params = {}
|
||||
|
||||
test_dataset = datasets.ETH3D(aug_params, split='testing', is_test=True)
|
||||
training_dataset = datasets.ETH3D(aug_params, split='training', is_test=True)
|
||||
dataset = test_dataset + training_dataset
|
||||
torch.backends.cudnn.benchmark = True
|
||||
|
||||
for test_id in tqdm(range(len(dataset))):
|
||||
img1, img2, imageL_file = dataset[test_id]
|
||||
image1 = img1[None].cuda()
|
||||
image2 = img2[None].cuda()
|
||||
|
||||
padder = InputPadder(image1.shape, divis_by=32)
|
||||
image1, image2 = padder.pad(image1, image2)
|
||||
|
||||
with autocast(enabled=mixed_prec):
|
||||
start = time.time()
|
||||
disp_pr = model(image1, image2, iters=iters, scale_iters=scale_iters, test_mode=True)
|
||||
end = time.time()
|
||||
runtime = end - start
|
||||
disp = padder.unpad(disp_pr).cpu().squeeze().numpy()
|
||||
disp[disp < 0] = 0
|
||||
disp[disp > 64] = 64
|
||||
|
||||
names = imageL_file.split("/")
|
||||
save_sub_path = os.path.join(save_path, "low_res_"+names[-3])
|
||||
makedirs(save_sub_path)
|
||||
|
||||
disp_path = os.path.join(save_sub_path, names[-2] + '.pfm')
|
||||
writePFM(disp_path, disp)
|
||||
|
||||
txt_path = os.path.join(save_sub_path, names[-2] + '.txt')
|
||||
with open(txt_path, 'wb') as time_file:
|
||||
time_file.write(StrToBytes('runtime ' + str(runtime)))
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def test_kitti(model, save_path, iters=32, scale_iters=3, split='15', mixed_prec=False):
|
||||
""" Peform testing on the KITTI-2015 (test) split """
|
||||
model.eval()
|
||||
aug_params = {}
|
||||
save_path = os.path.join(save_path, "disp_0")
|
||||
makedirs(save_path)
|
||||
|
||||
test_dataset = datasets.KITTI(aug_params, split=split, image_set='testing', is_test=True)
|
||||
|
||||
runtime_sum = 0.0
|
||||
runtime_count = 0
|
||||
|
||||
for test_id in tqdm(range(len(test_dataset))):
|
||||
img1, img2, imageL_file = test_dataset[test_id]
|
||||
image1 = img1[None].cuda()
|
||||
image2 = img2[None].cuda()
|
||||
|
||||
padder = InputPadder(image1.shape, divis_by=32)
|
||||
image1, image2 = padder.pad(image1, image2)
|
||||
|
||||
with autocast(enabled=mixed_prec):
|
||||
start = time.time()
|
||||
disp_pr = model(image1, image2, iters=iters, scale_iters=scale_iters, test_mode=True)
|
||||
end = time.time()
|
||||
runtime = end - start
|
||||
runtime_sum += runtime
|
||||
runtime_count += 1
|
||||
|
||||
disp = padder.unpad(disp_pr).cpu().squeeze().numpy()
|
||||
disp[disp < 0] = 0
|
||||
disp[disp > 240] = 240
|
||||
disp = np.uint16(disp*256)
|
||||
|
||||
name = imageL_file.split('/')[-1]
|
||||
path = os.path.join(save_path, name)
|
||||
cv2.imwrite(path, disp, [cv2.IMWRITE_PNG_COMPRESSION, 9])
|
||||
|
||||
print('The average runtime on Kitti test images is (you will need this for the submission): '
|
||||
+ str(runtime_sum / runtime_count) + " seconds")
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def test_middlebury(model, save_path, iters=32, scale_iters=8, split='F', mixed_prec=False, method_name="DEFOM-Stereo"):
|
||||
""" Peform validation using the Middlebury-V3 dataset """
|
||||
model.eval()
|
||||
aug_params = {}
|
||||
test_dataset = datasets.Middlebury(aug_params, split=split, image_set='test', is_test=True)
|
||||
training_dataset = datasets.Middlebury(aug_params, split=split, image_set='training', is_test=True)
|
||||
dataset = test_dataset + training_dataset
|
||||
torch.backends.cudnn.benchmark = True
|
||||
|
||||
for test_id in tqdm(range(len(dataset))):
|
||||
img1, img2, imageL_file = dataset[test_id]
|
||||
image1 = img1[None].cuda()
|
||||
image2 = img2[None].cuda()
|
||||
|
||||
padder = InputPadder(image1.shape, divis_by=32)
|
||||
image1, image2 = padder.pad(image1, image2)
|
||||
|
||||
with autocast(enabled=mixed_prec):
|
||||
start = time.time()
|
||||
disp_pr = model(image1, image2, iters=iters, scale_iters=scale_iters, test_mode=True)
|
||||
end = time.time()
|
||||
runtime = end - start
|
||||
disp = padder.unpad(disp_pr).cpu().squeeze().numpy()
|
||||
disp[disp < 0] = 0
|
||||
disp[disp > 800] = 800
|
||||
|
||||
names = imageL_file.split("/")
|
||||
save_sub_path = os.path.join(save_path, names[-3], names[-2])
|
||||
makedirs(save_sub_path)
|
||||
|
||||
disp_path = os.path.join(save_sub_path, 'disp0' + method_name + '.pfm')
|
||||
writePFM(disp_path, disp)
|
||||
|
||||
txt_path = os.path.join(save_sub_path, 'time' + method_name + '.txt')
|
||||
with open(txt_path, 'wb') as time_file:
|
||||
time_file.write(StrToBytes(str(runtime)))
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
parser = argparse.ArgumentParser()
|
||||
|
||||
parser.add_argument('--restore_ckpt', help="restore checkpoint", default=None)
|
||||
parser.add_argument('--datasets', nargs='+', type=str, help="dataset for evaluation", default=["kitti12", "kitti15"],
|
||||
choices=["eth3d", "kitti12", "kitti15"] + [f"middlebury_{s}" for s in 'FHQ'])
|
||||
parser.add_argument('--mixed_precision', action='store_true', help='use mixed precision')
|
||||
parser.add_argument('--valid_iters', type=int, default=32, help='number of disparity field updates during forward pass')
|
||||
parser.add_argument('--scale_iters', type=int, default=8, help="number of scaling updates to the disparity field in each forward pass.")
|
||||
parser.add_argument('--method_name', default="DEFOM-Stereo", help="the method to test")
|
||||
|
||||
# Architecure choices
|
||||
parser.add_argument('--dinov2_encoder', type=str, default='vits', choices=['vits', 'vitb', 'vitl', 'vitg'])
|
||||
parser.add_argument('--idepth_scale', type=float, default=0.5,
|
||||
help="the scale of inverse depth to initialize disparity")
|
||||
parser.add_argument('--hidden_dims', nargs='+', type=int, default=[128] * 3,
|
||||
help="hidden state and context dimensions")
|
||||
parser.add_argument('--corr_implementation', choices=["reg", "alt", "reg_cuda", "alt_cuda"], default="reg",
|
||||
help="correlation volume implementation")
|
||||
parser.add_argument('--shared_backbone', action='store_true',
|
||||
help="use a single backbone for the context and feature encoders")
|
||||
parser.add_argument('--corr_levels', type=int, default=2, help="number of levels in the correlation pyramid")
|
||||
parser.add_argument('--corr_radius', type=int, default=4, help="width of the correlation pyramid")
|
||||
parser.add_argument('--scale_list', type=float, nargs='+', default=[0.125, 0.25, 0.5, 0.75, 1.0, 1.25, 1.5, 2.0],
|
||||
help='the list of scaling factors of disparity')
|
||||
parser.add_argument('--scale_corr_radius', type=int, default=2,
|
||||
help="width of the correlation pyramid for scaled disparity")
|
||||
|
||||
parser.add_argument('--n_downsample', type=int, default=2, choices=[2, 3],
|
||||
help="resolution of the disparity field (1/2^K)")
|
||||
parser.add_argument('--context_norm', type=str, default="batch", choices=['group', 'batch', 'instance', 'none'],
|
||||
help="normalization of context encoder")
|
||||
parser.add_argument('--n_gru_layers', type=int, default=3, help="number of hidden GRU levels")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
model = DEFOMStereo(args)
|
||||
|
||||
logging.basicConfig(level=logging.INFO,
|
||||
format='%(asctime)s %(levelname)-8s [%(filename)s:%(lineno)d] %(message)s')
|
||||
|
||||
if args.restore_ckpt is not None:
|
||||
assert args.restore_ckpt.endswith(".pth")
|
||||
logging.info("Loading checkpoint...")
|
||||
checkpoint = torch.load(args.restore_ckpt, map_location='cuda')
|
||||
model.load_state_dict(checkpoint, strict=True)
|
||||
logging.info(f"Done loading checkpoint")
|
||||
|
||||
model.cuda()
|
||||
model.eval()
|
||||
|
||||
print(f"The model has {format(count_parameters(model)/1e6, '.2f')}M learnable parameters.")
|
||||
|
||||
# The CUDA implementations of the correlation volume prevent half-precision
|
||||
# rounding errors in the correlation lookup. This allows us to use mixed precision
|
||||
# in the entire forward pass, not just in the GRUs & feature extractors.
|
||||
use_mixed_precision = args.corr_implementation.endswith("_cuda")
|
||||
|
||||
if 'eth3d' in args.datasets:
|
||||
save_path = os.path.abspath(args.restore_ckpt).split('.')[0] + '_' + "eth3d"
|
||||
makedirs(save_path)
|
||||
test_eth3d(model, save_path, iters=args.valid_iters, scale_iters=args.scale_iters, mixed_prec=use_mixed_precision)
|
||||
|
||||
if 'kitti12' in args.datasets:
|
||||
save_path = os.path.abspath(args.restore_ckpt).split('.')[0] + '_' + "kitti12"
|
||||
makedirs(save_path)
|
||||
test_kitti(model, save_path, iters=args.valid_iters, scale_iters=args.scale_iters, mixed_prec=use_mixed_precision, split='12')
|
||||
|
||||
if 'kitti15' in args.datasets:
|
||||
save_path = os.path.abspath(args.restore_ckpt).split('.')[0] + '_' + "kitti15"
|
||||
makedirs(save_path)
|
||||
test_kitti(model, save_path, iters=args.valid_iters, scale_iters=args.scale_iters, mixed_prec=use_mixed_precision, split='15')
|
||||
|
||||
for s in 'FHQ':
|
||||
if f"middlebury_{s}" in args.datasets:
|
||||
save_path = os.path.abspath(args.restore_ckpt).split('.')[0] + '_' + f"middlebury_{s}"
|
||||
makedirs(save_path)
|
||||
test_middlebury(model, save_path, iters=args.valid_iters, scale_iters=args.scale_iters, split=s,
|
||||
method_name=args.method_name, mixed_prec=use_mixed_precision)
|
||||
|
||||
@@ -0,0 +1,16 @@
|
||||
--extra-index-url https://download.pytorch.org/whl/cu118
|
||||
torch==2.1.1
|
||||
torchvision==0.16.1
|
||||
xformers==0.0.23
|
||||
gradio_imageslider
|
||||
gradio==4.29.0
|
||||
matplotlib
|
||||
tensorboard
|
||||
scipy
|
||||
tqdm
|
||||
opt_einsum
|
||||
imageio
|
||||
scikit-image
|
||||
pillow
|
||||
timm
|
||||
gdown
|
||||
@@ -0,0 +1,6 @@
|
||||
cd checkpoints
|
||||
wget https://huggingface.co/depth-anything/Depth-Anything-V2-Small/resolve/main/depth_anything_v2_vits.pth
|
||||
wget https://huggingface.co/depth-anything/Depth-Anything-V2-Large/resolve/main/depth_anything_v2_vitl.pth
|
||||
cd ..
|
||||
|
||||
|
||||
@@ -0,0 +1,10 @@
|
||||
cd checkpoints
|
||||
gdown https://drive.google.com/uc?id=1XuAM4vqzura_6NKN70hMW5lFD4TafnDL
|
||||
gdown https://drive.google.com/uc?id=1FNt-SDysG5bUOmjZ91mzH2V_TXvLCvr5
|
||||
gdown https://drive.google.com/uc?id=1qyXKO-Nxq3ndl2H0deQpo6BSvwlGKYEg
|
||||
gdown https://drive.google.com/uc?id=1Dy1eGDdtkp2GQYQRTvMwR-3eAzaRCe_k
|
||||
gdown https://drive.google.com/uc?id=1duHLtUCDNIA76m6Fqwa7hv-aBMY-P3mg
|
||||
gdown https://drive.google.com/uc?id=1xEPS7gceJSFn_IHdzebBCgQaRNGwf1aG
|
||||
cd ..
|
||||
|
||||
|
||||
@@ -0,0 +1,97 @@
|
||||
#!/usr/bin/env bash
|
||||
|
||||
# evalutate on scene flow
|
||||
python evaluate_stereo.py \
|
||||
--restore_ckpt checkpoints/defomstereo_vitl_sceneflow.pth \
|
||||
--scale_iters 8 \
|
||||
--idepth_scale 0.5 \
|
||||
--corr_levels 2 \
|
||||
--corr_radius 4 \
|
||||
--scale_list 0.125 0.25 0.5 0.75 1.0 1.25 1.5 2.0 \
|
||||
--scale_corr_radius 2 \
|
||||
--datasets things \
|
||||
--dinov2_encoder vitl
|
||||
|
||||
python evaluate_stereo.py \
|
||||
--restore_ckpt checkpoints/defomstereo_vits_sceneflow.pth \
|
||||
--scale_iters 8 \
|
||||
--idepth_scale 0.5 \
|
||||
--corr_levels 2 \
|
||||
--corr_radius 4 \
|
||||
--scale_list 0.125 0.25 0.5 0.75 1.0 1.25 1.5 2.0 \
|
||||
--scale_corr_radius 2 \
|
||||
--datasets things \
|
||||
--dinov2_encoder vits
|
||||
|
||||
# evalutate on kitti12, kitti15, and eth3d
|
||||
python evaluate_stereo.py \
|
||||
--restore_ckpt checkpoints/defomstereo_vitl_sceneflow.pth \
|
||||
--scale_iters 8 \
|
||||
--idepth_scale 0.5 \
|
||||
--corr_levels 2 \
|
||||
--corr_radius 4 \
|
||||
--scale_list 0.125 0.25 0.5 0.75 1.0 1.25 1.5 2.0 \
|
||||
--scale_corr_radius 2 \
|
||||
--datasets kitti12 kitti15 eth3d \
|
||||
--dinov2_encoder vitl
|
||||
|
||||
python evaluate_stereo.py \
|
||||
--restore_ckpt checkpoints/defomstereo_vits_sceneflow.pth \
|
||||
--scale_iters 8 \
|
||||
--idepth_scale 0.5 \
|
||||
--corr_levels 2 \
|
||||
--corr_radius 4 \
|
||||
--scale_list 0.125 0.25 0.5 0.75 1.0 1.25 1.5 2.0 \
|
||||
--scale_corr_radius 2 \
|
||||
--datasets kitti12 kitti15 eth3d \
|
||||
--dinov2_encoder vits
|
||||
|
||||
|
||||
# evalutate on Middlebury; when evaluating defomstereo_vitl on Middlebury_F
|
||||
python evaluate_stereo.py \
|
||||
--restore_ckpt checkpoints/defomstereo_vitl_sceneflow.pth \
|
||||
--scale_iters 8 \
|
||||
--idepth_scale 0.5 \
|
||||
--corr_levels 2 \
|
||||
--corr_radius 4 \
|
||||
--scale_list 0.125 0.25 0.5 0.75 1.0 1.25 1.5 2.0 \
|
||||
--scale_corr_radius 2 \
|
||||
--datasets middlebury_F middlebury_H middlebury_Q \
|
||||
--dinov2_encoder vitl
|
||||
|
||||
python evaluate_stereo.py \
|
||||
--restore_ckpt checkpoints/defomstereo_vits_sceneflow.pth \
|
||||
--scale_iters 8 \
|
||||
--idepth_scale 0.5 \
|
||||
--corr_levels 2 \
|
||||
--corr_radius 4 \
|
||||
--scale_list 0.125 0.25 0.5 0.75 1.0 1.25 1.5 2.0 \
|
||||
--scale_corr_radius 2 \
|
||||
--datasets middlebury_F middlebury_H middlebury_Q \
|
||||
--dinov2_encoder vits
|
||||
|
||||
# evalutate on different region.
|
||||
python evaluate_stereo.py \
|
||||
--restore_ckpt checkpoints/defomstereo_vitl_sceneflow.pth \
|
||||
--scale_iters 8 \
|
||||
--idepth_scale 0.5 \
|
||||
--corr_levels 2 \
|
||||
--corr_radius 4 \
|
||||
--scale_list 0.125 0.25 0.5 0.75 1.0 1.25 1.5 2.0 \
|
||||
--scale_corr_radius 2 \
|
||||
--datasets middlebury_F middlebury_H middlebury_Q \
|
||||
--indetail \
|
||||
--dinov2_encoder vitl
|
||||
|
||||
python evaluate_stereo.py \
|
||||
--restore_ckpt checkpoints/defomstereo_vits_sceneflow.pth \
|
||||
--scale_iters 8 \
|
||||
--idepth_scale 0.5 \
|
||||
--corr_levels 2 \
|
||||
--corr_radius 4 \
|
||||
--scale_list 0.125 0.25 0.5 0.75 1.0 1.25 1.5 2.0 \
|
||||
--scale_corr_radius 2 \
|
||||
--datasets middlebury_F middlebury_H middlebury_Q \
|
||||
--indetail \
|
||||
--dinov2_encoder vits
|
||||
|
||||
@@ -0,0 +1,53 @@
|
||||
#!/usr/bin/env bash
|
||||
|
||||
# make submission to stereo benchmarks
|
||||
|
||||
# make submission for kitti12 and kitti15
|
||||
python make_submission.py \
|
||||
--restore_ckpt checkpoints/defomstereo_vitl_kitti.pth \
|
||||
--scale_iters 8 \
|
||||
--idepth_scale 0.5 \
|
||||
--corr_levels 2 \
|
||||
--corr_radius 4 \
|
||||
--scale_list 0.125 0.25 0.5 0.75 1.0 1.25 1.5 2.0 \
|
||||
--scale_corr_radius 2 \
|
||||
--datasets kitti12 kitti15 \
|
||||
--dinov2_encoder vitl
|
||||
|
||||
# make submission for eth3d
|
||||
python make_submission.py \
|
||||
--restore_ckpt checkpoints/defomstereo_vitl_eth3d.pth \
|
||||
--scale_iters 8 \
|
||||
--idepth_scale 0.5 \
|
||||
--corr_levels 2 \
|
||||
--corr_radius 4 \
|
||||
--scale_list 0.125 0.25 0.5 0.75 1.0 1.25 1.5 2.0 \
|
||||
--scale_corr_radius 2 \
|
||||
--datasets eth3d \
|
||||
--dinov2_encoder vitl
|
||||
|
||||
# make submission for middlebury
|
||||
python make_submission.py \
|
||||
--restore_ckpt checkpoints/defomstereo_vitl_middlebury.pth \
|
||||
--method_name DEFOM-Stereo \
|
||||
--scale_iters 8 \
|
||||
--idepth_scale 0.5 \
|
||||
--corr_levels 2 \
|
||||
--corr_radius 4 \
|
||||
--scale_list 0.125 0.25 0.5 0.75 1.0 1.25 1.5 2.0 \
|
||||
--scale_corr_radius 2 \
|
||||
--datasets middlebury_F \
|
||||
--dinov2_encoder vitl
|
||||
|
||||
# make submission for kitti15, middlebury and eth3d using the RVC model
|
||||
python make_submission.py \
|
||||
--restore_ckpt checkpoints/defomstereo_vits_rvc.pth \
|
||||
--method_name DEFOM-Stereo_RVC \
|
||||
--scale_iters 8 \
|
||||
--idepth_scale 0.5 \
|
||||
--corr_levels 2 \
|
||||
--corr_radius 4 \
|
||||
--scale_list 0.125 0.25 0.5 0.75 1.0 1.25 1.5 2.0 \
|
||||
--scale_corr_radius 2 \
|
||||
--datasets kitti15 middlebury_F eth3d \
|
||||
--dinov2_encoder vits
|
||||
@@ -0,0 +1,53 @@
|
||||
#!/usr/bin/env bash
|
||||
|
||||
# trained on 4 x 24GB 3090/4090 GPUs
|
||||
|
||||
CHECKPOINT_DIR=checkpoints/defomstereo_vitl_eth3d_pretrain && \
|
||||
mkdir -p ${CHECKPOINT_DIR} && \
|
||||
python -m torch.distributed.launch --nproc_per_node=4 --master_port=9994 train_stereo.py \
|
||||
--distributed \
|
||||
--launcher pytorch \
|
||||
--gpu_ids 0 1 2 3 \
|
||||
--name defomstereo_vitl_eth3d_pretrain \
|
||||
--batch_size 8 \
|
||||
--num_workers 8 \
|
||||
--train_datasets tartan_air sceneflow sintel_stereo eth3d instereo2k crestereo \
|
||||
--train_folds 1 1 50 1000 100 2 \
|
||||
--num_steps 300000 \
|
||||
--n_downsample 2 \
|
||||
--train_iters 18 \
|
||||
--scale_iters 8 \
|
||||
--idepth_scale 0.5 \
|
||||
--corr_levels 2 \
|
||||
--corr_radius 4 \
|
||||
--scale_list 0.125 0.25 0.5 0.75 1.0 1.25 1.5 2.0 \
|
||||
--scale_corr_radius 2 \
|
||||
--dinov2_encoder vitl \
|
||||
--image_size 384 512 \
|
||||
--resume_ckpt checkpoints/defomstereo_vitl_sceneflow.pth \
|
||||
2>&1 | tee -a ${CHECKPOINT_DIR}/train.log && \
|
||||
|
||||
CHECKPOINT_DIR=checkpoints/defomstereo_vitl_eth3d && \
|
||||
mkdir -p ${CHECKPOINT_DIR} && \
|
||||
python -m torch.distributed.launch --nproc_per_node=4 --master_port=9993 train_stereo.py \
|
||||
--distributed \
|
||||
--launcher pytorch \
|
||||
--gpu_ids 0 1 2 3 \
|
||||
--name defomstereo_vitl_eth3d \
|
||||
--batch_size 8 \
|
||||
--num_workers 8 \
|
||||
--train_datasets eth3d instereo2k crestereo \
|
||||
--train_folds 1000 10 1 \
|
||||
--num_steps 90000 \
|
||||
--n_downsample 2 \
|
||||
--train_iters 18 \
|
||||
--scale_iters 8 \
|
||||
--idepth_scale 0.5 \
|
||||
--corr_levels 2 \
|
||||
--corr_radius 4 \
|
||||
--scale_list 0.125 0.25 0.5 0.75 1.0 1.25 1.5 2.0 \
|
||||
--scale_corr_radius 2 \
|
||||
--dinov2_encoder vitl \
|
||||
--image_size 384 512 \
|
||||
--resume_ckpt checkpoints/defomstereo_vitl_eth3d_pretrain.pth \
|
||||
2>&1 | tee -a ${CHECKPOINT_DIR}/train.log
|
||||
@@ -0,0 +1,30 @@
|
||||
#!/usr/bin/env bash
|
||||
|
||||
# trained on 8 x 24GB 3090/4090 GPUs
|
||||
|
||||
CHECKPOINT_DIR=checkpoints/defomstereo_vitl_kitti && \
|
||||
mkdir -p ${CHECKPOINT_DIR} && \
|
||||
python -m torch.distributed.launch --nproc_per_node=8 --master_port=9992 train_stereo.py \
|
||||
--distributed \
|
||||
--launcher pytorch \
|
||||
--gpu_ids 0 1 2 3 4 5 6 7 \
|
||||
--name defomstereo_vitl_kitti \
|
||||
--batch_size 8 \
|
||||
--num_workers 4 \
|
||||
--train_datasets kitti12 kitti15 vkitti2 \
|
||||
--train_folds 50 50 1 \
|
||||
--num_steps 50000 \
|
||||
--n_downsample 2 \
|
||||
--train_iters 18 \
|
||||
--scale_iters 8 \
|
||||
--idepth_scale 0.5 \
|
||||
--corr_levels 2 \
|
||||
--corr_radius 4 \
|
||||
--scale_list 0.125 0.25 0.5 0.75 1.0 1.25 1.5 2.0 \
|
||||
--scale_corr_radius 2 \
|
||||
--dinov2_encoder vitl \
|
||||
--image_size 352 1216 \
|
||||
--resume_ckpt checkpoints/defomstereo_vitl_sceneflow.pth \
|
||||
2>&1 | tee -a ${CHECKPOINT_DIR}/train.log
|
||||
|
||||
|
||||
@@ -0,0 +1,53 @@
|
||||
#!/usr/bin/env bash
|
||||
|
||||
# trained on 4 x 24GB 3090/4090 GPUs
|
||||
CHECKPOINT_DIR=checkpoints/defomstereo_vitl_middlebury_pretrain && \
|
||||
mkdir -p ${CHECKPOINT_DIR} && \
|
||||
python -m torch.distributed.launch --nproc_per_node=4 --master_port=9993 train_stereo.py \
|
||||
--distributed \
|
||||
--launcher pytorch \
|
||||
--gpu_ids 0 1 2 3 \
|
||||
--name defomstereo_vitl_middlebury_pretrain \
|
||||
--batch_size 8 \
|
||||
--num_workers 8 \
|
||||
--train_datasets tartan_air sceneflow falling_things instereo2k carla_highres crestereo middlebury_2014 middlebury_2021 middlebury_H \
|
||||
--train_folds 1 1 1 50 50 1 200 200 200 \
|
||||
--num_steps 200000 \
|
||||
--n_downsample 2 \
|
||||
--train_iters 18 \
|
||||
--scale_iters 8 \
|
||||
--idepth_scale 0.5 \
|
||||
--corr_levels 2 \
|
||||
--corr_radius 4 \
|
||||
--scale_list 0.125 0.25 0.5 0.75 1.0 1.25 1.5 2.0 \
|
||||
--scale_corr_radius 2 \
|
||||
--dinov2_encoder vitl \
|
||||
--image_size 384 512 \
|
||||
--resume_ckpt checkpoints/defomstereo_vitl_sceneflow.pth \
|
||||
2>&1 | tee -a ${CHECKPOINT_DIR}/train.log && \
|
||||
|
||||
# trained on 8 x 24GB 3090/4090 GPUs
|
||||
CHECKPOINT_DIR=checkpoints/defomstereo_vitl_middlebury && \
|
||||
mkdir -p ${CHECKPOINT_DIR} && \
|
||||
python -m torch.distributed.launch --nproc_per_node=8 --master_port=9993 train_stereo.py \
|
||||
--distributed \
|
||||
--launcher pytorch \
|
||||
--gpu_ids 0 1 2 3 4 5 6 7 \
|
||||
--name defomstereo_vitl_middlebury \
|
||||
--batch_size 8 \
|
||||
--num_workers 4 \
|
||||
--train_datasets crestereo instereo2k carla_highres middlebury_2014 middlebury_2021 middlebury_H middlebury_F falling_things \
|
||||
--train_folds 1 50 50 200 200 200 200 5 \
|
||||
--num_steps 100000 \
|
||||
--n_downsample 2 \
|
||||
--train_iters 18 \
|
||||
--scale_iters 8 \
|
||||
--idepth_scale 0.5 \
|
||||
--corr_levels 2 \
|
||||
--corr_radius 4 \
|
||||
--scale_list 0.125 0.25 0.5 0.75 1.0 1.25 1.5 2.0 \
|
||||
--scale_corr_radius 2 \
|
||||
--dinov2_encoder vitl \
|
||||
--image_size 512 768 \
|
||||
--resume_ckpt checkpoints/defomstereo_vitl_middlebury_pretrain.pth \
|
||||
2>&1 | tee -a ${CHECKPOINT_DIR}/train.log
|
||||
@@ -0,0 +1,79 @@
|
||||
#!/usr/bin/env bash
|
||||
|
||||
# trained on 4 x 24GB 3090/4090 GPUs
|
||||
CHECKPOINT_DIR=checkpoints/defomstereo_vits_rvc_pretrain && \
|
||||
mkdir -p ${CHECKPOINT_DIR} && \
|
||||
python -m torch.distributed.launch --nproc_per_node=4 --master_port=9995 train_stereo.py \
|
||||
--distributed \
|
||||
--launcher pytorch \
|
||||
--gpu_ids 0 1 2 3 \
|
||||
--name defomstereo_vits_rvc_pretrain \
|
||||
--batch_size 8 \
|
||||
--num_workers 8 \
|
||||
--train_datasets tartan_air sceneflow irs 3dkenburns crestereo falling_things sintel_stereo vkitti2 carla_highres \
|
||||
--train_folds 1 1 1 1 1 1 3 3 80 \
|
||||
--num_steps 200000 \
|
||||
--n_downsample 2 \
|
||||
--train_iters 18 \
|
||||
--scale_iters 8 \
|
||||
--idepth_scale 0.5 \
|
||||
--corr_levels 2 \
|
||||
--corr_radius 4 \
|
||||
--scale_list 0.125 0.25 0.5 0.75 1.0 1.25 1.5 2.0 \
|
||||
--scale_corr_radius 2 \
|
||||
--dinov2_encoder vits \
|
||||
--image_size 384 768 \
|
||||
--resume_ckpt checkpoints/defomstereo_vits_sceneflow.pth \
|
||||
2>&1 | tee -a ${CHECKPOINT_DIR}/train.log && \
|
||||
|
||||
# trained on 4 x 24GB 3090/4090 GPUs
|
||||
CHECKPOINT_DIR=checkpoints/defomstereo_vits_rvc_pretrain2 && \
|
||||
mkdir -p ${CHECKPOINT_DIR} && \
|
||||
python -m torch.distributed.launch --nproc_per_node=4 --master_port=9996 train_stereo.py \
|
||||
--distributed \
|
||||
--launcher pytorch \
|
||||
--gpu_ids 0 1 2 3 \
|
||||
--name defomstereo_vits_rvc_pretrain2 \
|
||||
--batch_size 8 \
|
||||
--num_workers 8 \
|
||||
--train_datasets tartan_air irs 3dkenburns crestereo vkitti2 carla_highres kitti12 kitti15 middlebury_2005 middlebury_2006 middlebury_2014 middlebury_2021 middlebury_Q middlebury_H eth3d instereo2k booster \
|
||||
--train_folds 1 1 1 1 3 30 100 100 200 200 200 200 200 200 1000 20 10 \
|
||||
--num_steps 100000 \
|
||||
--n_downsample 2 \
|
||||
--train_iters 18 \
|
||||
--scale_iters 8 \
|
||||
--idepth_scale 0.5 \
|
||||
--corr_levels 2 \
|
||||
--corr_radius 4 \
|
||||
--scale_list 0.125 0.25 0.5 0.75 1.0 1.25 1.5 2.0 \
|
||||
--scale_corr_radius 2 \
|
||||
--dinov2_encoder vits \
|
||||
--image_size 384 768 \
|
||||
--resume_ckpt checkpoints/defomstereo_vits_rvc_pretrain.pth \
|
||||
2>&1 | tee -a ${CHECKPOINT_DIR}/train.log && \
|
||||
|
||||
# trained on 4 x 24GB 3090/4090 GPUs
|
||||
CHECKPOINT_DIR=checkpoints/defomstereo_vits_rvc && \
|
||||
mkdir -p ${CHECKPOINT_DIR} && \
|
||||
python -m torch.distributed.launch --nproc_per_node=4 --master_port=9997 train_stereo.py \
|
||||
--distributed \
|
||||
--launcher pytorch \
|
||||
--gpu_ids 0 1 2 3 \
|
||||
--name defomstereo_vits_rvc \
|
||||
--batch_size 8 \
|
||||
--num_workers 8 \
|
||||
--train_datasets tartan_air irs 3dkenburns crestereo vkitti2 carla_highres kitti12 kitti15 middlebury_2005 middlebury_2006 middlebury_2014 middlebury_2021 middlebury_Q middlebury_H eth3d instereo2k booster \
|
||||
--train_folds 1 1 1 1 3 30 2500 2500 200 200 200 200 200 200 1000 20 10 \
|
||||
--num_steps 20000 \
|
||||
--n_downsample 2 \
|
||||
--train_iters 18 \
|
||||
--scale_iters 8 \
|
||||
--idepth_scale 0.5 \
|
||||
--corr_levels 2 \
|
||||
--corr_radius 4 \
|
||||
--scale_list 0.125 0.25 0.5 0.75 1.0 1.25 1.5 2.0 \
|
||||
--scale_corr_radius 2 \
|
||||
--dinov2_encoder vits \
|
||||
--image_size 384 768 \
|
||||
--resume_ckpt checkpoints/defomstereo_vits_rvc_pretrain2.pth \
|
||||
2>&1 | tee -a ${CHECKPOINT_DIR}/train.log
|
||||
@@ -0,0 +1,28 @@
|
||||
#!/usr/bin/env bash
|
||||
|
||||
# trained on 4 x 24GB 3090/4090 GPUs
|
||||
|
||||
CHECKPOINT_DIR=checkpoints/defomstereo_vitl_sceneflow && \
|
||||
mkdir -p ${CHECKPOINT_DIR} && \
|
||||
python -m torch.distributed.launch --nproc_per_node=4 --master_port=9991 train_stereo.py \
|
||||
--distributed \
|
||||
--launcher pytorch \
|
||||
--gpu_ids 0 1 2 3 \
|
||||
--name defomstereo_vitl_sceneflow \
|
||||
--batch_size 8 \
|
||||
--num_workers 8 \
|
||||
--train_datasets sceneflow \
|
||||
--train_folds 1 \
|
||||
--num_steps 200000 \
|
||||
--n_downsample 2 \
|
||||
--train_iters 18 \
|
||||
--scale_iters 8 \
|
||||
--idepth_scale 0.5 \
|
||||
--corr_levels 2 \
|
||||
--corr_radius 4 \
|
||||
--scale_list 0.125 0.25 0.5 0.75 1.0 1.25 1.5 2.0 \
|
||||
--scale_corr_radius 2 \
|
||||
--dinov2_encoder vitl \
|
||||
2>&1 | tee -a ${CHECKPOINT_DIR}/train.log
|
||||
|
||||
|
||||
@@ -0,0 +1,29 @@
|
||||
#!/usr/bin/env bash
|
||||
|
||||
# trained on 2 x 24GB 3090/4090 GPUs
|
||||
|
||||
CHECKPOINT_DIR=checkpoints/defomstereo_vits_sceneflow && \
|
||||
mkdir -p ${CHECKPOINT_DIR} && \
|
||||
python -m torch.distributed.launch --nproc_per_node=2 --master_port=9990 train_stereo.py \
|
||||
--distributed \
|
||||
--launcher pytorch \
|
||||
--gpu_ids 0 1 \
|
||||
--name defomstereo_vits_sceneflow \
|
||||
--batch_size 8 \
|
||||
--num_workers 16 \
|
||||
--train_datasets sceneflow \
|
||||
--train_folds 1 \
|
||||
--num_steps 200000 \
|
||||
--mixed_precision \
|
||||
--n_downsample 2 \
|
||||
--train_iters 18 \
|
||||
--scale_iters 8 \
|
||||
--idepth_scale 0.5 \
|
||||
--corr_levels 2 \
|
||||
--corr_radius 4 \
|
||||
--scale_list 0.125 0.25 0.5 0.75 1.0 1.25 1.5 2.0 \
|
||||
--scale_corr_radius 2 \
|
||||
--dinov2_encoder vits \
|
||||
2>&1 | tee -a ${CHECKPOINT_DIR}/train.log
|
||||
|
||||
|
||||
@@ -0,0 +1,292 @@
|
||||
from __future__ import print_function, division
|
||||
import os
|
||||
import sys
|
||||
import logging
|
||||
import argparse
|
||||
import numpy as np
|
||||
from pathlib import Path
|
||||
from tqdm import tqdm
|
||||
|
||||
from torch.utils.tensorboard import SummaryWriter
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import torch.optim as optim
|
||||
from torch.utils.data import DataLoader
|
||||
|
||||
from utils.dist_utils import get_dist_info, init_dist, setup_for_distributed
|
||||
from utils.utils import *
|
||||
from core.defom_stereo import DEFOMStereo
|
||||
|
||||
from evaluate_stereo import validate_things, count_parameters
|
||||
import core.stereo_datasets as datasets
|
||||
|
||||
try:
|
||||
from torch.cuda.amp import GradScaler
|
||||
except:
|
||||
# dummy GradScaler for PyTorch < 1.6
|
||||
class GradScaler:
|
||||
def __init__(self):
|
||||
pass
|
||||
def scale(self, loss):
|
||||
return loss
|
||||
def unscale_(self, optimizer):
|
||||
pass
|
||||
def step(self, optimizer):
|
||||
optimizer.step()
|
||||
def update(self):
|
||||
pass
|
||||
|
||||
|
||||
def train(args):
|
||||
|
||||
seed_everything(args.seed)
|
||||
|
||||
if args.launcher == 'none':
|
||||
args.distributed = False
|
||||
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
||||
else:
|
||||
args.distributed = True
|
||||
|
||||
# adjust batch size for each gpu
|
||||
assert args.batch_size % torch.cuda.device_count() == 0
|
||||
args.batch_size = args.batch_size // torch.cuda.device_count()
|
||||
|
||||
dist_params = dict(backend='nccl')
|
||||
init_dist(args.launcher, **dist_params)
|
||||
# re-set gpu_ids with distributed training mode
|
||||
_, world_size = get_dist_info()
|
||||
args.gpu_ids = range(world_size)
|
||||
device = torch.device('cuda:{}'.format(args.local_rank))
|
||||
|
||||
setup_for_distributed(args.local_rank == 0)
|
||||
|
||||
model = DEFOMStereo(args).to(device)
|
||||
print("Parameter Count: %d, Trainable: %d" % count_parameters(model))
|
||||
|
||||
if args.distributed:
|
||||
process_group = torch.distributed.new_group(list(range(len(args.gpu_ids))))
|
||||
model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model, process_group)
|
||||
model = torch.nn.parallel.DistributedDataParallel(
|
||||
model.to(device),
|
||||
device_ids=[args.local_rank],
|
||||
output_device=args.local_rank,
|
||||
find_unused_parameters=True)
|
||||
model_without_ddp = model.module
|
||||
else:
|
||||
if torch.cuda.device_count() > 1:
|
||||
print('Use %d GPUs' % torch.cuda.device_count())
|
||||
model = torch.nn.DataParallel(model)
|
||||
model_without_ddp = model.module
|
||||
else:
|
||||
model_without_ddp = model
|
||||
|
||||
model_without_ddp.freeze_bn() # BatchNorm kept frozen if not distributed
|
||||
|
||||
start_epoch = 0
|
||||
start_step = 0
|
||||
optimizer, scheduler = fetch_optimizer(args, model)
|
||||
|
||||
if args.resume_ckpt:
|
||||
assert args.resume_ckpt.endswith(".pth")
|
||||
logging.info("Loading checkpoint: %s" % args.resume_ckpt)
|
||||
loc = 'cuda:{}'.format(args.local_rank) if torch.cuda.is_available() else 'cpu'
|
||||
checkpoint = torch.load(args.resume_ckpt, map_location=loc)
|
||||
if 'model' in checkpoint:
|
||||
model_without_ddp.load_state_dict(checkpoint['model'], strict=args.strict_resume)
|
||||
else:
|
||||
model_without_ddp.load_state_dict(checkpoint, strict=args.strict_resume)
|
||||
|
||||
if 'optimizer' in checkpoint and 'step' in checkpoint and 'epoch' in checkpoint and not \
|
||||
args.no_resume_optimizer:
|
||||
print('Load optimizer')
|
||||
start_step = checkpoint['step']
|
||||
start_epoch = checkpoint['epoch']
|
||||
del optimizer, scheduler
|
||||
optimizer, scheduler = fetch_optimizer(args, model, start_step, checkpoint)
|
||||
|
||||
train_data = datasets.fetch_dataset(args)
|
||||
if args.distributed:
|
||||
train_sampler = torch.utils.data.distributed.DistributedSampler(
|
||||
train_data,
|
||||
num_replicas=torch.cuda.device_count(),
|
||||
rank=args.local_rank
|
||||
)
|
||||
else:
|
||||
train_sampler = None
|
||||
train_loader = DataLoader(dataset=train_data, batch_size=args.batch_size, shuffle=train_sampler is None,
|
||||
num_workers=args.num_workers, pin_memory=True, drop_last=True,
|
||||
sampler=train_sampler)
|
||||
|
||||
total_steps = start_step
|
||||
epoch = start_epoch
|
||||
logger = Logger(model, scheduler, args.name)
|
||||
logger.total_steps = total_steps
|
||||
|
||||
model.train()
|
||||
scaler = GradScaler(enabled=args.mixed_precision)
|
||||
should_keep_training = True
|
||||
|
||||
while should_keep_training:
|
||||
|
||||
# mannually change random seed for shuffling every epoch
|
||||
if args.distributed:
|
||||
train_sampler.set_epoch(epoch)
|
||||
|
||||
if total_steps == start_step:
|
||||
epoch_start_step = start_step - len(train_loader)*start_epoch
|
||||
else:
|
||||
epoch_start_step = 0
|
||||
|
||||
for i_batch, data_blob in enumerate(tqdm(train_loader, initial=epoch_start_step)):
|
||||
optimizer.zero_grad()
|
||||
image1 = data_blob["img1"].cuda()
|
||||
image2 = data_blob["img2"].cuda()
|
||||
disp_gt = data_blob["disp"].cuda()
|
||||
valid = data_blob["valid"].cuda()
|
||||
|
||||
assert model.training
|
||||
disp_predictions = model(image1, image2, iters=args.train_iters, scale_iters=args.scale_iters)
|
||||
assert model.training
|
||||
|
||||
loss, metrics = sequence_loss(disp_predictions, disp_gt, valid)
|
||||
|
||||
scaler.scale(loss).backward()
|
||||
scaler.unscale_(optimizer)
|
||||
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
|
||||
|
||||
scaler.step(optimizer)
|
||||
scheduler.step()
|
||||
scaler.update()
|
||||
|
||||
total_steps += 1
|
||||
|
||||
if args.local_rank == 0:
|
||||
logger.writer.add_scalar("train/live_loss", loss.item(), total_steps)
|
||||
logger.writer.add_scalar(f'train/learning_rate', optimizer.param_groups[0]['lr'], total_steps)
|
||||
logger.push(metrics)
|
||||
|
||||
if total_steps % args.save_latest_ckpt_freq == 0:
|
||||
save_path = Path('checkpoints/%s/checkpoint_latest.pth' % (args.name))
|
||||
logging.info(f"Saving file {save_path.absolute()}")
|
||||
save_dict = { 'model': model_without_ddp.state_dict(),
|
||||
'optimizer': optimizer.state_dict(),
|
||||
'step': total_steps,
|
||||
'epoch': epoch}
|
||||
torch.save(save_dict, save_path)
|
||||
|
||||
if total_steps % args.save_ckpt_freq == 0:
|
||||
save_path = Path('checkpoints/%s/%s_%6d.pth' % (args.name, args.name, total_steps))
|
||||
logging.info(f"Saving file {save_path.absolute()}")
|
||||
torch.save(model_without_ddp.state_dict(), save_path)
|
||||
|
||||
if total_steps % args.val_freq == 0:
|
||||
|
||||
# visualizing training results with tensorboard
|
||||
disp = disp_predictions[-1]
|
||||
|
||||
for j in range(min(4, args.batch_size)): # write a maxmimum of four images
|
||||
logger.writer.add_image("image1/{}".format(j), image1[j].data.type(torch.uint8), total_steps)
|
||||
logger.writer.add_image("image2/{}".format(j), image2[j].data.type(torch.uint8), total_steps)
|
||||
logger.writer.add_image("disp/{}".format(j),
|
||||
(disp[j]).data.type(torch.uint8), total_steps)
|
||||
logger.writer.add_image("gt_disp/{}".format(j),
|
||||
(disp_gt[j]).data.type(torch.uint8), total_steps)
|
||||
|
||||
results = validate_things(model_without_ddp, args.valid_iters, args.scale_iters)
|
||||
logger.write_dict(results)
|
||||
model.train()
|
||||
if not args.distributed: model_without_ddp.freeze_bn()
|
||||
|
||||
if total_steps > args.num_steps:
|
||||
should_keep_training = False
|
||||
break
|
||||
|
||||
epoch += 1
|
||||
|
||||
if len(train_loader) >= 10000:
|
||||
save_path = Path('checkpoints/%s/%d_epoch_%s.pth.gz' % (args.name, total_steps, args.name))
|
||||
logging.info(f"Saving file {save_path}")
|
||||
torch.save(model_without_ddp.state_dict(), save_path)
|
||||
|
||||
print("FINISHED TRAINING")
|
||||
logger.close()
|
||||
PATH = 'checkpoints/%s.pth' % args.name
|
||||
torch.save(model_without_ddp.state_dict(), PATH)
|
||||
|
||||
return PATH
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--name', default='defom-stereo', help="name your experiment")
|
||||
|
||||
# resume pretrained model or resume training
|
||||
parser.add_argument('--resume_ckpt', default=None, type=str,
|
||||
help='resume from pretrained model or resume from unexpectedly terminated training')
|
||||
parser.add_argument('--strict_resume', action='store_true',
|
||||
help='strict resume while loading pretrained weights')
|
||||
parser.add_argument('--no_resume_optimizer', action='store_true')
|
||||
|
||||
# Training parameters
|
||||
parser.add_argument('--batch_size', type=int, default=8, help="batch size used during training.")
|
||||
parser.add_argument('--num_workers', default=8, type=int)
|
||||
parser.add_argument('--train_datasets', nargs='+', default=['sceneflow'], help="training datasets.")
|
||||
parser.add_argument('--train_folds', type=int, nargs='+', default=[1], help="training datasets' folds.")
|
||||
parser.add_argument('--lr', type=float, default=0.0002, help="max learning rate.")
|
||||
parser.add_argument('--image_size', type=int, nargs='+', default=[320, 736], help="size of the random image crops used during training.")
|
||||
parser.add_argument('--train_iters', type=int, default=18, help="number of updates to the disparity field in each forward pass.")
|
||||
parser.add_argument('--scale_iters', type=int, default=8, help="number of scaling updates to the disparity field in each forward pass.")
|
||||
parser.add_argument('--wdecay', type=float, default=.00001, help="Weight decay in optimizer.")
|
||||
parser.add_argument('--mixed_precision', action='store_true', help='use mixed precision')
|
||||
parser.add_argument('--seed', default=1234, type=int)
|
||||
|
||||
# log
|
||||
parser.add_argument('--num_steps', type=int, default=200000, help="length of training schedule.")
|
||||
parser.add_argument('--save_ckpt_freq', default=10000, type=int, help='Save checkpoint frequency (steps)')
|
||||
parser.add_argument('--save_latest_ckpt_freq', default=1000, type=int)
|
||||
parser.add_argument('--val_freq', default=10000, type=int, help='validation frequency in terms of training steps')
|
||||
|
||||
# distributed training
|
||||
parser.add_argument('--distributed', action='store_true')
|
||||
parser.add_argument('--local-rank', type=int, default=0)
|
||||
parser.add_argument('--launcher', default='none', type=str)
|
||||
parser.add_argument('--gpu_ids', default=0, type=int, nargs='+')
|
||||
|
||||
# Validation parameters
|
||||
parser.add_argument('--valid_iters', type=int, default=32, help='number of disparity field updates during validation forward pass')
|
||||
|
||||
# Raft Architecure choices
|
||||
parser.add_argument('--dinov2_encoder', type=str, default='vits', choices=['vits', 'vitb', 'vitl', 'vitg'])
|
||||
parser.add_argument('--idepth_scale', type=float, default=0.5, help="the scale of inverse depth to initialize disparity")
|
||||
parser.add_argument('--corr_implementation', choices=["reg", "alt", "reg_cuda", "alt_cuda"], default="reg", help="correlation volume implementation")
|
||||
parser.add_argument('--corr_levels', type=int, default=2, help="number of levels in the correlation pyramid")
|
||||
parser.add_argument('--corr_radius', type=int, default=4, help="width of the correlation pyramid")
|
||||
|
||||
parser.add_argument('--scale_list', type=float, nargs='+', default=[0.125, 0.25, 0.5, 0.75, 1.0, 1.25, 1.5, 2.0],
|
||||
help='the list of scaling factors of disparity')
|
||||
parser.add_argument('--scale_corr_radius', type=int, default=2, help="width of the correlation pyramid for scaled disparity")
|
||||
|
||||
parser.add_argument('--n_downsample', type=int, default=2, choices=[2, 3], help="resolution of the disparity field (1/2^K)")
|
||||
parser.add_argument('--context_norm', type=str, default="batch", choices=['group', 'batch', 'instance', 'none'], help="normalization of context encoder")
|
||||
parser.add_argument('--n_gru_layers', type=int, default=3, help="number of hidden GRU levels")
|
||||
parser.add_argument('--hidden_dims', nargs='+', type=int, default=[128]*3, help="hidden state and context dimensions")
|
||||
|
||||
# Data augmentation
|
||||
parser.add_argument('--img_gamma', type=float, nargs='+', default=None, help="gamma range")
|
||||
parser.add_argument('--saturation_range', type=float, nargs='+', default=[0.0, 1.4], help='color saturation')
|
||||
parser.add_argument('--do_flip', default='v', choices=['v', 'None'], help='flip the images vertically')
|
||||
parser.add_argument('--spatial_scale', type=float, nargs='+', default=[-0.2, 0.4], help='re-scale the images randomly')
|
||||
parser.add_argument('--noyjitter', action='store_true', help='don\'t simulate imperfect rectification')
|
||||
args = parser.parse_args()
|
||||
|
||||
logging.basicConfig(level=logging.INFO,
|
||||
stream=sys.stdout,
|
||||
format='%(asctime)s %(levelname)-8s [%(filename)s:%(lineno)d] %(message)s')
|
||||
|
||||
if 'LOCAL_RANK' not in os.environ:
|
||||
os.environ['LOCAL_RANK'] = str(args.local_rank)
|
||||
|
||||
Path("checkpoints/"+args.name).mkdir(exist_ok=True, parents=True)
|
||||
|
||||
train(args)
|
||||
@@ -0,0 +1,105 @@
|
||||
# Copyright (c) OpenMMLab. All rights reserved.
|
||||
# https://github.com/open-mmlab/mmcv/blob/7540cf73ac7e5d1e14d0ffbd9b6759e83929ecfc/mmcv/runner/dist_utils.py
|
||||
|
||||
import os
|
||||
import subprocess
|
||||
|
||||
import torch
|
||||
import torch.multiprocessing as mp
|
||||
from torch import distributed as dist
|
||||
|
||||
|
||||
def init_dist(launcher, backend='nccl', **kwargs):
|
||||
if mp.get_start_method(allow_none=True) is None:
|
||||
mp.set_start_method('spawn')
|
||||
if launcher == 'pytorch':
|
||||
_init_dist_pytorch(backend, **kwargs)
|
||||
elif launcher == 'mpi':
|
||||
_init_dist_mpi(backend, **kwargs)
|
||||
elif launcher == 'slurm':
|
||||
_init_dist_slurm(backend, **kwargs)
|
||||
else:
|
||||
raise ValueError(f'Invalid launcher type: {launcher}')
|
||||
|
||||
|
||||
def _init_dist_pytorch(backend, **kwargs):
|
||||
# TODO: use local_rank instead of rank % num_gpus
|
||||
rank = int(os.environ['RANK'])
|
||||
num_gpus = torch.cuda.device_count()
|
||||
torch.cuda.set_device(rank % num_gpus)
|
||||
dist.init_process_group(backend=backend, **kwargs)
|
||||
|
||||
|
||||
def _init_dist_mpi(backend, **kwargs):
|
||||
# TODO: use local_rank instead of rank % num_gpus
|
||||
rank = int(os.environ['OMPI_COMM_WORLD_RANK'])
|
||||
num_gpus = torch.cuda.device_count()
|
||||
torch.cuda.set_device(rank % num_gpus)
|
||||
dist.init_process_group(backend=backend, **kwargs)
|
||||
|
||||
|
||||
def _init_dist_slurm(backend, port=None):
|
||||
"""Initialize slurm distributed training environment.
|
||||
If argument ``port`` is not specified, then the master port will be system
|
||||
environment variable ``MASTER_PORT``. If ``MASTER_PORT`` is not in system
|
||||
environment variable, then a default port ``29500`` will be used.
|
||||
Args:
|
||||
backend (str): Backend of torch.distributed.
|
||||
port (int, optional): Master port. Defaults to None.
|
||||
"""
|
||||
proc_id = int(os.environ['SLURM_PROCID'])
|
||||
ntasks = int(os.environ['SLURM_NTASKS'])
|
||||
node_list = os.environ['SLURM_NODELIST']
|
||||
num_gpus = torch.cuda.device_count()
|
||||
torch.cuda.set_device(proc_id % num_gpus)
|
||||
addr = subprocess.getoutput(
|
||||
f'scontrol show hostname {node_list} | head -n1')
|
||||
# specify master port
|
||||
if port is not None:
|
||||
os.environ['MASTER_PORT'] = str(port)
|
||||
elif 'MASTER_PORT' in os.environ:
|
||||
pass # use MASTER_PORT in the environment variable
|
||||
else:
|
||||
# 29500 is torch.distributed default port
|
||||
os.environ['MASTER_PORT'] = '29500'
|
||||
# use MASTER_ADDR in the environment variable if it already exists
|
||||
if 'MASTER_ADDR' not in os.environ:
|
||||
os.environ['MASTER_ADDR'] = addr
|
||||
os.environ['WORLD_SIZE'] = str(ntasks)
|
||||
os.environ['LOCAL_RANK'] = str(proc_id % num_gpus)
|
||||
os.environ['RANK'] = str(proc_id)
|
||||
dist.init_process_group(backend=backend)
|
||||
|
||||
|
||||
def get_dist_info():
|
||||
# if (TORCH_VERSION != 'parrots'
|
||||
# and digit_version(TORCH_VERSION) < digit_version('1.0')):
|
||||
# initialized = dist._initialized
|
||||
# else:
|
||||
if dist.is_available():
|
||||
initialized = dist.is_initialized()
|
||||
else:
|
||||
initialized = False
|
||||
if initialized:
|
||||
rank = dist.get_rank()
|
||||
world_size = dist.get_world_size()
|
||||
else:
|
||||
rank = 0
|
||||
world_size = 1
|
||||
return rank, world_size
|
||||
|
||||
|
||||
# from DETR repo
|
||||
def setup_for_distributed(is_master):
|
||||
"""
|
||||
This function disables printing when not in master process
|
||||
"""
|
||||
import builtins as __builtin__
|
||||
builtin_print = __builtin__.print
|
||||
|
||||
def print(*args, **kwargs):
|
||||
force = kwargs.pop('force', False)
|
||||
if is_master or force:
|
||||
builtin_print(*args, **kwargs)
|
||||
|
||||
__builtin__.print = print
|
||||
@@ -0,0 +1,119 @@
|
||||
import logging
|
||||
import sys
|
||||
|
||||
import torch
|
||||
import torch.optim as optim
|
||||
from torch.utils.tensorboard import SummaryWriter
|
||||
|
||||
import numpy as np
|
||||
import random
|
||||
|
||||
def seed_everything(seed):
|
||||
torch.manual_seed(seed) # Current CPU
|
||||
torch.cuda.manual_seed(seed) # Current GPU
|
||||
np.random.seed(seed) # Numpy module
|
||||
random.seed(seed) # Python random module
|
||||
torch.backends.cudnn.benchmark = False # Close optimization
|
||||
torch.backends.cudnn.deterministic = True # Close optimization
|
||||
torch.cuda.manual_seed_all(seed) # All GPU (Optional)
|
||||
|
||||
|
||||
def sequence_loss(flow_preds, flow_gt, valid, loss_gamma=0.9, max_flow=700):
|
||||
""" Loss function defined over sequence of flow predictions """
|
||||
|
||||
n_predictions = len(flow_preds)
|
||||
assert n_predictions >= 1
|
||||
flow_loss = 0.0
|
||||
|
||||
# exlude invalid pixels and extremely large diplacements
|
||||
mag = torch.sum(flow_gt ** 2, dim=1, keepdim=True).sqrt()
|
||||
|
||||
# exclude extremly large displacements
|
||||
valid = ((valid >= 0.5) & (mag < max_flow))
|
||||
assert valid.shape == flow_gt.shape, [valid.shape, flow_gt.shape]
|
||||
assert not torch.isinf(flow_gt[valid.bool()]).any()
|
||||
|
||||
for i in range(n_predictions):
|
||||
assert not torch.isnan(flow_preds[i]).any() and not torch.isinf(flow_preds[i]).any()
|
||||
# We adjust the loss_gamma so it is consistent for any number of RAFT-Stereo iterations
|
||||
adjusted_loss_gamma = loss_gamma ** (15 / (n_predictions))
|
||||
i_weight = adjusted_loss_gamma ** (n_predictions - i)
|
||||
i_loss = (flow_preds[i] - flow_gt).abs()
|
||||
assert i_loss.shape == valid.shape, [i_loss.shape, valid.shape, flow_gt.shape, flow_preds[i].shape]
|
||||
flow_loss += i_weight * i_loss[valid.bool()].mean()
|
||||
|
||||
epe = torch.sum((flow_preds[-1] - flow_gt) ** 2, dim=1).sqrt()
|
||||
epe = epe.view(-1)[valid.view(-1)]
|
||||
|
||||
metrics = {
|
||||
'epe': epe.mean().item(),
|
||||
'1px': (epe < 1).float().mean().item(),
|
||||
'3px': (epe < 3).float().mean().item(),
|
||||
'5px': (epe < 5).float().mean().item(),
|
||||
}
|
||||
|
||||
return flow_loss, metrics
|
||||
|
||||
|
||||
def fetch_optimizer(args, model, last_epoch=-1, checkpoint=None):
|
||||
""" Create the optimizer and learning rate scheduler """
|
||||
trainable_params = filter(lambda p: p.requires_grad, model.parameters())
|
||||
optimizer = optim.AdamW(trainable_params, lr=args.lr, weight_decay=args.wdecay, eps=1e-8)
|
||||
if checkpoint is not None:
|
||||
optimizer.load_state_dict(checkpoint['optimizer'])
|
||||
|
||||
scheduler = optim.lr_scheduler.OneCycleLR(optimizer, args.lr, args.num_steps + 100, pct_start=0.01,
|
||||
cycle_momentum=False, anneal_strategy='linear', last_epoch=last_epoch)
|
||||
|
||||
return optimizer, scheduler
|
||||
|
||||
|
||||
class Logger:
|
||||
SUM_FREQ = 100
|
||||
|
||||
def __init__(self, model, scheduler, name):
|
||||
self.model = model
|
||||
self.scheduler = scheduler
|
||||
self.total_steps = 0
|
||||
self.running_loss = {}
|
||||
self.log_dir = 'runs/' + name
|
||||
self.writer = SummaryWriter(log_dir=self.log_dir)
|
||||
|
||||
def _print_training_status(self):
|
||||
metrics_data = [self.running_loss[k] / Logger.SUM_FREQ for k in sorted(self.running_loss.keys())]
|
||||
training_str = "[{:6d}, {:10.7f}] ".format(self.total_steps + 1, self.scheduler.get_last_lr()[0])
|
||||
metrics_str = ("{:10.4f}, " * len(metrics_data)).format(*metrics_data)
|
||||
|
||||
# print the training status
|
||||
logging.info(f"Training Metrics ({self.total_steps}): {training_str + metrics_str}")
|
||||
|
||||
if self.writer is None:
|
||||
self.writer = SummaryWriter(log_dir=self.log_dir)
|
||||
|
||||
for k in self.running_loss:
|
||||
self.writer.add_scalar("train/" + k, self.running_loss[k] / Logger.SUM_FREQ, self.total_steps)
|
||||
self.running_loss[k] = 0.0
|
||||
|
||||
def push(self, metrics):
|
||||
self.total_steps += 1
|
||||
|
||||
for key in metrics:
|
||||
if key not in self.running_loss:
|
||||
self.running_loss[key] = 0.0
|
||||
|
||||
self.running_loss[key] += metrics[key]
|
||||
|
||||
if self.total_steps % Logger.SUM_FREQ == Logger.SUM_FREQ - 1:
|
||||
self._print_training_status()
|
||||
self.running_loss = {}
|
||||
|
||||
def write_dict(self, results):
|
||||
if self.writer is None:
|
||||
self.writer = SummaryWriter(log_dir=self.log_dir)
|
||||
|
||||
for key in results:
|
||||
self.writer.add_scalar("valid/" + key, results[key], self.total_steps)
|
||||
|
||||
def close(self):
|
||||
self.writer.close()
|
||||
|
||||
Reference in New Issue
Block a user