Добавлены пропсы конвейера и стереодвижки, задействованные в прогоне
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,99 @@
|
||||
"""
|
||||
Circular-cross-section criterion (spec Category D "не подходит для сортировки без доупаковки").
|
||||
|
||||
K = r_inscribed / R_circumscribed for a cross-section outline; an object "has a circle in
|
||||
section" when K > 0.8 in ANY of its principal cross-sections. Implemented so it works on a
|
||||
(partial) point cloud measured under the dimension-estimation cameras: the circle CENTRE is
|
||||
fit to the boundary (Kasa) so an occluded-bottom arc still reads correctly, and a low
|
||||
circle-fit residual is required so an elongated rounded blob is not mistaken for a circle.
|
||||
|
||||
This module both (a) exposes `circular_section_K(points)` for the pipeline and (b) validates
|
||||
it against ground truth on every mesh, including the ones that are NOT round (boxes, etc.).
|
||||
Run inside Isaac Sim via isaacsim_send.py (reads /World/CVObjects mesh geometry).
|
||||
"""
|
||||
import omni.usd, numpy as np, json
|
||||
from pxr import UsdGeom, Usd, Gf
|
||||
stage=omni.usd.get_context().get_stage()
|
||||
K_ROUND=0.80 # spec threshold
|
||||
|
||||
def _kasa(P):
|
||||
x,y=P[:,0],P[:,1]; A=np.c_[2*x,2*y,np.ones(len(x))]; b=x*x+y*y
|
||||
s,*_=np.linalg.lstsq(A,b,rcond=None); cx,cy,cc=s; r=np.sqrt(max(cc+cx*cx+cy*cy,1e-12))
|
||||
return cx,cy,r,np.abs(np.hypot(x-cx,y-cy)-r).mean()
|
||||
|
||||
def section_K(xy):
|
||||
"""True r_inscribed/R_circumscribed of a full cross-section outline via the convex-hull
|
||||
incenter (largest inscribed circle) and the circumscribed radius from that centre.
|
||||
K=1 for a circle, b/a for an ellipse, 0.707 for a square, short/long for a rectangle."""
|
||||
from scipy.spatial import ConvexHull
|
||||
if len(xy)<20: return None,0.0,1.0
|
||||
try: h=ConvexHull(xy)
|
||||
except Exception: return None,0.0,1.0
|
||||
V=xy[h.vertices] # CCW hull vertices
|
||||
A=V; B=np.roll(V,-1,axis=0); E=B-A; L=np.linalg.norm(E,axis=1)+1e-12
|
||||
mn=xy.min(0); mx=xy.max(0)
|
||||
G=np.stack(np.meshgrid(np.linspace(mn[0],mx[0],40),np.linspace(mn[1],mx[1],40)),-1).reshape(-1,2)
|
||||
# signed distance from each grid point to each hull edge (CCW -> interior side positive)
|
||||
d=(E[:,0][None,:]*(G[:,1][:,None]-A[:,1][None,:]) - E[:,1][None,:]*(G[:,0][:,None]-A[:,0][None,:]))/L[None,:]
|
||||
inside=(d>0).all(1)
|
||||
if inside.sum()<3: return 0.0,1.0,1.0
|
||||
rin=float(d[inside].min(1).max()) # max inscribed circle radius (its own centre)
|
||||
# min enclosing circle radius (its own centre): grid centre minimising max distance to hull vertices
|
||||
Gd=np.stack(np.meshgrid(np.linspace(mn[0],mx[0],48),np.linspace(mn[1],mx[1],48)),-1).reshape(-1,2)
|
||||
Rout=float(np.linalg.norm(Gd[:,None,:]-V[None,:,:],axis=2).max(1).min())
|
||||
return rin/max(Rout,1e-9),1.0,0.0
|
||||
|
||||
def circular_section_K(points):
|
||||
"""Max K over cross-sections sampled along each principal axis (a circle in ANY section
|
||||
-> round). Returns (max_K, is_round, best_section)."""
|
||||
if len(points)<60: return 0.0,False,None
|
||||
c=points.mean(0); Q=points-c; _,_,V=np.linalg.svd(Q,full_matrices=False); proj=Q@V.T
|
||||
best=0.0; best_sec=None
|
||||
for a in range(3):
|
||||
o=[i for i in range(3) if i!=a]; ca=proj[:,a]; sp=np.ptp(ca)+1e-9
|
||||
for frac in (0.25,0.375,0.5,0.625,0.75): # sample slices along the axis
|
||||
lvl=np.percentile(ca,frac*100)
|
||||
sl=proj[np.abs(ca-lvl)<0.07*sp][:,o]
|
||||
K,cov,rr=section_K(sl)
|
||||
if K is not None and rr<0.15 and K>best:
|
||||
best=K; best_sec=(a,round(frac,2),round(cov,2),rr)
|
||||
return round(best,3), (best>K_ROUND), best_sec
|
||||
|
||||
# ---------- validate on all meshes (full GT geometry) ----------
|
||||
def mesh_points(nm, nsamp=40000):
|
||||
"""Dense, uniform surface sample (barycentric) so thin cross-sections are well populated."""
|
||||
root=stage.GetPrimAtPath(f"/World/CVObjects/{nm}"); rng=np.random.default_rng(0); out=[]
|
||||
for m in Usd.PrimRange(root):
|
||||
if m.GetTypeName()!="Mesh": continue
|
||||
P=np.array(UsdGeom.Mesh(m).GetPointsAttr().Get(),dtype=np.float64)
|
||||
idx=np.array(UsdGeom.Mesh(m).GetFaceVertexIndicesAttr().Get())
|
||||
if len(idx)%3: continue
|
||||
tris=P[idx].reshape(-1,3,3); v0,v1,v2=tris[:,0],tris[:,1],tris[:,2]
|
||||
area=0.5*np.linalg.norm(np.cross(v1-v0,v2-v0),axis=1); s=area.sum()
|
||||
if s<=0: continue
|
||||
ti=rng.choice(len(tris),nsamp,p=area/s)
|
||||
r1=np.sqrt(rng.random(nsamp)); r2=rng.random(nsamp)
|
||||
out.append(((1-r1)[:,None]*v0[ti]+(r1*(1-r2))[:,None]*v1[ti]+(r1*r2)[:,None]*v2[ti]))
|
||||
return np.concatenate(out) if out else np.zeros((0,3))
|
||||
|
||||
objs=[p.GetName() for p in stage.GetPrimAtPath("/World/CVObjects").GetChildren()]
|
||||
rows=[]; tp=fp=tn=fn=0
|
||||
for nm in sorted(objs):
|
||||
mp=stage.GetPrimAtPath(f"/World/CVObjects/{nm}/Mesh")
|
||||
gt_k=mp.GetCustomDataByKey("gt_k_round"); gt_zone=mp.GetCustomDataByKey("gt_zone")
|
||||
P=mesh_points(nm)
|
||||
maxK,is_round,sec=circular_section_K(P)
|
||||
gt_round=(gt_k is not None and gt_k>0.80)
|
||||
ok = (is_round==gt_round)
|
||||
if is_round and gt_round: tp+=1
|
||||
elif is_round and not gt_round: fp+=1
|
||||
elif not is_round and not gt_round: tn+=1
|
||||
else: fn+=1
|
||||
rows.append((nm,gt_zone,round(gt_k,3) if gt_k else None,maxK,is_round,gt_round,ok))
|
||||
print(f" {nm:16s} zone={gt_zone} GTk={gt_k:.2f} measuredK={maxK:.2f} round={str(is_round):5s} GTround={str(gt_round):5s} {'OK' if ok else 'MISS'}")
|
||||
prec=tp/(tp+fp) if tp+fp else 0; rec=tp/(tp+fn) if tp+fn else 0
|
||||
print(f"\ncircle-in-section detection vs GT k_round>0.8:")
|
||||
print(f" round: precision={prec:.2f} recall={rec:.2f} | TP={tp} FP={fp} TN={tn} FN={fn} acc={(tp+tn)/len(rows):.2f} ({tp+tn}/{len(rows)})")
|
||||
non_round=[r for r in rows if not r[5]]
|
||||
print(f" non-round objects correctly rejected: {sum(1 for r in non_round if not r[4])}/{len(non_round)} ({[r[0] for r in non_round if r[4]]} wrongly flagged)")
|
||||
json.dump([list(r) for r in rows],open("/home/dasha/isaac_assets/cv/circular_section_results.json","w"),indent=2)
|
||||
@@ -0,0 +1,281 @@
|
||||
[
|
||||
[
|
||||
"backpack",
|
||||
"C",
|
||||
0.82,
|
||||
0.795,
|
||||
false,
|
||||
true,
|
||||
false
|
||||
],
|
||||
[
|
||||
"bag",
|
||||
"D",
|
||||
0.896,
|
||||
0.889,
|
||||
true,
|
||||
true,
|
||||
true
|
||||
],
|
||||
[
|
||||
"banana",
|
||||
"D",
|
||||
0.94,
|
||||
0.917,
|
||||
true,
|
||||
true,
|
||||
true
|
||||
],
|
||||
[
|
||||
"barrel",
|
||||
"C",
|
||||
0.995,
|
||||
0.933,
|
||||
true,
|
||||
true,
|
||||
true
|
||||
],
|
||||
[
|
||||
"bolts_cluster",
|
||||
"B",
|
||||
0.718,
|
||||
0.708,
|
||||
false,
|
||||
false,
|
||||
true
|
||||
],
|
||||
[
|
||||
"bottle",
|
||||
"D",
|
||||
0.995,
|
||||
0.935,
|
||||
true,
|
||||
true,
|
||||
true
|
||||
],
|
||||
[
|
||||
"box_300x200x200",
|
||||
"B",
|
||||
0.72,
|
||||
0.685,
|
||||
false,
|
||||
false,
|
||||
true
|
||||
],
|
||||
[
|
||||
"box_400x400x300",
|
||||
"C",
|
||||
0.716,
|
||||
0.688,
|
||||
false,
|
||||
false,
|
||||
true
|
||||
],
|
||||
[
|
||||
"broom",
|
||||
"C",
|
||||
0.966,
|
||||
0.552,
|
||||
false,
|
||||
true,
|
||||
false
|
||||
],
|
||||
[
|
||||
"bucket",
|
||||
"D",
|
||||
0.995,
|
||||
0.934,
|
||||
true,
|
||||
true,
|
||||
true
|
||||
],
|
||||
[
|
||||
"chip_bag",
|
||||
"D",
|
||||
0.811,
|
||||
0.589,
|
||||
false,
|
||||
true,
|
||||
false
|
||||
],
|
||||
[
|
||||
"cone",
|
||||
"C",
|
||||
0.991,
|
||||
0.931,
|
||||
true,
|
||||
true,
|
||||
true
|
||||
],
|
||||
[
|
||||
"cylinder",
|
||||
"D",
|
||||
0.867,
|
||||
0.705,
|
||||
false,
|
||||
true,
|
||||
false
|
||||
],
|
||||
[
|
||||
"detergent",
|
||||
"B",
|
||||
0.742,
|
||||
0.684,
|
||||
false,
|
||||
false,
|
||||
true
|
||||
],
|
||||
[
|
||||
"headphones",
|
||||
"D",
|
||||
0.807,
|
||||
0.76,
|
||||
false,
|
||||
true,
|
||||
false
|
||||
],
|
||||
[
|
||||
"helmet",
|
||||
"D",
|
||||
0.895,
|
||||
0.803,
|
||||
true,
|
||||
true,
|
||||
true
|
||||
],
|
||||
[
|
||||
"lunchbox",
|
||||
"B",
|
||||
0.646,
|
||||
0.619,
|
||||
false,
|
||||
false,
|
||||
true
|
||||
],
|
||||
[
|
||||
"mug",
|
||||
"D",
|
||||
0.985,
|
||||
0.858,
|
||||
true,
|
||||
true,
|
||||
true
|
||||
],
|
||||
[
|
||||
"office_chair",
|
||||
"C",
|
||||
0.996,
|
||||
0.791,
|
||||
false,
|
||||
true,
|
||||
false
|
||||
],
|
||||
[
|
||||
"pallet",
|
||||
"C",
|
||||
0.711,
|
||||
0.685,
|
||||
false,
|
||||
false,
|
||||
true
|
||||
],
|
||||
[
|
||||
"parcel_box",
|
||||
"B",
|
||||
0.699,
|
||||
0.675,
|
||||
false,
|
||||
false,
|
||||
true
|
||||
],
|
||||
[
|
||||
"pen",
|
||||
"C",
|
||||
0.842,
|
||||
0.828,
|
||||
true,
|
||||
true,
|
||||
true
|
||||
],
|
||||
[
|
||||
"perfume",
|
||||
"D",
|
||||
0.924,
|
||||
0.878,
|
||||
true,
|
||||
true,
|
||||
true
|
||||
],
|
||||
[
|
||||
"pillow",
|
||||
"C",
|
||||
0.905,
|
||||
0.698,
|
||||
false,
|
||||
true,
|
||||
false
|
||||
],
|
||||
[
|
||||
"plate",
|
||||
"D",
|
||||
0.998,
|
||||
0.936,
|
||||
true,
|
||||
true,
|
||||
true
|
||||
],
|
||||
[
|
||||
"pouf",
|
||||
"C",
|
||||
0.994,
|
||||
0.935,
|
||||
true,
|
||||
true,
|
||||
true
|
||||
],
|
||||
[
|
||||
"sneaker",
|
||||
"B",
|
||||
0.706,
|
||||
0.69,
|
||||
false,
|
||||
false,
|
||||
true
|
||||
],
|
||||
[
|
||||
"tire",
|
||||
"C",
|
||||
0.994,
|
||||
0.933,
|
||||
true,
|
||||
true,
|
||||
true
|
||||
],
|
||||
[
|
||||
"tool_case",
|
||||
"B",
|
||||
0.454,
|
||||
0.43,
|
||||
false,
|
||||
false,
|
||||
true
|
||||
],
|
||||
[
|
||||
"umbrella",
|
||||
"C",
|
||||
0.628,
|
||||
0.879,
|
||||
true,
|
||||
false,
|
||||
false
|
||||
],
|
||||
[
|
||||
"watch",
|
||||
"C",
|
||||
0.995,
|
||||
0.668,
|
||||
false,
|
||||
true,
|
||||
false
|
||||
]
|
||||
]
|
||||
@@ -0,0 +1,2 @@
|
||||
# Auto detect text files and perform LF normalization
|
||||
* text=auto
|
||||
@@ -0,0 +1,156 @@
|
||||
# Byte-compiled / optimized / DLL files
|
||||
__pycache__/
|
||||
*.py[cod]
|
||||
*$py.class
|
||||
|
||||
# C extensions
|
||||
*.so
|
||||
|
||||
# Distribution / packaging
|
||||
.Python
|
||||
build/
|
||||
develop-eggs/
|
||||
dist/
|
||||
downloads/
|
||||
eggs/
|
||||
.eggs/
|
||||
lib/
|
||||
lib64/
|
||||
parts/
|
||||
sdist/
|
||||
var/
|
||||
wheels/
|
||||
share/python-wheels/
|
||||
*.egg-info/
|
||||
.installed.cfg
|
||||
*.egg
|
||||
MANIFEST
|
||||
|
||||
# PyInstaller
|
||||
# Usually these files are written by a python script from a template
|
||||
# before PyInstaller builds the exe, so as to inject date/other infos into it.
|
||||
*.manifest
|
||||
*.spec
|
||||
|
||||
# Installer logs
|
||||
pip-log.txt
|
||||
pip-delete-this-directory.txt
|
||||
|
||||
# Unit test / coverage reports
|
||||
htmlcov/
|
||||
.tox/
|
||||
.nox/
|
||||
.coverage
|
||||
.coverage.*
|
||||
.cache
|
||||
nosetests.xml
|
||||
coverage.xml
|
||||
*.cover
|
||||
*.py,cover
|
||||
.hypothesis/
|
||||
.pytest_cache/
|
||||
cover/
|
||||
|
||||
# Translations
|
||||
*.mo
|
||||
*.pot
|
||||
|
||||
# Django stuff:
|
||||
*.log
|
||||
local_settings.py
|
||||
db.sqlite3
|
||||
db.sqlite3-journal
|
||||
|
||||
# Flask stuff:
|
||||
instance/
|
||||
.webassets-cache
|
||||
|
||||
# Scrapy stuff:
|
||||
.scrapy
|
||||
|
||||
# Sphinx documentation
|
||||
docs/_build/
|
||||
|
||||
# PyBuilder
|
||||
.pybuilder/
|
||||
target/
|
||||
|
||||
# Jupyter Notebook
|
||||
.ipynb_checkpoints
|
||||
|
||||
# IPython
|
||||
profile_default/
|
||||
ipython_config.py
|
||||
|
||||
# pyenv
|
||||
# For a library or package, you might want to ignore these files since the code is
|
||||
# intended to run in multiple environments; otherwise, check them in:
|
||||
# .python-version
|
||||
|
||||
# pipenv
|
||||
# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
|
||||
# However, in case of collaboration, if having platform-specific dependencies or dependencies
|
||||
# having no cross-platform support, pipenv may install dependencies that don't work, or not
|
||||
# install all needed dependencies.
|
||||
#Pipfile.lock
|
||||
|
||||
# poetry
|
||||
# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
|
||||
# This is especially recommended for binary packages to ensure reproducibility, and is more
|
||||
# commonly ignored for libraries.
|
||||
# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
|
||||
#poetry.lock
|
||||
|
||||
# PEP 582; used by e.g. github.com/David-OConnor/pyflow
|
||||
__pypackages__/
|
||||
|
||||
# Celery stuff
|
||||
celerybeat-schedule
|
||||
celerybeat.pid
|
||||
|
||||
# SageMath parsed files
|
||||
*.sage.py
|
||||
|
||||
# Environments
|
||||
.env
|
||||
.venv
|
||||
env/
|
||||
venv/
|
||||
ENV/
|
||||
env.bak/
|
||||
venv.bak/
|
||||
|
||||
# Spyder project settings
|
||||
.spyderproject
|
||||
.spyproject
|
||||
|
||||
# Rope project settings
|
||||
.ropeproject
|
||||
|
||||
# mkdocs documentation
|
||||
/site
|
||||
|
||||
# mypy
|
||||
.mypy_cache/
|
||||
.dmypy.json
|
||||
dmypy.json
|
||||
|
||||
# Pyre type checker
|
||||
.pyre/
|
||||
|
||||
# pytype static type analyzer
|
||||
.pytype/
|
||||
|
||||
# Cython debug symbols
|
||||
cython_debug/
|
||||
|
||||
# PyCharm
|
||||
# JetBrains specific template is maintainted in a separate JetBrains.gitignore that can
|
||||
# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
|
||||
# and can be added to the global gitignore or merged into this file. For a more nuclear
|
||||
# option (not recommended) you can uncomment the following to ignore the entire idea folder.
|
||||
#.idea/
|
||||
|
||||
vis_results/
|
||||
models/*
|
||||
test_data/*
|
||||
@@ -0,0 +1,46 @@
|
||||
# CREStereo-Pytorch
|
||||
Non-official Pytorch implementation of the CREStereo (CVPR 2022 Oral) model converted from the original MegEngine implementation.
|
||||
|
||||
|
||||
|
||||
**update 2023/01/03**:
|
||||
- enable DistributedDataParallel (DDP) training, training time is much faster than before.
|
||||
|
||||
```shell
|
||||
# train DDP
|
||||
# change 'dist' to True in /cfgs/train.yaml file
|
||||
python -m torch.distributed.launch --nproc_per_node=8 train.py
|
||||
# train DP
|
||||
# change 'dist' to False in /cfgs/train.yaml file
|
||||
python train.py
|
||||
```
|
||||
|
||||
# Important
|
||||
- This is just an effort to try to implement the CREStereo model into Pytorch from MegEngine due to the issues of the framework to convert to other formats (https://github.com/megvii-research/CREStereo/issues/3).
|
||||
- I am not the author of the paper, and I am don't fully understand what the model is doing. Therefore, there might be small differences with the original model that might impact the performance.
|
||||
- I have not added any license, since the repository uses code from different repositories. Check the License section below for more detail.
|
||||
|
||||
# Pretrained model
|
||||
- Download the model from [here](https://drive.google.com/file/d/1D2s1v4VhJlNz98FQpFxf_kBAKQVN_7xo/view?usp=sharing) and save it into the **[models](https://github.com/ibaiGorordo/CREStereo-Pytorch/tree/main/models)** folder.
|
||||
- The model was converted from the original **[MegEngine weights](https://drive.google.com/file/d/1Wx_-zDQh7BUFBmN9im_26DFpnf3AkXj4/view)** using the `convert_weights.py` script. Place the MegEngine weights (crestereo_eth3d.mge) file into the **[models](https://github.com/ibaiGorordo/CREStereo-Pytorch/tree/main/models)** folder before the conversion.
|
||||
|
||||
# ONNX Conversion
|
||||
- After either downloading the pretrained weights or training your own model, you will have a `models/crestereo_eth3d.pth` file. If you want to run your model with ONNX, you need to run the convert_to_onnx.py script. The script has two parts:
|
||||
1. Convert the model to an ONNX model that takes in left, right images as well as an initial flow estimate (takes a few seconds)
|
||||
2. Convert the model to an ONNX model that takes in left, right images and NO initial flow estimate (takes several minutes and requires pytorch >= 1.12)
|
||||
(afaik) You will need both models to get the same results as you do from test_model.py.
|
||||
- Run the test_onnx_model.py script to verify your models work as expected!
|
||||
- NOTE: although the test_model.py script works with any size images as input, once you have converted your
|
||||
Pytorch model into ONNX models, you must provide them with the image sizes used at conversion time or it will not work.
|
||||
|
||||
# Licences:
|
||||
- CREStereo (Apache License 2.0): https://github.com/megvii-research/CREStereo/blob/master/LICENSE
|
||||
- RAFT (BSD 3-Clause):https://github.com/princeton-vl/RAFT/blob/master/LICENSE
|
||||
- LoFTR (Apache License 2.0):https://github.com/zju3dv/LoFTR/blob/master/LICENSE
|
||||
|
||||
# References:
|
||||
- CREStereo: https://github.com/megvii-research/CREStereo
|
||||
- RAFT: https://github.com/princeton-vl/RAFT
|
||||
- LoFTR: https://github.com/zju3dv/LoFTR
|
||||
- Grid sample replacement: https://zenn.dev/pinto0309/scraps/7d4032067d0160
|
||||
- torch2mge: https://github.com/MegEngine/torch2mge
|
||||
@@ -0,0 +1,20 @@
|
||||
seed: 0
|
||||
mixed_precision: false
|
||||
base_lr: 4.0e-4
|
||||
|
||||
nr_gpus: 8
|
||||
batch_size: 4
|
||||
n_total_epoch: 600
|
||||
minibatch_per_epoch: 500
|
||||
|
||||
loadmodel: ~
|
||||
log_dir: "./train_log"
|
||||
model_save_freq_epoch: 1
|
||||
|
||||
max_disp: 256
|
||||
image_width: 512
|
||||
image_height: 384
|
||||
training_data_path: "./stereo_trainset/crestereo"
|
||||
|
||||
log_level: "logging.INFO"
|
||||
dist: True # True for DDP, False for DP
|
||||
@@ -0,0 +1,49 @@
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import numpy as np
|
||||
import cv2
|
||||
from imread_from_url import imread_from_url
|
||||
|
||||
from nets import Model
|
||||
|
||||
if __name__ == '__main__':
|
||||
|
||||
model_path = "models/crestereo_eth3d.pth"
|
||||
|
||||
model = Model(max_disp=256, mixed_precision=False, test_mode=True)
|
||||
model.load_state_dict(torch.load(model_path), strict=True)
|
||||
model.eval()
|
||||
|
||||
in_h, in_w = (480, 640)
|
||||
t1_half = torch.rand(1, 3, in_h//2, in_w//2)
|
||||
t2_half = torch.rand(1, 3, in_h//2, in_w//2)
|
||||
|
||||
t1 = torch.rand(1, 3, in_h, in_w)
|
||||
t2 = torch.rand(1, 3, in_h, in_w)
|
||||
flow_init = torch.rand(1, 2, in_h//2, in_w//2)
|
||||
|
||||
# Export the model
|
||||
torch.onnx.export(model,
|
||||
(t1, t2, flow_init),
|
||||
"crestereo.onnx", # where to save the model (can be a file or file-like object)
|
||||
export_params=True, # store the trained parameter weights inside the model file
|
||||
opset_version=12, # the ONNX version to export the model to
|
||||
do_constant_folding=True, # whether to execute constant folding for optimization
|
||||
input_names = ['left', 'right','flow_init'], # the model's input names
|
||||
output_names = ['output'])
|
||||
|
||||
# Export the model without init_flow (it takes a lot of time)
|
||||
# !! Does not work prior to pytorch 1.12 (confirmed working on pytorch 2.0.0)
|
||||
# Ref: https://github.com/pytorch/pytorch/pull/73760
|
||||
torch.onnx.export(model,
|
||||
(t1_half, t2_half),
|
||||
"crestereo_without_flow.onnx", # where to save the model (can be a file or file-like object)
|
||||
export_params=True, # store the trained parameter weights inside the model file
|
||||
opset_version=12, # the ONNX version to export the model to
|
||||
do_constant_folding=True, # whether to execute constant folding for optimization
|
||||
input_names = ['left', 'right'], # the model's input names
|
||||
output_names = ['output'])
|
||||
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,26 @@
|
||||
import copy
|
||||
import torch
|
||||
import numpy as np
|
||||
import megengine as mge
|
||||
|
||||
from nets import Model
|
||||
|
||||
# Read Megengine parameters
|
||||
pretrained_dict = mge.load("models/crestereo_eth3d.mge")
|
||||
|
||||
model = Model(max_disp=256, mixed_precision=False, test_mode=True)
|
||||
model.eval()
|
||||
|
||||
state_dict = model.state_dict()
|
||||
for key, value in pretrained_dict['state_dict'].items():
|
||||
|
||||
print(f"Converting {key}")
|
||||
# Fix shape mismatch
|
||||
if value.shape[0] == 1:
|
||||
value = np.squeeze(value)
|
||||
|
||||
state_dict[key] = torch.tensor(value)
|
||||
|
||||
output_path = "models/crestereo_eth3d.pth"
|
||||
torch.save(state_dict, output_path)
|
||||
print(f"\nModel saved to: {output_path}")
|
||||
@@ -0,0 +1,215 @@
|
||||
import os
|
||||
import cv2
|
||||
import glob
|
||||
import numpy as np
|
||||
from PIL import Image, ImageEnhance
|
||||
|
||||
from torch.utils.data import Dataset
|
||||
|
||||
|
||||
class Augmentor:
|
||||
def __init__(
|
||||
self,
|
||||
image_height=384,
|
||||
image_width=512,
|
||||
max_disp=256,
|
||||
scale_min=0.6,
|
||||
scale_max=1.0,
|
||||
seed=0,
|
||||
):
|
||||
super().__init__()
|
||||
self.image_height = image_height
|
||||
self.image_width = image_width
|
||||
self.max_disp = max_disp
|
||||
self.scale_min = scale_min
|
||||
self.scale_max = scale_max
|
||||
self.rng = np.random.RandomState(seed)
|
||||
|
||||
def chromatic_augmentation(self, img):
|
||||
random_brightness = np.random.uniform(0.8, 1.2)
|
||||
random_contrast = np.random.uniform(0.8, 1.2)
|
||||
random_gamma = np.random.uniform(0.8, 1.2)
|
||||
|
||||
img = Image.fromarray(img)
|
||||
|
||||
enhancer = ImageEnhance.Brightness(img)
|
||||
img = enhancer.enhance(random_brightness)
|
||||
enhancer = ImageEnhance.Contrast(img)
|
||||
img = enhancer.enhance(random_contrast)
|
||||
|
||||
gamma_map = [
|
||||
255 * 1.0 * pow(ele / 255.0, random_gamma) for ele in range(256)
|
||||
] * 3
|
||||
img = img.point(gamma_map) # use PIL's point-function to accelerate this part
|
||||
|
||||
img_ = np.array(img)
|
||||
|
||||
return img_
|
||||
|
||||
def __call__(self, left_img, right_img, left_disp):
|
||||
# 1. chromatic augmentation
|
||||
left_img = self.chromatic_augmentation(left_img)
|
||||
right_img = self.chromatic_augmentation(right_img)
|
||||
|
||||
# 2. spatial augmentation
|
||||
# 2.1) rotate & vertical shift for right image
|
||||
if self.rng.binomial(1, 0.5):
|
||||
angle, pixel = 0.1, 2
|
||||
px = self.rng.uniform(-pixel, pixel)
|
||||
ag = self.rng.uniform(-angle, angle)
|
||||
image_center = (
|
||||
self.rng.uniform(0, right_img.shape[0]),
|
||||
self.rng.uniform(0, right_img.shape[1]),
|
||||
)
|
||||
rot_mat = cv2.getRotationMatrix2D(image_center, ag, 1.0)
|
||||
right_img = cv2.warpAffine(
|
||||
right_img, rot_mat, right_img.shape[1::-1], flags=cv2.INTER_LINEAR
|
||||
)
|
||||
trans_mat = np.float32([[1, 0, 0], [0, 1, px]])
|
||||
right_img = cv2.warpAffine(
|
||||
right_img, trans_mat, right_img.shape[1::-1], flags=cv2.INTER_LINEAR
|
||||
)
|
||||
|
||||
# 2.2) random resize
|
||||
resize_scale = self.rng.uniform(self.scale_min, self.scale_max)
|
||||
|
||||
left_img = cv2.resize(
|
||||
left_img,
|
||||
None,
|
||||
fx=resize_scale,
|
||||
fy=resize_scale,
|
||||
interpolation=cv2.INTER_LINEAR,
|
||||
)
|
||||
right_img = cv2.resize(
|
||||
right_img,
|
||||
None,
|
||||
fx=resize_scale,
|
||||
fy=resize_scale,
|
||||
interpolation=cv2.INTER_LINEAR,
|
||||
)
|
||||
|
||||
disp_mask = (left_disp < float(self.max_disp / resize_scale)) & (left_disp > 0)
|
||||
disp_mask = disp_mask.astype("float32")
|
||||
disp_mask = cv2.resize(
|
||||
disp_mask,
|
||||
None,
|
||||
fx=resize_scale,
|
||||
fy=resize_scale,
|
||||
interpolation=cv2.INTER_LINEAR,
|
||||
)
|
||||
|
||||
left_disp = (
|
||||
cv2.resize(
|
||||
left_disp,
|
||||
None,
|
||||
fx=resize_scale,
|
||||
fy=resize_scale,
|
||||
interpolation=cv2.INTER_LINEAR,
|
||||
)
|
||||
* resize_scale
|
||||
)
|
||||
|
||||
# 2.3) random crop
|
||||
h, w, c = left_img.shape
|
||||
dx = w - self.image_width
|
||||
dy = h - self.image_height
|
||||
dy = self.rng.randint(min(0, dy), max(0, dy) + 1)
|
||||
dx = self.rng.randint(min(0, dx), max(0, dx) + 1)
|
||||
|
||||
M = np.float32([[1.0, 0.0, -dx], [0.0, 1.0, -dy]])
|
||||
left_img = cv2.warpAffine(
|
||||
left_img,
|
||||
M,
|
||||
(self.image_width, self.image_height),
|
||||
flags=cv2.INTER_LINEAR,
|
||||
borderValue=0,
|
||||
)
|
||||
right_img = cv2.warpAffine(
|
||||
right_img,
|
||||
M,
|
||||
(self.image_width, self.image_height),
|
||||
flags=cv2.INTER_LINEAR,
|
||||
borderValue=0,
|
||||
)
|
||||
left_disp = cv2.warpAffine(
|
||||
left_disp,
|
||||
M,
|
||||
(self.image_width, self.image_height),
|
||||
flags=cv2.INTER_LINEAR,
|
||||
borderValue=0,
|
||||
)
|
||||
disp_mask = cv2.warpAffine(
|
||||
disp_mask,
|
||||
M,
|
||||
(self.image_width, self.image_height),
|
||||
flags=cv2.INTER_LINEAR,
|
||||
borderValue=0,
|
||||
)
|
||||
|
||||
# 3. add random occlusion to right image
|
||||
if self.rng.binomial(1, 0.5):
|
||||
sx = int(self.rng.uniform(50, 100))
|
||||
sy = int(self.rng.uniform(50, 100))
|
||||
cx = int(self.rng.uniform(sx, right_img.shape[0] - sx))
|
||||
cy = int(self.rng.uniform(sy, right_img.shape[1] - sy))
|
||||
right_img[cx - sx : cx + sx, cy - sy : cy + sy] = np.mean(
|
||||
np.mean(right_img, 0), 0
|
||||
)[np.newaxis, np.newaxis]
|
||||
|
||||
return left_img, right_img, left_disp, disp_mask
|
||||
|
||||
|
||||
class CREStereoDataset(Dataset):
|
||||
def __init__(self, root):
|
||||
super().__init__()
|
||||
self.imgs = glob.glob(os.path.join(root, "**/*_left.jpg"), recursive=True)
|
||||
self.augmentor = Augmentor(
|
||||
image_height=384,
|
||||
image_width=512,
|
||||
max_disp=256,
|
||||
scale_min=0.6,
|
||||
scale_max=1.0,
|
||||
seed=0,
|
||||
)
|
||||
self.rng = np.random.RandomState(0)
|
||||
|
||||
def get_disp(self, path):
|
||||
disp = cv2.imread(path, cv2.IMREAD_UNCHANGED)
|
||||
return disp.astype(np.float32) / 32
|
||||
|
||||
def __getitem__(self, index):
|
||||
# find path
|
||||
left_path = self.imgs[index]
|
||||
prefix = left_path[: left_path.rfind("_")]
|
||||
right_path = prefix + "_right.jpg"
|
||||
left_disp_path = prefix + "_left.disp.png"
|
||||
right_disp_path = prefix + "_right.disp.png"
|
||||
|
||||
# read img, disp
|
||||
left_img = cv2.imread(left_path, cv2.IMREAD_COLOR)
|
||||
right_img = cv2.imread(right_path, cv2.IMREAD_COLOR)
|
||||
left_disp = self.get_disp(left_disp_path)
|
||||
right_disp = self.get_disp(right_disp_path)
|
||||
|
||||
if self.rng.binomial(1, 0.5):
|
||||
left_img, right_img = np.fliplr(right_img), np.fliplr(left_img)
|
||||
left_disp, right_disp = np.fliplr(right_disp), np.fliplr(left_disp)
|
||||
left_disp[left_disp == np.inf] = 0
|
||||
|
||||
# augmentaion
|
||||
left_img, right_img, left_disp, disp_mask = self.augmentor(
|
||||
left_img, right_img, left_disp
|
||||
)
|
||||
|
||||
left_img = left_img.transpose(2, 0, 1).astype("uint8")
|
||||
right_img = right_img.transpose(2, 0, 1).astype("uint8")
|
||||
|
||||
return {
|
||||
"left": left_img,
|
||||
"right": right_img,
|
||||
"disparity": left_disp,
|
||||
"mask": disp_mask,
|
||||
}
|
||||
|
||||
def __len__(self):
|
||||
return len(self.imgs)
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 388 KiB |
@@ -0,0 +1,44 @@
|
||||
import pickle
|
||||
import numpy as np
|
||||
import megengine as mge
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
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
|
||||
ygrid = 2*ygrid/(H-1) - 1
|
||||
|
||||
grid = torch.cat([xgrid, ygrid], dim=-1)
|
||||
img = F.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 test_bilinear_sampler():
|
||||
# Getting back the megengine objects:
|
||||
with open('test_data/bilinear_sampler_test.pickle', 'rb') as f:
|
||||
right_feature_prev, coords, right_feature = pickle.load(f)
|
||||
|
||||
right_feature_prev = torch.tensor(right_feature_prev.numpy())
|
||||
coords = torch.tensor(coords.numpy())
|
||||
right_feature = right_feature.numpy()
|
||||
|
||||
# Test Pytorch
|
||||
right_feature_pytorch = bilinear_sampler(right_feature_prev, coords).numpy()
|
||||
|
||||
error = np.mean(right_feature_pytorch-right_feature)
|
||||
print(f"test_coords_grid - Avg. Error: {error}, \n \
|
||||
Original shape: {coords.numpy().shape},\n \
|
||||
Obtained shape: {right_feature_pytorch.shape}, Expected shape: {right_feature.shape}")
|
||||
|
||||
if __name__ == '__main__':
|
||||
|
||||
test_bilinear_sampler()
|
||||
@@ -0,0 +1,29 @@
|
||||
import pickle
|
||||
import numpy as np
|
||||
import megengine as mge
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
def coords_grid(batch, ht, wd, device):
|
||||
coords = torch.meshgrid(torch.arange(ht, device=device), torch.arange(wd, device=device), indexing='ij')
|
||||
coords = torch.stack(coords[::-1], dim=0).float()
|
||||
return coords[None].repeat(batch, 1, 1, 1)
|
||||
|
||||
def test_coords_grid():
|
||||
# Getting back the megengine objects:
|
||||
with open('test_data/coords_grid_test.pickle', 'rb') as f:
|
||||
batch, ht, wd, coords = pickle.load(f)
|
||||
|
||||
coords = coords.numpy()
|
||||
|
||||
# Test Pytorch
|
||||
coords_pytorch = coords_grid(batch, ht, wd, 'cpu').numpy()
|
||||
|
||||
error = np.mean(coords_pytorch-coords)
|
||||
print(f"test_coords_grid - Avg. Error: {error}, \n \
|
||||
Obtained shape: {coords_pytorch.shape}, Expected shape: {coords.shape}")
|
||||
|
||||
if __name__ == '__main__':
|
||||
|
||||
test_coords_grid()
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -0,0 +1,51 @@
|
||||
import pickle
|
||||
import numpy as np
|
||||
import megengine as mge
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
def manual_pad(x, pady, padx):
|
||||
|
||||
pad = (padx, padx, pady, pady)
|
||||
return F.pad(torch.tensor(x), pad, "replicate")
|
||||
|
||||
|
||||
def test_pad_1_1():
|
||||
# Getting back the megengine objects:
|
||||
with open('test_data/manual_pad_test1_1.pickle', 'rb') as f:
|
||||
right_feature, pady, padx, right_pad = pickle.load(f)
|
||||
|
||||
right_feature = right_feature.numpy()
|
||||
right_pad = right_pad.numpy()
|
||||
|
||||
# Test Pytorch
|
||||
right_pad_pytorch = manual_pad(right_feature, pady, padx).numpy()
|
||||
|
||||
error = np.mean(right_pad_pytorch-right_pad)
|
||||
print(f"test_pad_1_1 - Avg. Error: {error}, \n \
|
||||
Orig. shape: {right_feature.shape}, \n \
|
||||
Padded shape: {right_pad_pytorch.shape}, Expected shape: {right_pad.shape}")
|
||||
|
||||
def test_pad_0_4():
|
||||
# Getting back the megengine objects:
|
||||
with open('test_data/manual_pad_test0_4.pickle', 'rb') as f:
|
||||
right_feature, pady, padx, right_pad = pickle.load(f)
|
||||
|
||||
right_feature = right_feature.numpy()
|
||||
right_pad = right_pad.numpy()
|
||||
|
||||
# Test Pytorch
|
||||
right_pad_pytorch = manual_pad(right_feature, pady, padx).numpy()
|
||||
|
||||
error = np.mean(right_pad_pytorch-right_pad)
|
||||
print(f"test_pad_0_4 - Avg. Error: {error}, \n \
|
||||
Orig. shape: {right_feature.shape}, \n \
|
||||
Padded shape: {right_pad_pytorch.shape}, Expected shape: {right_pad.shape}")
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
|
||||
test_pad_1_1()
|
||||
|
||||
test_pad_0_4()
|
||||
@@ -0,0 +1,30 @@
|
||||
import pickle
|
||||
import numpy as np
|
||||
import megengine as mge
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
def test_meshgrid():
|
||||
# Getting back the megengine objects:
|
||||
with open('test_data/meshgrid_np_test.pkl', 'rb') as f:
|
||||
rx, dilatex, ry, dilatey, x_grid, y_grid = pickle.load(f)
|
||||
|
||||
x_grid = x_grid.numpy()
|
||||
y_grid = y_grid.numpy()
|
||||
|
||||
# Test Pytorch
|
||||
x_grid_pytorch, y_grid_pytorch = torch.meshgrid(torch.arange(-rx, rx + 1, dilatex, device='cpu'),
|
||||
torch.arange(-ry, ry + 1, dilatey, device='cpu'), indexing='xy')
|
||||
|
||||
|
||||
error_x = np.mean(x_grid_pytorch.numpy()-x_grid)
|
||||
error_y = np.mean(y_grid_pytorch.numpy()-y_grid)
|
||||
print(f"test_meshgrid (X) - Avg. Error: {error_x}, \n \
|
||||
Obtained shape: {x_grid_pytorch.numpy().shape}, Expected shape: {x_grid.shape}")
|
||||
print(f"test_meshgrid (Y) - Avg. Error: {error_y}, \n \
|
||||
Obtained shape: {y_grid_pytorch.numpy().shape}, Expected shape: {y_grid.shape}")
|
||||
|
||||
if __name__ == '__main__':
|
||||
|
||||
test_meshgrid()
|
||||
@@ -0,0 +1,31 @@
|
||||
import pickle
|
||||
import numpy as np
|
||||
import megengine as mge
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
def test_offset():
|
||||
# Getting back the megengine objects:
|
||||
with open('test_data/offset_test.pkl', 'rb') as f:
|
||||
x_grid, y_grid, reshape_shape, transpose_order, expand_size, repeat_size, repeat_axis, offsets = pickle.load(f)
|
||||
|
||||
x_grid = torch.tensor(x_grid.numpy())
|
||||
y_grid = torch.tensor(y_grid.numpy())
|
||||
offsets_mge = offsets.numpy()
|
||||
N = repeat_size
|
||||
|
||||
# Test Pytorch
|
||||
offsets = torch.stack((x_grid, y_grid))
|
||||
offsets = offsets.reshape(2, -1).permute(1, 0)
|
||||
for d in sorted((0, 2, 3)):
|
||||
offsets = offsets.unsqueeze(d)
|
||||
offsets = offsets.repeat_interleave(N, dim=0)
|
||||
|
||||
error = np.mean(offsets.numpy()-offsets_mge)
|
||||
print(f"test_offset - Avg. Error: {error}, \n \
|
||||
Obtained shape: {offsets.numpy().shape}, Expected shape: {offsets_mge.shape}")
|
||||
|
||||
if __name__ == '__main__':
|
||||
|
||||
test_offset()
|
||||
@@ -0,0 +1,47 @@
|
||||
import pickle
|
||||
import numpy as np
|
||||
import megengine as mge
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
def test_split():
|
||||
# Getting back the megengine objects:
|
||||
with open('test_data/split_test.pkl', 'rb') as f:
|
||||
left_feature, size, axis, lefts = pickle.load(f)
|
||||
|
||||
left_feature = torch.tensor(left_feature.numpy())
|
||||
|
||||
# Test Pytorch
|
||||
lefts_pytorch = torch.split(left_feature, left_feature.shape[axis]//size, dim=axis)
|
||||
|
||||
for i, (left_pytorch, left) in enumerate(zip(lefts_pytorch, lefts)):
|
||||
|
||||
error = np.mean(left_pytorch.numpy()-left.numpy())
|
||||
print(f"test_split {i} - Avg. Error: {error}, \n \
|
||||
Obtained shape: {left_pytorch.numpy().shape}, Expected shape: {left.numpy().shape}\n")
|
||||
|
||||
def test_split_list():
|
||||
# Getting back the megengine objects:
|
||||
with open('test_data/split_test_list.pkl', 'rb') as f:
|
||||
fmap1, size, axis, net, inp = pickle.load(f)
|
||||
|
||||
fmap1 = torch.tensor(fmap1.numpy())
|
||||
net = net.numpy()
|
||||
inp = inp.numpy()
|
||||
|
||||
# Test Pytorch
|
||||
net_pytorch, inp_pytorch = torch.split(fmap1, [size[0],size[0]], dim=axis)
|
||||
|
||||
error_net = np.mean(net_pytorch.numpy()-net)
|
||||
error_inp = np.mean(inp_pytorch.numpy()-inp)
|
||||
print(f"test_split_list (net) - Avg. Error: {error_net}, \n \
|
||||
Obtained shape: {net_pytorch.numpy().shape}, Expected shape: {net.shape}\n")
|
||||
print(f"test_split_list (inp) - Avg. Error: {error_inp}, \n \
|
||||
Obtained shape: {inp_pytorch.numpy().shape}, Expected shape: {inp.shape}\n")
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
|
||||
test_split()
|
||||
test_split_list()
|
||||
@@ -0,0 +1 @@
|
||||
from .crestereo import CREStereo as Model
|
||||
@@ -0,0 +1,2 @@
|
||||
from .transformer import LocalFeatureTransformer
|
||||
from .position_encoding import PositionEncodingSine
|
||||
@@ -0,0 +1,81 @@
|
||||
"""
|
||||
Linear Transformer proposed in "Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention"
|
||||
Modified from: https://github.com/idiap/fast-transformers/blob/master/fast_transformers/attention/linear_attention.py
|
||||
"""
|
||||
|
||||
import torch
|
||||
from torch.nn import Module, Dropout
|
||||
|
||||
|
||||
def elu_feature_map(x):
|
||||
return torch.nn.functional.elu(x) + 1
|
||||
|
||||
|
||||
class LinearAttention(Module):
|
||||
def __init__(self, eps=1e-6):
|
||||
super().__init__()
|
||||
self.feature_map = elu_feature_map
|
||||
self.eps = eps
|
||||
|
||||
def forward(self, queries, keys, values, q_mask=None, kv_mask=None):
|
||||
""" Multi-Head linear attention proposed in "Transformers are RNNs"
|
||||
Args:
|
||||
queries: [N, L, H, D]
|
||||
keys: [N, S, H, D]
|
||||
values: [N, S, H, D]
|
||||
q_mask: [N, L]
|
||||
kv_mask: [N, S]
|
||||
Returns:
|
||||
queried_values: (N, L, H, D)
|
||||
"""
|
||||
Q = self.feature_map(queries)
|
||||
K = self.feature_map(keys)
|
||||
|
||||
# set padded position to zero
|
||||
if q_mask is not None:
|
||||
Q = Q * q_mask[:, :, None, None]
|
||||
if kv_mask is not None:
|
||||
K = K * kv_mask[:, :, None, None]
|
||||
values = values * kv_mask[:, :, None, None]
|
||||
|
||||
v_length = values.size(1)
|
||||
values = values / v_length # prevent fp16 overflow
|
||||
KV = torch.einsum("nshd,nshv->nhdv", K, values) # (S,D)' @ S,V
|
||||
Z = 1 / (torch.einsum("nlhd,nhd->nlh", Q, K.sum(dim=1)) + self.eps)
|
||||
queried_values = torch.einsum("nlhd,nhdv,nlh->nlhv", Q, KV, Z) * v_length
|
||||
|
||||
return queried_values.contiguous()
|
||||
|
||||
|
||||
class FullAttention(Module):
|
||||
def __init__(self, use_dropout=False, attention_dropout=0.1):
|
||||
super().__init__()
|
||||
self.use_dropout = use_dropout
|
||||
self.dropout = Dropout(attention_dropout)
|
||||
|
||||
def forward(self, queries, keys, values, q_mask=None, kv_mask=None):
|
||||
""" Multi-head scaled dot-product attention, a.k.a full attention.
|
||||
Args:
|
||||
queries: [N, L, H, D]
|
||||
keys: [N, S, H, D]
|
||||
values: [N, S, H, D]
|
||||
q_mask: [N, L]
|
||||
kv_mask: [N, S]
|
||||
Returns:
|
||||
queried_values: (N, L, H, D)
|
||||
"""
|
||||
|
||||
# Compute the unnormalized attention and apply the masks
|
||||
QK = torch.einsum("nlhd,nshd->nlsh", queries, keys)
|
||||
if kv_mask is not None:
|
||||
QK.masked_fill_(~(q_mask[:, :, None, None] * kv_mask[:, None, :, None]), float('-inf'))
|
||||
|
||||
# Compute the attention and the weighted average
|
||||
softmax_temp = 1. / queries.size(3)**.5 # sqrt(D)
|
||||
A = torch.softmax(softmax_temp * QK, dim=2)
|
||||
if self.use_dropout:
|
||||
A = self.dropout(A)
|
||||
|
||||
queried_values = torch.einsum("nlsh,nshd->nlhd", A, values)
|
||||
|
||||
return queried_values.contiguous()
|
||||
@@ -0,0 +1,41 @@
|
||||
import math
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
|
||||
class PositionEncodingSine(nn.Module):
|
||||
"""
|
||||
This is a sinusoidal position encoding that generalized to 2-dimensional images
|
||||
"""
|
||||
|
||||
def __init__(self, d_model, max_shape=(256, 256), temp_bug_fix=False):
|
||||
"""
|
||||
Args:
|
||||
max_shape (tuple): for 1/8 featmap, the max length of 256 corresponds to 2048 pixels
|
||||
temp_bug_fix (bool): As noted in this [issue](https://github.com/zju3dv/LoFTR/issues/41),
|
||||
the original implementation of LoFTR includes a bug in the pos-enc impl, which has little impact
|
||||
on the final performance. For now, we keep both impls for backward compatability.
|
||||
We will remove the buggy impl after re-training all variants of our released models.
|
||||
"""
|
||||
super().__init__()
|
||||
pe = torch.zeros((d_model, *max_shape))
|
||||
y_position = torch.ones(max_shape).cumsum(0).float().unsqueeze(0)
|
||||
x_position = torch.ones(max_shape).cumsum(1).float().unsqueeze(0)
|
||||
if temp_bug_fix:
|
||||
div_term = torch.exp(torch.arange(0, d_model//2, 2).float() * (-math.log(10000.0) / (d_model//2)))
|
||||
else: # a buggy implementation (for backward compatability only)
|
||||
div_term = torch.exp(torch.arange(0, d_model//2, 2).float() * (-math.log(10000.0) / d_model//2))
|
||||
div_term = div_term[:, None, None] # [C//4, 1, 1]
|
||||
pe[0::4, :, :] = torch.sin(x_position * div_term)
|
||||
pe[1::4, :, :] = torch.cos(x_position * div_term)
|
||||
pe[2::4, :, :] = torch.sin(y_position * div_term)
|
||||
pe[3::4, :, :] = torch.cos(y_position * div_term)
|
||||
|
||||
self.register_buffer('pe', pe.unsqueeze(0), persistent=False) # [1, C, H, W]
|
||||
|
||||
def forward(self, x):
|
||||
"""
|
||||
Args:
|
||||
x: [N, C, H, W]
|
||||
"""
|
||||
return x + self.pe[:, :, :x.size(2), :x.size(3)].to(x.device)
|
||||
@@ -0,0 +1,100 @@
|
||||
import copy
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from .linear_attention import LinearAttention, FullAttention
|
||||
|
||||
#Ref: https://github.com/zju3dv/LoFTR/blob/master/src/loftr/loftr_module/transformer.py
|
||||
class LoFTREncoderLayer(nn.Module):
|
||||
def __init__(self,
|
||||
d_model,
|
||||
nhead,
|
||||
attention='linear'):
|
||||
super(LoFTREncoderLayer, self).__init__()
|
||||
|
||||
self.dim = d_model // nhead
|
||||
self.nhead = nhead
|
||||
|
||||
# multi-head attention
|
||||
self.q_proj = nn.Linear(d_model, d_model, bias=False)
|
||||
self.k_proj = nn.Linear(d_model, d_model, bias=False)
|
||||
self.v_proj = nn.Linear(d_model, d_model, bias=False)
|
||||
self.attention = LinearAttention() if attention == 'linear' else FullAttention()
|
||||
self.merge = nn.Linear(d_model, d_model, bias=False)
|
||||
|
||||
# feed-forward network
|
||||
self.mlp = nn.Sequential(
|
||||
nn.Linear(d_model*2, d_model*2, bias=False),
|
||||
nn.ReLU(),
|
||||
nn.Linear(d_model*2, d_model, bias=False),
|
||||
)
|
||||
|
||||
# norm and dropout
|
||||
self.norm1 = nn.LayerNorm(d_model)
|
||||
self.norm2 = nn.LayerNorm(d_model)
|
||||
|
||||
def forward(self, x, source, x_mask=None, source_mask=None):
|
||||
"""
|
||||
Args:
|
||||
x (torch.Tensor): [N, L, C]
|
||||
source (torch.Tensor): [N, S, C]
|
||||
x_mask (torch.Tensor): [N, L] (optional)
|
||||
source_mask (torch.Tensor): [N, S] (optional)
|
||||
"""
|
||||
bs = x.size(0)
|
||||
query, key, value = x, source, source
|
||||
|
||||
# multi-head attention
|
||||
query = self.q_proj(query).view(bs, -1, self.nhead, self.dim) # [N, L, (H, D)]
|
||||
key = self.k_proj(key).view(bs, -1, self.nhead, self.dim) # [N, S, (H, D)]
|
||||
value = self.v_proj(value).view(bs, -1, self.nhead, self.dim)
|
||||
message = self.attention(query, key, value, q_mask=x_mask, kv_mask=source_mask) # [N, L, (H, D)]
|
||||
message = self.merge(message.view(bs, -1, self.nhead*self.dim)) # [N, L, C]
|
||||
message = self.norm1(message)
|
||||
|
||||
# feed-forward network
|
||||
message = self.mlp(torch.cat([x, message], dim=2))
|
||||
message = self.norm2(message)
|
||||
|
||||
return x + message
|
||||
|
||||
|
||||
class LocalFeatureTransformer(nn.Module):
|
||||
"""A Local Feature Transformer (LoFTR) module."""
|
||||
|
||||
def __init__(self, d_model, nhead, layer_names, attention):
|
||||
super(LocalFeatureTransformer, self).__init__()
|
||||
|
||||
self.d_model = d_model
|
||||
self.nhead = nhead
|
||||
self.layer_names = layer_names
|
||||
encoder_layer = LoFTREncoderLayer(d_model, nhead, attention)
|
||||
self.layers = nn.ModuleList([copy.deepcopy(encoder_layer) for _ in range(len(self.layer_names))])
|
||||
self._reset_parameters()
|
||||
|
||||
def _reset_parameters(self):
|
||||
for p in self.parameters():
|
||||
if p.dim() > 1:
|
||||
nn.init.xavier_uniform_(p)
|
||||
|
||||
def forward(self, feat0, feat1, mask0=None, mask1=None):
|
||||
"""
|
||||
Args:
|
||||
feat0 (torch.Tensor): [N, L, C]
|
||||
feat1 (torch.Tensor): [N, S, C]
|
||||
mask0 (torch.Tensor): [N, L] (optional)
|
||||
mask1 (torch.Tensor): [N, S] (optional)
|
||||
"""
|
||||
assert self.d_model == feat0.size(2), "the feature number of src and transformer must be equal"
|
||||
|
||||
for layer, name in zip(self.layers, self.layer_names):
|
||||
|
||||
if name == 'self':
|
||||
feat0 = layer(feat0, feat0, mask0, mask0)
|
||||
feat1 = layer(feat1, feat1, mask1, mask1)
|
||||
elif name == 'cross':
|
||||
feat0 = layer(feat0, feat1, mask0, mask1)
|
||||
feat1 = layer(feat1, feat0, mask1, mask0)
|
||||
else:
|
||||
raise KeyError
|
||||
|
||||
return feat0, feat1
|
||||
@@ -0,0 +1,148 @@
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from .utils import bilinear_sampler, coords_grid, manual_pad
|
||||
|
||||
class AGCL:
|
||||
"""
|
||||
Implementation of Adaptive Group Correlation Layer (AGCL).
|
||||
"""
|
||||
|
||||
def __init__(self, fmap1, fmap2, att=None):
|
||||
self.fmap1 = fmap1
|
||||
self.fmap2 = fmap2
|
||||
|
||||
self.att = att
|
||||
|
||||
self.coords = coords_grid(fmap1.shape[0], fmap1.shape[2], fmap1.shape[3], fmap1.device)
|
||||
|
||||
def __call__(self, flow, extra_offset, small_patch=False, iter_mode=False):
|
||||
if iter_mode:
|
||||
corr = self.corr_iter(self.fmap1, self.fmap2, flow, small_patch)
|
||||
else:
|
||||
corr = self.corr_att_offset(
|
||||
self.fmap1, self.fmap2, flow, extra_offset, small_patch
|
||||
)
|
||||
return corr
|
||||
|
||||
def get_correlation(self, left_feature, right_feature, psize=(3, 3), dilate=(1, 1)):
|
||||
|
||||
N, C, H, W = left_feature.shape
|
||||
|
||||
di_y, di_x = dilate[0], dilate[1]
|
||||
pady, padx = psize[0] // 2 * di_y, psize[1] // 2 * di_x
|
||||
|
||||
right_pad = manual_pad(right_feature, pady, padx)
|
||||
|
||||
corr_list = []
|
||||
for h in range(0, pady * 2 + 1, di_y):
|
||||
for w in range(0, padx * 2 + 1, di_x):
|
||||
right_crop = right_pad[:, :, h : h + H, w : w + W]
|
||||
assert right_crop.shape == left_feature.shape
|
||||
corr = torch.mean(left_feature * right_crop, dim=1, keepdims=True)
|
||||
corr_list.append(corr)
|
||||
|
||||
corr_final = torch.cat(corr_list, dim=1)
|
||||
|
||||
return corr_final
|
||||
|
||||
def corr_iter(self, left_feature, right_feature, flow, small_patch):
|
||||
|
||||
coords = self.coords + flow
|
||||
coords = coords.permute(0, 2, 3, 1)
|
||||
right_feature = bilinear_sampler(right_feature, coords)
|
||||
|
||||
if small_patch:
|
||||
psize_list = [(3, 3), (3, 3), (3, 3), (3, 3)]
|
||||
dilate_list = [(1, 1), (1, 1), (1, 1), (1, 1)]
|
||||
else:
|
||||
psize_list = [(1, 9), (1, 9), (1, 9), (1, 9)]
|
||||
dilate_list = [(1, 1), (1, 1), (1, 1), (1, 1)]
|
||||
|
||||
N, C, H, W = left_feature.shape
|
||||
lefts = torch.split(left_feature, left_feature.shape[1]//4, dim=1)
|
||||
rights = torch.split(right_feature, right_feature.shape[1]//4, dim=1)
|
||||
|
||||
corrs = []
|
||||
for i in range(len(psize_list)):
|
||||
corr = self.get_correlation(
|
||||
lefts[i], rights[i], psize_list[i], dilate_list[i]
|
||||
)
|
||||
corrs.append(corr)
|
||||
|
||||
final_corr = torch.cat(corrs, dim=1)
|
||||
|
||||
return final_corr
|
||||
|
||||
def corr_att_offset(
|
||||
self, left_feature, right_feature, flow, extra_offset, small_patch
|
||||
):
|
||||
|
||||
N, C, H, W = left_feature.shape
|
||||
|
||||
if self.att is not None:
|
||||
left_feature = left_feature.permute(0, 2, 3, 1).reshape(N, H * W, C) # 'n c h w -> n (h w) c'
|
||||
right_feature = right_feature.permute(0, 2, 3, 1).reshape(N, H * W, C) # 'n c h w -> n (h w) c'
|
||||
# 'n (h w) c -> n c h w'
|
||||
left_feature, right_feature = self.att(left_feature, right_feature)
|
||||
# 'n (h w) c -> n c h w'
|
||||
left_feature, right_feature = [
|
||||
x.reshape(N, H, W, C).permute(0, 3, 1, 2)
|
||||
for x in [left_feature, right_feature]
|
||||
]
|
||||
|
||||
lefts = torch.split(left_feature, left_feature.shape[1]//4, dim=1)
|
||||
rights = torch.split(right_feature, right_feature.shape[1]//4, dim=1)
|
||||
|
||||
C = C // 4
|
||||
|
||||
if small_patch:
|
||||
psize_list = [(3, 3), (3, 3), (3, 3), (3, 3)]
|
||||
dilate_list = [(1, 1), (1, 1), (1, 1), (1, 1)]
|
||||
else:
|
||||
psize_list = [(1, 9), (1, 9), (1, 9), (1, 9)]
|
||||
dilate_list = [(1, 1), (1, 1), (1, 1), (1, 1)]
|
||||
|
||||
search_num = 9
|
||||
extra_offset = extra_offset.reshape(N, search_num, 2, H, W).permute(0, 1, 3, 4, 2) # [N, search_num, 1, 1, 2]
|
||||
|
||||
corrs = []
|
||||
for i in range(len(psize_list)):
|
||||
left_feature, right_feature = lefts[i], rights[i]
|
||||
psize, dilate = psize_list[i], dilate_list[i]
|
||||
|
||||
psizey, psizex = psize[0], psize[1]
|
||||
dilatey, dilatex = dilate[0], dilate[1]
|
||||
|
||||
ry = psizey // 2 * dilatey
|
||||
rx = psizex // 2 * dilatex
|
||||
x_grid, y_grid = torch.meshgrid(torch.arange(-rx, rx + 1, dilatex, device=self.fmap1.device),
|
||||
torch.arange(-ry, ry + 1, dilatey, device=self.fmap1.device), indexing='xy')
|
||||
|
||||
offsets = torch.stack((x_grid, y_grid))
|
||||
offsets = offsets.reshape(2, -1).permute(1, 0)
|
||||
for d in sorted((0, 2, 3)):
|
||||
offsets = offsets.unsqueeze(d)
|
||||
offsets = offsets.repeat_interleave(N, dim=0)
|
||||
offsets = offsets + extra_offset
|
||||
|
||||
coords = self.coords + flow # [N, 2, H, W]
|
||||
coords = coords.permute(0, 2, 3, 1) # [N, H, W, 2]
|
||||
coords = torch.unsqueeze(coords, 1) + offsets
|
||||
coords = coords.reshape(N, -1, W, 2) # [N, search_num*H, W, 2]
|
||||
|
||||
right_feature = bilinear_sampler(
|
||||
right_feature, coords
|
||||
) # [N, C, search_num*H, W]
|
||||
right_feature = right_feature.reshape(N, C, -1, H, W) # [N, C, search_num, H, W]
|
||||
left_feature = left_feature.unsqueeze(2).repeat_interleave(right_feature.shape[2], dim=2)
|
||||
|
||||
corr = torch.mean(left_feature * right_feature, dim=1)
|
||||
|
||||
corrs.append(corr)
|
||||
|
||||
final_corr = torch.cat(corrs, dim=1)
|
||||
|
||||
return final_corr
|
||||
@@ -0,0 +1,258 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from .update import BasicUpdateBlock
|
||||
from .extractor import BasicEncoder
|
||||
from .corr import AGCL
|
||||
|
||||
from .attention import PositionEncodingSine, LocalFeatureTransformer
|
||||
|
||||
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
|
||||
|
||||
#Ref: https://github.com/princeton-vl/RAFT/blob/master/core/raft.py
|
||||
class CREStereo(nn.Module):
|
||||
def __init__(self, max_disp=192, mixed_precision=False, test_mode=False):
|
||||
super(CREStereo, self).__init__()
|
||||
|
||||
self.max_flow = max_disp
|
||||
self.mixed_precision = mixed_precision
|
||||
self.test_mode = test_mode
|
||||
|
||||
self.hidden_dim = 128
|
||||
self.context_dim = 128
|
||||
self.dropout = 0
|
||||
|
||||
self.fnet = BasicEncoder(output_dim=256, norm_fn='instance', dropout=self.dropout)
|
||||
self.update_block = BasicUpdateBlock(hidden_dim=self.hidden_dim, cor_planes=4 * 9, mask_size=4)
|
||||
|
||||
# loftr
|
||||
self.self_att_fn = LocalFeatureTransformer(
|
||||
d_model=256, nhead=8, layer_names=["self"] * 1, attention="linear"
|
||||
)
|
||||
self.cross_att_fn = LocalFeatureTransformer(
|
||||
d_model=256, nhead=8, layer_names=["cross"] * 1, attention="linear"
|
||||
)
|
||||
|
||||
# adaptive search
|
||||
self.search_num = 9
|
||||
self.conv_offset_16 = nn.Conv2d(
|
||||
256, self.search_num * 2, kernel_size=3, stride=1, padding=1
|
||||
)
|
||||
self.conv_offset_8 = nn.Conv2d(
|
||||
256, self.search_num * 2, kernel_size=3, stride=1, padding=1
|
||||
)
|
||||
self.range_16 = 1
|
||||
self.range_8 = 1
|
||||
|
||||
def freeze_bn(self):
|
||||
for m in self.modules():
|
||||
if isinstance(m, nn.BatchNorm2d):
|
||||
m.eval()
|
||||
|
||||
def convex_upsample(self, flow, mask, rate=4):
|
||||
""" Upsample flow field [H/8, W/8, 2] -> [H, W, 2] using convex combination """
|
||||
N, _, H, W = flow.shape
|
||||
# print(flow.shape, mask.shape, rate)
|
||||
mask = mask.view(N, 1, 9, rate, rate, H, W)
|
||||
mask = torch.softmax(mask, dim=2)
|
||||
|
||||
up_flow = F.unfold(rate * flow, [3,3], padding=1)
|
||||
up_flow = up_flow.view(N, 2, 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, 2, rate*H, rate*W)
|
||||
|
||||
def zero_init(self, fmap):
|
||||
N, C, H, W = fmap.shape
|
||||
_x = torch.zeros([N, 1, H, W], dtype=torch.float32)
|
||||
_y = torch.zeros([N, 1, H, W], dtype=torch.float32)
|
||||
zero_flow = torch.cat((_x, _y), dim=1).to(fmap.device)
|
||||
return zero_flow
|
||||
|
||||
def forward(self, image1, image2, flow_init=None, iters=10, upsample=True, test_mode=False):
|
||||
""" Estimate optical flow between pair of frames """
|
||||
|
||||
image1 = 2 * (image1 / 255.0) - 1.0
|
||||
image2 = 2 * (image2 / 255.0) - 1.0
|
||||
|
||||
image1 = image1.contiguous()
|
||||
image2 = image2.contiguous()
|
||||
|
||||
hdim = self.hidden_dim
|
||||
cdim = self.context_dim
|
||||
|
||||
# run the feature network
|
||||
with autocast(enabled=self.mixed_precision):
|
||||
fmap1, fmap2 = self.fnet([image1, image2])
|
||||
|
||||
fmap1 = fmap1.float()
|
||||
fmap2 = fmap2.float()
|
||||
|
||||
with autocast(enabled=self.mixed_precision):
|
||||
|
||||
# 1/4 -> 1/8
|
||||
# feature
|
||||
fmap1_dw8 = F.avg_pool2d(fmap1, 2, stride=2)
|
||||
fmap2_dw8 = F.avg_pool2d(fmap2, 2, stride=2)
|
||||
|
||||
# offset
|
||||
offset_dw8 = self.conv_offset_8(fmap1_dw8)
|
||||
offset_dw8 = self.range_8 * (torch.sigmoid(offset_dw8) - 0.5) * 2.0
|
||||
|
||||
# context
|
||||
net, inp = torch.split(fmap1, [hdim,hdim], dim=1)
|
||||
net = torch.tanh(net)
|
||||
inp = F.relu(inp)
|
||||
net_dw8 = F.avg_pool2d(net, 2, stride=2)
|
||||
inp_dw8 = F.avg_pool2d(inp, 2, stride=2)
|
||||
|
||||
# 1/4 -> 1/16
|
||||
# feature
|
||||
fmap1_dw16 = F.avg_pool2d(fmap1, 4, stride=4)
|
||||
fmap2_dw16 = F.avg_pool2d(fmap2, 4, stride=4)
|
||||
offset_dw16 = self.conv_offset_16(fmap1_dw16)
|
||||
offset_dw16 = self.range_16 * (torch.sigmoid(offset_dw16) - 0.5) * 2.0
|
||||
|
||||
# context
|
||||
net_dw16 = F.avg_pool2d(net, 4, stride=4)
|
||||
inp_dw16 = F.avg_pool2d(inp, 4, stride=4)
|
||||
|
||||
# positional encoding and self-attention
|
||||
pos_encoding_fn_small = PositionEncodingSine(
|
||||
d_model=256, max_shape=(image1.shape[2] // 16, image1.shape[3] // 16)
|
||||
)
|
||||
# 'n c h w -> n (h w) c'
|
||||
x_tmp = pos_encoding_fn_small(fmap1_dw16)
|
||||
fmap1_dw16 = x_tmp.permute(0, 2, 3, 1).reshape(x_tmp.shape[0], x_tmp.shape[2] * x_tmp.shape[3], x_tmp.shape[1])
|
||||
# 'n c h w -> n (h w) c'
|
||||
x_tmp = pos_encoding_fn_small(fmap2_dw16)
|
||||
fmap2_dw16 = x_tmp.permute(0, 2, 3, 1).reshape(x_tmp.shape[0], x_tmp.shape[2] * x_tmp.shape[3], x_tmp.shape[1])
|
||||
|
||||
fmap1_dw16, fmap2_dw16 = self.self_att_fn(fmap1_dw16, fmap2_dw16)
|
||||
fmap1_dw16, fmap2_dw16 = [
|
||||
x.reshape(x.shape[0], image1.shape[2] // 16, -1, x.shape[2]).permute(0, 3, 1, 2)
|
||||
for x in [fmap1_dw16, fmap2_dw16]
|
||||
]
|
||||
|
||||
corr_fn = AGCL(fmap1, fmap2)
|
||||
corr_fn_dw8 = AGCL(fmap1_dw8, fmap2_dw8)
|
||||
corr_fn_att_dw16 = AGCL(fmap1_dw16, fmap2_dw16, att=self.cross_att_fn)
|
||||
|
||||
# Cascaded refinement (1/16 + 1/8 + 1/4)
|
||||
predictions = []
|
||||
flow = None
|
||||
flow_up = None
|
||||
if flow_init is not None:
|
||||
scale = fmap1.shape[2] / flow_init.shape[2]
|
||||
flow = -scale * F.interpolate(
|
||||
flow_init,
|
||||
size=(fmap1.shape[2], fmap1.shape[3]),
|
||||
mode="bilinear",
|
||||
align_corners=True,
|
||||
)
|
||||
else:
|
||||
# zero initialization
|
||||
flow_dw16 = self.zero_init(fmap1_dw16)
|
||||
|
||||
# Recurrent Update Module
|
||||
# RUM: 1/16
|
||||
for itr in range(iters // 2):
|
||||
if itr % 2 == 0:
|
||||
small_patch = False
|
||||
else:
|
||||
small_patch = True
|
||||
|
||||
flow_dw16 = flow_dw16.detach()
|
||||
out_corrs = corr_fn_att_dw16(
|
||||
flow_dw16, offset_dw16, small_patch=small_patch
|
||||
)
|
||||
|
||||
with autocast(enabled=self.mixed_precision):
|
||||
net_dw16, up_mask, delta_flow = self.update_block(
|
||||
net_dw16, inp_dw16, out_corrs, flow_dw16
|
||||
)
|
||||
|
||||
flow_dw16 = flow_dw16 + delta_flow
|
||||
flow = self.convex_upsample(flow_dw16, up_mask, rate=4)
|
||||
flow_up = -4 * F.interpolate(
|
||||
flow,
|
||||
size=(4 * flow.shape[2], 4 * flow.shape[3]),
|
||||
mode="bilinear",
|
||||
align_corners=True,
|
||||
)
|
||||
predictions.append(flow_up)
|
||||
|
||||
scale = fmap1_dw8.shape[2] / flow.shape[2]
|
||||
flow_dw8 = -scale * F.interpolate(
|
||||
flow,
|
||||
size=(fmap1_dw8.shape[2], fmap1_dw8.shape[3]),
|
||||
mode="bilinear",
|
||||
align_corners=True,
|
||||
)
|
||||
|
||||
# RUM: 1/8
|
||||
for itr in range(iters // 2):
|
||||
if itr % 2 == 0:
|
||||
small_patch = False
|
||||
else:
|
||||
small_patch = True
|
||||
|
||||
flow_dw8 = flow_dw8.detach()
|
||||
out_corrs = corr_fn_dw8(flow_dw8, offset_dw8, small_patch=small_patch)
|
||||
|
||||
with autocast(enabled=self.mixed_precision):
|
||||
net_dw8, up_mask, delta_flow = self.update_block(
|
||||
net_dw8, inp_dw8, out_corrs, flow_dw8
|
||||
)
|
||||
|
||||
flow_dw8 = flow_dw8 + delta_flow
|
||||
flow = self.convex_upsample(flow_dw8, up_mask, rate=4)
|
||||
flow_up = -2 * F.interpolate(
|
||||
flow,
|
||||
size=(2 * flow.shape[2], 2 * flow.shape[3]),
|
||||
mode="bilinear",
|
||||
align_corners=True,
|
||||
)
|
||||
predictions.append(flow_up)
|
||||
|
||||
scale = fmap1.shape[2] / flow.shape[2]
|
||||
flow = -scale * F.interpolate(
|
||||
flow,
|
||||
size=(fmap1.shape[2], fmap1.shape[3]),
|
||||
mode="bilinear",
|
||||
align_corners=True,
|
||||
)
|
||||
|
||||
# RUM: 1/4
|
||||
for itr in range(iters):
|
||||
if itr % 2 == 0:
|
||||
small_patch = False
|
||||
else:
|
||||
small_patch = True
|
||||
|
||||
flow = flow.detach()
|
||||
out_corrs = corr_fn(flow, None, small_patch=small_patch, iter_mode=True)
|
||||
|
||||
with autocast(enabled=self.mixed_precision):
|
||||
net, up_mask, delta_flow = self.update_block(net, inp, out_corrs, flow)
|
||||
|
||||
flow = flow + delta_flow
|
||||
flow_up = -self.convex_upsample(flow, up_mask, rate=4)
|
||||
predictions.append(flow_up)
|
||||
|
||||
if self.test_mode:
|
||||
return flow_up
|
||||
|
||||
return predictions
|
||||
@@ -0,0 +1,123 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
# Ref: https://github.com/princeton-vl/RAFT/blob/master/core/extractor.py
|
||||
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)
|
||||
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)
|
||||
self.norm3 = nn.BatchNorm2d(planes)
|
||||
|
||||
elif norm_fn == 'instance':
|
||||
self.norm1 = nn.InstanceNorm2d(planes, affine=False)
|
||||
self.norm2 = nn.InstanceNorm2d(planes, affine=False)
|
||||
self.norm3 = nn.InstanceNorm2d(planes, affine=False)
|
||||
|
||||
elif norm_fn == 'none':
|
||||
self.norm1 = nn.Sequential()
|
||||
self.norm2 = nn.Sequential()
|
||||
self.norm3 = nn.Sequential()
|
||||
|
||||
self.downsample = nn.Sequential(
|
||||
nn.Conv2d(in_planes, planes, kernel_size=1, stride=stride), self.norm3)
|
||||
|
||||
|
||||
def forward(self, x):
|
||||
y = x
|
||||
y = self.relu(self.norm1(self.conv1(y)))
|
||||
y = self.relu(self.norm2(self.conv2(y)))
|
||||
|
||||
x = self.downsample(x)
|
||||
|
||||
return self.relu(x+y)
|
||||
|
||||
|
||||
class BasicEncoder(nn.Module):
|
||||
def __init__(self, output_dim=128, norm_fn='batch', dropout=0.0):
|
||||
super(BasicEncoder, self).__init__()
|
||||
self.norm_fn = norm_fn
|
||||
|
||||
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, affine=False)
|
||||
|
||||
elif self.norm_fn == 'none':
|
||||
self.norm1 = nn.Sequential()
|
||||
|
||||
self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=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=2)
|
||||
self.layer3 = self._make_layer(128, stride=1)
|
||||
|
||||
# output convolution
|
||||
self.conv2 = nn.Conv2d(128, output_dim, kernel_size=1)
|
||||
|
||||
self.dropout = None
|
||||
if dropout > 0:
|
||||
self.dropout = nn.Dropout2d(p=dropout)
|
||||
|
||||
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):
|
||||
|
||||
# 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)
|
||||
|
||||
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 = self.conv2(x)
|
||||
|
||||
if self.dropout is not None:
|
||||
x = self.dropout(x)
|
||||
|
||||
if is_list:
|
||||
x = torch.split(x, x.shape[0]//2, dim=0)
|
||||
|
||||
return x
|
||||
@@ -0,0 +1,91 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
#Ref: https://github.com/princeton-vl/RAFT/blob/master/core/update.py
|
||||
class FlowHead(nn.Module):
|
||||
def __init__(self, input_dim=128, hidden_dim=256):
|
||||
super(FlowHead, self).__init__()
|
||||
self.conv1 = nn.Conv2d(input_dim, hidden_dim, 3, padding=1)
|
||||
self.conv2 = nn.Conv2d(hidden_dim, 2, 3, padding=1)
|
||||
self.relu = nn.ReLU(inplace=True)
|
||||
|
||||
def forward(self, x):
|
||||
return self.conv2(self.relu(self.conv1(x)))
|
||||
|
||||
|
||||
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
|
||||
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):
|
||||
super(BasicMotionEncoder, self).__init__()
|
||||
|
||||
self.convc1 = nn.Conv2d(cor_planes, 256, 1, padding=0)
|
||||
self.convc2 = nn.Conv2d(256, 192, 3, padding=1)
|
||||
self.convf1 = nn.Conv2d(2, 128, 7, padding=3)
|
||||
self.convf2 = nn.Conv2d(128, 64, 3, padding=1)
|
||||
self.conv = nn.Conv2d(64+192, 128-2, 3, padding=1)
|
||||
|
||||
def forward(self, flow, corr):
|
||||
cor = F.relu(self.convc1(corr))
|
||||
cor = F.relu(self.convc2(cor))
|
||||
flo = F.relu(self.convf1(flow))
|
||||
flo = F.relu(self.convf2(flo))
|
||||
|
||||
cor_flo = torch.cat([cor, flo], dim=1)
|
||||
out = F.relu(self.conv(cor_flo))
|
||||
return torch.cat([out, flow], dim=1)
|
||||
|
||||
|
||||
class BasicUpdateBlock(nn.Module):
|
||||
def __init__(self, hidden_dim, cor_planes, mask_size=8):
|
||||
super(BasicUpdateBlock, self).__init__()
|
||||
|
||||
self.encoder = BasicMotionEncoder(cor_planes)
|
||||
self.gru = SepConvGRU(hidden_dim=hidden_dim, input_dim=128+hidden_dim)
|
||||
self.flow_head = FlowHead(hidden_dim, hidden_dim=256)
|
||||
|
||||
self.mask = nn.Sequential(
|
||||
nn.Conv2d(128, 256, 3, padding=1),
|
||||
nn.ReLU(inplace=True),
|
||||
nn.Conv2d(256, mask_size**2 *9, 1, padding=0))
|
||||
|
||||
def forward(self, net, inp, corr, flow, upsample=True):
|
||||
# print(inp.shape, corr.shape, flow.shape)
|
||||
motion_features = self.encoder(flow, corr)
|
||||
# print(motion_features.shape, inp.shape)
|
||||
inp = torch.cat((inp, motion_features), dim=1)
|
||||
|
||||
net = self.gru(net, inp)
|
||||
delta_flow = self.flow_head(net)
|
||||
|
||||
# scale mask to balence gradients
|
||||
mask = .25 * self.mask(net)
|
||||
return net, mask, delta_flow
|
||||
@@ -0,0 +1 @@
|
||||
from .utils import bilinear_sampler, coords_grid, manual_pad
|
||||
@@ -0,0 +1,108 @@
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import numpy as np
|
||||
|
||||
#Ref: https://github.com/princeton-vl/RAFT/blob/master/core/utils/utils.py
|
||||
|
||||
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
|
||||
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, device):
|
||||
coords = torch.meshgrid(torch.arange(ht, device=device), torch.arange(wd, device=device), indexing='ij')
|
||||
coords = torch.stack(coords[::-1], dim=0).float()
|
||||
return coords[None].repeat(batch, 1, 1, 1)
|
||||
|
||||
def manual_pad(x, pady, padx):
|
||||
|
||||
pad = (padx, padx, pady, pady)
|
||||
return F.pad(x.clone().detach(), pad, "replicate")
|
||||
|
||||
# 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)
|
||||
@@ -0,0 +1,82 @@
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import numpy as np
|
||||
import cv2
|
||||
from imread_from_url import imread_from_url
|
||||
|
||||
from nets import Model
|
||||
|
||||
device = 'cuda'
|
||||
|
||||
#Ref: https://github.com/megvii-research/CREStereo/blob/master/test.py
|
||||
def inference(left, right, model, n_iter=20):
|
||||
|
||||
print("Model Forwarding...")
|
||||
imgL = left.transpose(2, 0, 1)
|
||||
imgR = right.transpose(2, 0, 1)
|
||||
imgL = np.ascontiguousarray(imgL[None, :, :, :])
|
||||
imgR = np.ascontiguousarray(imgR[None, :, :, :])
|
||||
|
||||
imgL = torch.tensor(imgL.astype("float32")).to(device)
|
||||
imgR = torch.tensor(imgR.astype("float32")).to(device)
|
||||
|
||||
imgL_dw2 = F.interpolate(
|
||||
imgL,
|
||||
size=(imgL.shape[2] // 2, imgL.shape[3] // 2),
|
||||
mode="bilinear",
|
||||
align_corners=True,
|
||||
)
|
||||
imgR_dw2 = F.interpolate(
|
||||
imgR,
|
||||
size=(imgL.shape[2] // 2, imgL.shape[3] // 2),
|
||||
mode="bilinear",
|
||||
align_corners=True,
|
||||
)
|
||||
# print(imgR_dw2.shape)
|
||||
with torch.inference_mode():
|
||||
pred_flow_dw2 = model(imgL_dw2, imgR_dw2, iters=n_iter, flow_init=None)
|
||||
|
||||
pred_flow = model(imgL, imgR, iters=n_iter, flow_init=pred_flow_dw2)
|
||||
pred_disp = torch.squeeze(pred_flow[:, 0, :, :]).cpu().detach().numpy()
|
||||
|
||||
return pred_disp
|
||||
|
||||
if __name__ == '__main__':
|
||||
|
||||
left_img = imread_from_url("https://raw.githubusercontent.com/megvii-research/CREStereo/master/img/test/left.png")
|
||||
right_img = imread_from_url("https://raw.githubusercontent.com/megvii-research/CREStereo/master/img/test/right.png")
|
||||
|
||||
in_h, in_w = left_img.shape[:2]
|
||||
|
||||
# Resize image in case the GPU memory overflows
|
||||
eval_h, eval_w = (in_h,in_w)
|
||||
assert eval_h%8 == 0, "input height should be divisible by 8"
|
||||
assert eval_w%8 == 0, "input width should be divisible by 8"
|
||||
|
||||
imgL = cv2.resize(left_img, (eval_w, eval_h), interpolation=cv2.INTER_LINEAR)
|
||||
imgR = cv2.resize(right_img, (eval_w, eval_h), interpolation=cv2.INTER_LINEAR)
|
||||
|
||||
model_path = "models/crestereo_eth3d.pth"
|
||||
|
||||
model = Model(max_disp=256, mixed_precision=False, test_mode=True)
|
||||
model.load_state_dict(torch.load(model_path), strict=True)
|
||||
model.to(device)
|
||||
model.eval()
|
||||
|
||||
pred = inference(imgL, imgR, model, n_iter=20)
|
||||
|
||||
t = float(in_w) / float(eval_w)
|
||||
disp = cv2.resize(pred, (in_w, in_h), interpolation=cv2.INTER_LINEAR) * t
|
||||
|
||||
disp_vis = (disp - disp.min()) / (disp.max() - disp.min()) * 255.0
|
||||
disp_vis = disp_vis.astype("uint8")
|
||||
disp_vis = cv2.applyColorMap(disp_vis, cv2.COLORMAP_INFERNO)
|
||||
|
||||
combined_img = np.hstack((left_img, disp_vis))
|
||||
cv2.namedWindow("output", cv2.WINDOW_NORMAL)
|
||||
cv2.imshow("output", combined_img)
|
||||
cv2.imwrite("output.jpg", disp_vis)
|
||||
cv2.waitKey(0)
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,78 @@
|
||||
import numpy as np
|
||||
import cv2
|
||||
|
||||
import onnxruntime
|
||||
|
||||
# Ref: https://github.com/megvii-research/CREStereo/blob/master/test.py
|
||||
def inference(left, right, model, no_flow_model):
|
||||
# Get onnx model layer names (see convert_to_onnx.py for what these are)
|
||||
input1_name = model.get_inputs()[0].name
|
||||
input2_name = model.get_inputs()[1].name
|
||||
input3_name = model.get_inputs()[2].name
|
||||
output_name = model.get_outputs()[0].name
|
||||
|
||||
# Decimate the image to half the original size for flow estimation network
|
||||
imgL_dw2 = cv2.resize(
|
||||
left, (left.shape[1] // 2, left.shape[0] // 2), interpolation=cv2.INTER_LINEAR)
|
||||
imgR_dw2 = cv2.resize(
|
||||
right, (right.shape[1] // 2, right.shape[0] // 2), interpolation=cv2.INTER_LINEAR)
|
||||
|
||||
# Reshape inputs to match what is expected
|
||||
imgL = left.transpose(2, 0, 1)
|
||||
imgR = right.transpose(2, 0, 1)
|
||||
imgL = np.ascontiguousarray(imgL[None, :, :, :]).astype("float32")
|
||||
imgR = np.ascontiguousarray(imgR[None, :, :, :]).astype("float32")
|
||||
|
||||
imgL_dw2 = imgL_dw2.transpose(2, 0, 1)
|
||||
imgR_dw2 = imgR_dw2.transpose(2, 0, 1)
|
||||
imgL_dw2 = np.ascontiguousarray(imgL_dw2[None, :, :, :]).astype("float32")
|
||||
imgR_dw2 = np.ascontiguousarray(imgR_dw2[None, :, :, :]).astype("float32")
|
||||
|
||||
print("Model Forwarding...")
|
||||
# First pass it just to get the flow
|
||||
pred_flow_dw2 = no_flow_model.run(
|
||||
[output_name], {input1_name: imgL_dw2, input2_name: imgR_dw2})[0]
|
||||
# Second pass gets us the disparity
|
||||
pred_disp = model.run([output_name], {
|
||||
input1_name: imgL, input2_name: imgR, input3_name: pred_flow_dw2})[0]
|
||||
|
||||
return np.squeeze(pred_disp[:, 0, :, :])
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
|
||||
left_img = cv2.imread("left.png")
|
||||
right_img = cv2.imread("right.png")
|
||||
|
||||
in_h, in_w = left_img.shape[:2]
|
||||
|
||||
# Resize images
|
||||
eval_h, eval_w = (in_h, in_w)
|
||||
assert eval_h % 8 == 0, "input height should be divisible by 8"
|
||||
assert eval_w % 8 == 0, "input width should be divisible by 8"
|
||||
|
||||
imgL = cv2.resize(left_img, (eval_w, eval_h),
|
||||
interpolation=cv2.INTER_LINEAR)
|
||||
imgR = cv2.resize(right_img, (eval_w, eval_h),
|
||||
interpolation=cv2.INTER_LINEAR)
|
||||
|
||||
no_flow_model_path = "models/crestereo_without_flow.onnx"
|
||||
model_path = "models/crestereo.onnx"
|
||||
|
||||
model = onnxruntime.InferenceSession(model_path)
|
||||
no_flow_model = onnxruntime.InferenceSession(no_flow_model_path)
|
||||
|
||||
pred = inference(imgL, imgR, model, no_flow_model)
|
||||
|
||||
t = float(in_w) / float(eval_w)
|
||||
disp = cv2.resize(pred, (eval_w, eval_h),
|
||||
interpolation=cv2.INTER_LINEAR) * t
|
||||
disp_vis = (disp - disp.min()) / (disp.max() - disp.min()) * 255.0
|
||||
disp_vis = disp_vis.astype("uint8")
|
||||
disp_vis = cv2.applyColorMap(disp_vis, cv2.COLORMAP_INFERNO)
|
||||
|
||||
combined_img = np.hstack((left_img, disp_vis))
|
||||
cv2.namedWindow("output", cv2.WINDOW_NORMAL)
|
||||
cv2.imshow("output", combined_img)
|
||||
cv2.imwrite("output.jpg", disp_vis)
|
||||
cv2.waitKey(0)
|
||||
@@ -0,0 +1,492 @@
|
||||
import argparse
|
||||
import os
|
||||
import shutil
|
||||
import sys
|
||||
import time
|
||||
import logging
|
||||
from collections import namedtuple
|
||||
from itertools import repeat
|
||||
|
||||
import yaml
|
||||
from tensorboardX import SummaryWriter
|
||||
|
||||
from nets import Model
|
||||
from dataset import CREStereoDataset
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.optim as optim
|
||||
import torch.backends.cudnn as cudnn
|
||||
import torch.distributed as dist
|
||||
from torch.nn.parallel import DistributedDataParallel as DDP
|
||||
from torch.utils.data import DataLoader, RandomSampler
|
||||
|
||||
|
||||
def parse_yaml(file_path: str) -> namedtuple:
|
||||
"""Parse yaml configuration file and return the object in `namedtuple`."""
|
||||
with open(file_path, "rb") as f:
|
||||
cfg: dict = yaml.safe_load(f)
|
||||
args = namedtuple("train_args", cfg.keys())(*cfg.values())
|
||||
# save cfg into train_log
|
||||
ensure_dir(args.log_dir)
|
||||
dst_file = os.path.join(args.log_dir, file_path.split('/')[-1])
|
||||
shutil.copy2(file_path, dst_file)
|
||||
return args
|
||||
|
||||
|
||||
def format_time(elapse):
|
||||
elapse = int(elapse)
|
||||
hour = elapse // 3600
|
||||
minute = elapse % 3600 // 60
|
||||
seconds = elapse % 60
|
||||
return "{:02d}:{:02d}:{:02d}".format(hour, minute, seconds)
|
||||
|
||||
|
||||
def ensure_dir(path):
|
||||
if not os.path.exists(path):
|
||||
os.makedirs(path, exist_ok=True)
|
||||
|
||||
|
||||
def adjust_learning_rate(optimizer, epoch):
|
||||
|
||||
warm_up = 0.02
|
||||
const_range = 0.6
|
||||
min_lr_rate = 0.05
|
||||
|
||||
if epoch <= args.n_total_epoch * warm_up:
|
||||
lr = (1 - min_lr_rate) * args.base_lr / (
|
||||
args.n_total_epoch * warm_up
|
||||
) * epoch + min_lr_rate * args.base_lr
|
||||
elif args.n_total_epoch * warm_up < epoch <= args.n_total_epoch * const_range:
|
||||
lr = args.base_lr
|
||||
else:
|
||||
lr = (min_lr_rate - 1) * args.base_lr / (
|
||||
(1 - const_range) * args.n_total_epoch
|
||||
) * epoch + (1 - min_lr_rate * const_range) / (1 - const_range) * args.base_lr
|
||||
|
||||
for param_group in optimizer.param_groups:
|
||||
param_group['lr'] = lr
|
||||
|
||||
def sequence_loss(flow_preds, flow_gt, valid, gamma=0.8):
|
||||
'''
|
||||
valid: (2, 384, 512) (B, H, W) -> (B, 1, H, W)
|
||||
flow_preds[0]: (B, 2, H, W)
|
||||
flow_gt: (B, 2, H, W)
|
||||
'''
|
||||
n_predictions = len(flow_preds)
|
||||
flow_loss = 0.0
|
||||
for i in range(n_predictions):
|
||||
i_weight = gamma ** (n_predictions - i - 1)
|
||||
i_loss = torch.abs(flow_preds[i] - flow_gt)
|
||||
flow_loss += i_weight * (valid.unsqueeze(1) * i_loss).mean()
|
||||
|
||||
return flow_loss
|
||||
|
||||
def repeater(data_loader):
|
||||
for loader in repeat(data_loader):
|
||||
for data in loader:
|
||||
yield data
|
||||
|
||||
def train_dist(args, world_size):
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--local_rank",type=int)
|
||||
FLAGS = parser.parse_args()
|
||||
local_rank = FLAGS.local_rank
|
||||
# directory check
|
||||
log_model_dir = os.path.join(args.log_dir, "models")
|
||||
ensure_dir(log_model_dir)
|
||||
|
||||
# distributed init and model / optimizer
|
||||
torch.cuda.set_device(local_rank)
|
||||
dist.init_process_group(backend='nccl') # nccl is highly recommanded
|
||||
model = Model(
|
||||
max_disp=args.max_disp, mixed_precision=args.mixed_precision, test_mode=False
|
||||
)
|
||||
# sync batch norm
|
||||
model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model).to(local_rank)
|
||||
model = DDP(model, device_ids=[local_rank], output_device=local_rank)
|
||||
optimizer = optim.Adam(model.parameters(), lr=0.1, betas=(0.9, 0.999))
|
||||
|
||||
if dist.get_rank() == 0:
|
||||
# tensorboard
|
||||
tb_log = SummaryWriter(os.path.join(args.log_dir, "train.events"))
|
||||
|
||||
# worklog
|
||||
logging.basicConfig(level=eval(args.log_level))
|
||||
worklog = logging.getLogger("train_logger")
|
||||
worklog.propagate = False
|
||||
fileHandler = logging.FileHandler(
|
||||
os.path.join(args.log_dir, "worklog.txt"), mode="a", encoding="utf8"
|
||||
)
|
||||
formatter = logging.Formatter(
|
||||
fmt="%(asctime)s %(message)s", datefmt="%Y/%m/%d %H:%M:%S"
|
||||
)
|
||||
fileHandler.setFormatter(formatter)
|
||||
consoleHandler = logging.StreamHandler(sys.stdout)
|
||||
formatter = logging.Formatter(
|
||||
fmt="\x1b[32m%(asctime)s\x1b[0m %(message)s", datefmt="%Y/%m/%d %H:%M:%S"
|
||||
)
|
||||
consoleHandler.setFormatter(formatter)
|
||||
worklog.handlers = [fileHandler, consoleHandler]
|
||||
|
||||
# params stat
|
||||
worklog.info(f"Use {world_size} GPU(s)")
|
||||
worklog.info("Params: %s" % sum([p.numel() for p in model.parameters()]))
|
||||
|
||||
# load pretrained model if exist
|
||||
chk_path = os.path.join(log_model_dir, "latest.pth")
|
||||
if args.loadmodel is not None:
|
||||
chk_path = args.loadmodel
|
||||
elif not os.path.exists(chk_path):
|
||||
chk_path = None
|
||||
|
||||
if chk_path is not None:
|
||||
if dist.get_rank() == 0:
|
||||
worklog.info(f"loading model: {chk_path}")
|
||||
# map_location=torch.device('cpu') make more balance memory usage
|
||||
state_dict = torch.load(chk_path, map_location=torch.device('cpu'))
|
||||
model.module.load_state_dict(state_dict['state_dict'])
|
||||
optimizer.load_state_dict(state_dict['optim_state_dict'])
|
||||
resume_epoch_idx = state_dict["epoch"]
|
||||
resume_iters = state_dict["iters"]
|
||||
start_epoch_idx = resume_epoch_idx + 1
|
||||
start_iters = resume_iters
|
||||
else:
|
||||
start_epoch_idx = 1
|
||||
start_iters = 0
|
||||
|
||||
# datasets
|
||||
dataset = CREStereoDataset(args.training_data_path)
|
||||
# dataset = MixDataset("train",
|
||||
# data_path=args.data["train"]["data_path"],
|
||||
# fields=args.data["train"]["fields"],
|
||||
# filelists=args.data["train"]["filelists"],
|
||||
# input_size=(args.data["train"]["input_size"][0], args.data["train"]["input_size"][1]))
|
||||
if dist.get_rank() == 0:
|
||||
worklog.info(f"Dataset size: {len(dataset)}")
|
||||
train_sampler = torch.utils.data.distributed.DistributedSampler(dataset)
|
||||
dataloader = torch.utils.data.DataLoader(dataset,
|
||||
batch_size=args.batch_size, num_workers=4, sampler=train_sampler)
|
||||
|
||||
# counter
|
||||
cur_iters = start_iters
|
||||
total_iters = args.minibatch_per_epoch * args.n_total_epoch
|
||||
t0 = time.perf_counter()
|
||||
for epoch_idx in range(start_epoch_idx, args.n_total_epoch + 1):
|
||||
dataloader.sampler.set_epoch(epoch_idx)
|
||||
# adjust learning rate
|
||||
epoch_total_train_loss = 0
|
||||
adjust_learning_rate(optimizer, epoch_idx)
|
||||
model.train()
|
||||
|
||||
t1 = time.perf_counter()
|
||||
|
||||
for batch_idx, mini_batch_data in enumerate(dataloader):
|
||||
|
||||
if batch_idx % args.minibatch_per_epoch == 0 and batch_idx != 0:
|
||||
break
|
||||
cur_iters += 1
|
||||
|
||||
# parse data
|
||||
left, right, gt_disp, valid_mask = (
|
||||
mini_batch_data["left"].to(local_rank),
|
||||
mini_batch_data["right"].to(local_rank),
|
||||
mini_batch_data["disparity"].to(local_rank),
|
||||
mini_batch_data["mask"].to(local_rank),
|
||||
)
|
||||
|
||||
t2 = time.perf_counter()
|
||||
optimizer.zero_grad()
|
||||
|
||||
# pre-process
|
||||
gt_disp = torch.unsqueeze(gt_disp, dim=1) # [2, 384, 512] -> [2, 1, 384, 512]
|
||||
gt_flow = torch.cat([gt_disp, gt_disp * 0], dim=1) # [2, 2, 384, 512]
|
||||
|
||||
# forward
|
||||
flow_predictions = model(left, right)
|
||||
|
||||
# loss & backword
|
||||
loss = sequence_loss(
|
||||
flow_predictions, gt_flow, valid_mask, gamma=0.8
|
||||
).to(local_rank)
|
||||
|
||||
# loss stats
|
||||
loss_item = loss.data.item()
|
||||
epoch_total_train_loss += loss_item
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
t3 = time.perf_counter()
|
||||
if dist.get_rank() == 0:
|
||||
if cur_iters % 10 == 0:
|
||||
tdata = t2 - t1
|
||||
time_train_passed = t3 - t0
|
||||
time_iter_passed = t3 - t1
|
||||
step_passed = cur_iters - start_iters
|
||||
eta = (
|
||||
(total_iters - cur_iters)
|
||||
/ max(step_passed, 1e-7)
|
||||
* time_train_passed
|
||||
)
|
||||
|
||||
meta_info = list()
|
||||
meta_info.append("{:.2g} b/s".format(1.0 / time_iter_passed))
|
||||
meta_info.append("passed:{}".format(format_time(time_train_passed)))
|
||||
meta_info.append("eta:{}".format(format_time(eta)))
|
||||
meta_info.append(
|
||||
"data_time:{:.2g}".format(tdata / time_iter_passed)
|
||||
)
|
||||
|
||||
meta_info.append(
|
||||
"lr:{:.5g}".format(optimizer.param_groups[0]["lr"])
|
||||
)
|
||||
meta_info.append(
|
||||
"[{}/{}:{}/{}]".format(
|
||||
epoch_idx,
|
||||
args.n_total_epoch,
|
||||
batch_idx,
|
||||
args.minibatch_per_epoch,
|
||||
)
|
||||
)
|
||||
loss_info = list()
|
||||
loss_info.append("{}:{:.4g}".format("total_loss", loss_item))
|
||||
# exp_name = ['\n' + os.path.basename(os.getcwd())]
|
||||
|
||||
info = [",".join(meta_info+loss_info)]
|
||||
worklog.info("".join(info))
|
||||
|
||||
# minibatch loss
|
||||
tb_log.add_scalar("train/loss_batch", loss_item, cur_iters)
|
||||
tb_log.add_scalar(
|
||||
"train/lr", optimizer.param_groups[0]["lr"], cur_iters
|
||||
)
|
||||
tb_log.flush()
|
||||
|
||||
t1 = time.perf_counter()
|
||||
|
||||
if dist.get_rank() == 0:
|
||||
# epoch loss
|
||||
tb_log.add_scalar(
|
||||
"train/loss",
|
||||
epoch_total_train_loss / args.minibatch_per_epoch,
|
||||
epoch_idx,
|
||||
)
|
||||
tb_log.flush()
|
||||
|
||||
# save model params
|
||||
ckp_data = {
|
||||
"epoch": epoch_idx,
|
||||
"iters": cur_iters,
|
||||
"batch_size": args.batch_size * world_size,
|
||||
"epoch_size": args.minibatch_per_epoch,
|
||||
"train_loss": epoch_total_train_loss / args.minibatch_per_epoch,
|
||||
"state_dict": model.module.state_dict(),
|
||||
"optim_state_dict": optimizer.state_dict(),
|
||||
}
|
||||
torch.save(ckp_data, os.path.join(log_model_dir, "latest.pth"))
|
||||
if epoch_idx % args.model_save_freq_epoch == 0:
|
||||
save_path = os.path.join(log_model_dir, "epoch-%d.pth" % epoch_idx)
|
||||
worklog.info(f"Model params saved: {save_path}")
|
||||
torch.save(ckp_data, save_path)
|
||||
if dist.get_rank() == 0:
|
||||
worklog.info("Training is done, exit.")
|
||||
|
||||
def train(args, world_size):
|
||||
# directory check
|
||||
log_model_dir = os.path.join(args.log_dir, "models")
|
||||
ensure_dir(log_model_dir)
|
||||
|
||||
# model / optimizer
|
||||
model = Model(
|
||||
max_disp=args.max_disp, mixed_precision=args.mixed_precision, test_mode=False
|
||||
)
|
||||
model = nn.DataParallel(model,device_ids=[i for i in range(world_size)])
|
||||
model.cuda()
|
||||
optimizer = optim.Adam(model.parameters(), lr=0.1, betas=(0.9, 0.999))
|
||||
|
||||
tb_log = SummaryWriter(os.path.join(args.log_dir, "train.events"))
|
||||
|
||||
# worklog
|
||||
logging.basicConfig(level=eval(args.log_level))
|
||||
worklog = logging.getLogger("train_logger")
|
||||
worklog.propagate = False
|
||||
fileHandler = logging.FileHandler(
|
||||
os.path.join(args.log_dir, "worklog.txt"), mode="a", encoding="utf8"
|
||||
)
|
||||
formatter = logging.Formatter(
|
||||
fmt="%(asctime)s %(message)s", datefmt="%Y/%m/%d %H:%M:%S"
|
||||
)
|
||||
fileHandler.setFormatter(formatter)
|
||||
consoleHandler = logging.StreamHandler(sys.stdout)
|
||||
formatter = logging.Formatter(
|
||||
fmt="\x1b[32m%(asctime)s\x1b[0m %(message)s", datefmt="%Y/%m/%d %H:%M:%S"
|
||||
)
|
||||
consoleHandler.setFormatter(formatter)
|
||||
worklog.handlers = [fileHandler, consoleHandler]
|
||||
|
||||
# params stat
|
||||
worklog.info(f"Use {world_size} GPU(s)")
|
||||
worklog.info("Params: %s" % sum([p.numel() for p in model.parameters()]))
|
||||
|
||||
# load pretrained model if exist
|
||||
chk_path = os.path.join(log_model_dir, "latest.pth")
|
||||
if args.loadmodel is not None:
|
||||
chk_path = args.loadmodel
|
||||
elif not os.path.exists(chk_path):
|
||||
chk_path = None
|
||||
|
||||
if chk_path is not None:
|
||||
worklog.info(f"loading model: {chk_path}")
|
||||
state_dict = torch.load(chk_path)
|
||||
model.module.load_state_dict(state_dict['state_dict'])
|
||||
optimizer.load_state_dict(state_dict['optim_state_dict'])
|
||||
resume_epoch_idx = state_dict["epoch"]
|
||||
resume_iters = state_dict["iters"]
|
||||
start_epoch_idx = resume_epoch_idx + 1
|
||||
start_iters = resume_iters
|
||||
else:
|
||||
start_epoch_idx = 1
|
||||
start_iters = 0
|
||||
|
||||
# datasets
|
||||
dataset = CREStereoDataset(args.training_data_path)
|
||||
sampler = RandomSampler(dataset, replacement=False)
|
||||
worklog.info(f"Dataset size: {len(dataset)}")
|
||||
dataloader = DataLoader(dataset, sampler=sampler, batch_size=args.batch_size*world_size,
|
||||
num_workers=0, drop_last=True, persistent_workers=False, pin_memory=True)
|
||||
dataloader = repeater(dataloader)
|
||||
|
||||
# counter
|
||||
cur_iters = start_iters
|
||||
total_iters = args.minibatch_per_epoch * args.n_total_epoch
|
||||
t0 = time.perf_counter()
|
||||
for epoch_idx in range(start_epoch_idx, args.n_total_epoch + 1):
|
||||
|
||||
# adjust learning rate
|
||||
epoch_total_train_loss = 0
|
||||
adjust_learning_rate(optimizer, epoch_idx)
|
||||
model.train()
|
||||
|
||||
t1 = time.perf_counter()
|
||||
|
||||
# for mini_batch_data in dataloader:
|
||||
for batch_idx, mini_batch_data in enumerate(dataloader):
|
||||
|
||||
if batch_idx % args.minibatch_per_epoch == 0 and batch_idx != 0:
|
||||
break
|
||||
cur_iters += 1
|
||||
|
||||
# parse data
|
||||
left, right, gt_disp, valid_mask = (
|
||||
mini_batch_data["left"].cuda(),
|
||||
mini_batch_data["right"].cuda(),
|
||||
mini_batch_data["disparity"].cuda(),
|
||||
mini_batch_data["mask"].cuda(),
|
||||
)
|
||||
|
||||
t2 = time.perf_counter()
|
||||
optimizer.zero_grad()
|
||||
|
||||
# pre-process
|
||||
gt_disp = torch.unsqueeze(gt_disp, dim=1) # [2, 384, 512] -> [2, 1, 384, 512]
|
||||
gt_flow = torch.cat([gt_disp, gt_disp * 0], dim=1) # [2, 2, 384, 512]
|
||||
|
||||
# forward
|
||||
flow_predictions = model(left, right)
|
||||
|
||||
# loss & backword
|
||||
loss = sequence_loss(
|
||||
flow_predictions, gt_flow, valid_mask, gamma=0.8
|
||||
)
|
||||
|
||||
# loss stats
|
||||
loss_item = loss.data.item()
|
||||
epoch_total_train_loss += loss_item
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
t3 = time.perf_counter()
|
||||
|
||||
if cur_iters % 10 == 0:
|
||||
tdata = t2 - t1
|
||||
time_train_passed = t3 - t0
|
||||
time_iter_passed = t3 - t1
|
||||
step_passed = cur_iters - start_iters
|
||||
eta = (
|
||||
(total_iters - cur_iters)
|
||||
/ max(step_passed, 1e-7)
|
||||
* time_train_passed
|
||||
)
|
||||
|
||||
meta_info = list()
|
||||
meta_info.append("{:.2g} b/s".format(1.0 / time_iter_passed))
|
||||
meta_info.append("passed:{}".format(format_time(time_train_passed)))
|
||||
meta_info.append("eta:{}".format(format_time(eta)))
|
||||
meta_info.append(
|
||||
"data_time:{:.2g}".format(tdata / time_iter_passed)
|
||||
)
|
||||
|
||||
meta_info.append(
|
||||
"lr:{:.5g}".format(optimizer.param_groups[0]["lr"])
|
||||
)
|
||||
meta_info.append(
|
||||
"[{}/{}:{}/{}]".format(
|
||||
epoch_idx,
|
||||
args.n_total_epoch,
|
||||
batch_idx,
|
||||
args.minibatch_per_epoch,
|
||||
)
|
||||
)
|
||||
loss_info = [" ==> {}:{:.4g}".format("loss", loss_item)]
|
||||
# exp_name = ['\n' + os.path.basename(os.getcwd())]
|
||||
|
||||
info = [",".join(meta_info)] + loss_info
|
||||
worklog.info("".join(info))
|
||||
|
||||
# minibatch loss
|
||||
tb_log.add_scalar("train/loss_batch", loss_item, cur_iters)
|
||||
tb_log.add_scalar(
|
||||
"train/lr", optimizer.param_groups[0]["lr"], cur_iters
|
||||
)
|
||||
tb_log.flush()
|
||||
|
||||
t1 = time.perf_counter()
|
||||
|
||||
tb_log.add_scalar(
|
||||
"train/loss",
|
||||
epoch_total_train_loss / args.minibatch_per_epoch,
|
||||
epoch_idx,
|
||||
)
|
||||
tb_log.flush()
|
||||
|
||||
# save model params
|
||||
ckp_data = {
|
||||
"epoch": epoch_idx,
|
||||
"iters": cur_iters,
|
||||
"batch_size": args.batch_size*world_size,
|
||||
"epoch_size": args.minibatch_per_epoch,
|
||||
"train_loss": epoch_total_train_loss / args.minibatch_per_epoch,
|
||||
"state_dict": model.module.state_dict(),
|
||||
"optim_state_dict": optimizer.state_dict(),
|
||||
}
|
||||
torch.save(ckp_data, os.path.join(log_model_dir, "latest.pth"))
|
||||
if epoch_idx % args.model_save_freq_epoch == 0:
|
||||
save_path = os.path.join(log_model_dir, "epoch-%d.pth" % epoch_idx)
|
||||
worklog.info(f"Model params saved: {save_path}")
|
||||
torch.save(ckp_data, save_path)
|
||||
|
||||
worklog.info("Training is done, exit.")
|
||||
|
||||
def main(args):
|
||||
# initial info
|
||||
torch.manual_seed(args.seed)
|
||||
torch.cuda.manual_seed(args.seed)
|
||||
world_size = torch.cuda.device_count() # number of GPU(s)
|
||||
cudnn.benchmark = True
|
||||
if args.dist and world_size > 1:
|
||||
train_dist(args, world_size)
|
||||
else:
|
||||
train(args, world_size)
|
||||
|
||||
if __name__ == "__main__":
|
||||
# train configuration
|
||||
args = parse_yaml("cfgs/train.yaml")
|
||||
main(args)
|
||||
@@ -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()
|
||||
|
||||
@@ -0,0 +1,14 @@
|
||||
__pycache__
|
||||
*.vscode
|
||||
cfg.yaml
|
||||
*.pth
|
||||
*.pkl
|
||||
*.pyc
|
||||
*.so
|
||||
*.egg-info
|
||||
# Ignore weight files but keep the folder structure
|
||||
weights/*
|
||||
!weights/.gitkeep
|
||||
output/
|
||||
.claude/
|
||||
cpp/build/
|
||||
@@ -0,0 +1,94 @@
|
||||
Copyright (c) 2026-Present, NVIDIA Corporation & affiliates. All rights reserved.
|
||||
|
||||
|
||||
=======================================================================
|
||||
|
||||
1. Definitions
|
||||
|
||||
"Licensor" means any person or entity that distributes its Work.
|
||||
|
||||
"Software" means the original work of authorship made available under
|
||||
this License.
|
||||
|
||||
"Work" means the Software and any additions to or derivative works of
|
||||
the Software that are made available under this License.
|
||||
|
||||
The terms "reproduce," "reproduction," "derivative works," and
|
||||
"distribution" have the meaning as provided under U.S. copyright law;
|
||||
provided, however, that for the purposes of this License, derivative
|
||||
works shall not include works that remain separable from, or merely
|
||||
link (or bind by name) to the interfaces of, the Work.
|
||||
|
||||
Works, including the Software, are "made available" under this License
|
||||
by including in or with the Work either (a) a copyright notice
|
||||
referencing the applicability of this License to the Work, or (b) a
|
||||
copy of this License.
|
||||
|
||||
2. License Grants
|
||||
|
||||
2.1 Copyright Grant. Subject to the terms and conditions of this
|
||||
License, each Licensor grants to you a perpetual, worldwide,
|
||||
non-exclusive, royalty-free, copyright license to reproduce,
|
||||
prepare derivative works of, publicly display, publicly perform,
|
||||
sublicense and distribute its Work and any resulting derivative
|
||||
works in any form.
|
||||
|
||||
3. Limitations
|
||||
|
||||
3.1 Redistribution. You may reproduce or distribute the Work only
|
||||
if (a) you do so under this License, (b) you include a complete
|
||||
copy of this License with your distribution, and (c) you retain
|
||||
without modification any copyright, patent, trademark, or
|
||||
attribution notices that are present in the Work.
|
||||
|
||||
3.2 Derivative Works. You may specify that additional or different
|
||||
terms apply to the use, reproduction, and distribution of your
|
||||
derivative works of the Work ("Your Terms") only if (a) Your Terms
|
||||
provide that the use limitation in Section 3.3 applies to your
|
||||
derivative works, and (b) you identify the specific derivative
|
||||
works that are subject to Your Terms. Notwithstanding Your Terms,
|
||||
this License (including the redistribution requirements in Section
|
||||
3.1) will continue to apply to the Work itself.
|
||||
|
||||
3.3 Use Limitation. The Work and any derivative works thereof only
|
||||
may be used or intended for use non-commercially. Notwithstanding
|
||||
the foregoing, NVIDIA and its affiliates may use the Work and any
|
||||
derivative works commercially. As used herein, "non-commercially"
|
||||
means for research purposes only.
|
||||
|
||||
3.4 Patent Claims. If you bring or threaten to bring a patent claim
|
||||
against any Licensor (including any claim, cross-claim or
|
||||
counterclaim in a lawsuit) to enforce any patents that you allege
|
||||
are infringed by any Work, then your rights under this License from
|
||||
such Licensor (including the grant in Section 2.1) will terminate
|
||||
immediately.
|
||||
|
||||
3.5 Trademarks. This License does not grant any rights to use any
|
||||
Licensor�s or its affiliates� names, logos, or trademarks, except
|
||||
as necessary to reproduce the notices described in this License.
|
||||
|
||||
3.6 Termination. If you violate any term of this License, then your
|
||||
rights under this License (including the grant in Section 2.1) will
|
||||
terminate immediately.
|
||||
|
||||
4. Disclaimer of Warranty.
|
||||
|
||||
THE WORK IS PROVIDED "AS IS" WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
KIND, EITHER EXPRESS OR IMPLIED, INCLUDING WARRANTIES OR CONDITIONS OF
|
||||
MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE, TITLE OR
|
||||
NON-INFRINGEMENT. YOU BEAR THE RISK OF UNDERTAKING ANY ACTIVITIES UNDER
|
||||
THIS LICENSE.
|
||||
|
||||
5. Limitation of Liability.
|
||||
|
||||
EXCEPT AS PROHIBITED BY APPLICABLE LAW, IN NO EVENT AND UNDER NO LEGAL
|
||||
THEORY, WHETHER IN TORT (INCLUDING NEGLIGENCE), CONTRACT, OR OTHERWISE
|
||||
SHALL ANY LICENSOR BE LIABLE TO YOU FOR DAMAGES, INCLUDING ANY DIRECT,
|
||||
INDIRECT, SPECIAL, INCIDENTAL, OR CONSEQUENTIAL DAMAGES ARISING OUT OF
|
||||
OR RELATED TO THIS LICENSE, THE USE OR INABILITY TO USE THE WORK
|
||||
(INCLUDING BUT NOT LIMITED TO LOSS OF GOODWILL, BUSINESS INTERRUPTION,
|
||||
LOST PROFITS OR DATA, COMPUTER FAILURE OR MALFUNCTION, OR ANY OTHER
|
||||
COMMERCIAL DAMAGES OR LOSSES), EVEN IF THE LICENSOR HAS BEEN ADVISED OF
|
||||
THE POSSIBILITY OF SUCH DAMAGES.
|
||||
|
||||
=======================================================================
|
||||
Executable
+88
@@ -0,0 +1,88 @@
|
||||
import os, sys, torch, imageio, logging, importlib, argparse
|
||||
import cv2
|
||||
import numpy as np
|
||||
import yaml
|
||||
try:
|
||||
import open3d as o3d
|
||||
except:
|
||||
o3d = None
|
||||
|
||||
AMP_DTYPE = torch.float16
|
||||
|
||||
def set_logging_format(level=logging.INFO):
|
||||
importlib.reload(logging)
|
||||
FORMAT = '%(message)s'
|
||||
logging.basicConfig(level=level, format=FORMAT, datefmt='%m-%d|%H:%M:%S')
|
||||
|
||||
|
||||
def set_seed(random_seed):
|
||||
import torch,random
|
||||
np.random.seed(random_seed)
|
||||
random.seed(random_seed)
|
||||
torch.manual_seed(random_seed)
|
||||
torch.cuda.manual_seed_all(random_seed)
|
||||
torch.backends.cudnn.deterministic = True
|
||||
torch.backends.cudnn.benchmark = False
|
||||
|
||||
|
||||
def toOpen3dCloud(points,colors=None,normals=None):
|
||||
cloud = o3d.geometry.PointCloud()
|
||||
cloud.points = o3d.utility.Vector3dVector(points.astype(np.float64))
|
||||
if colors is not None:
|
||||
if colors.max()>1:
|
||||
colors = colors/255.0
|
||||
cloud.colors = o3d.utility.Vector3dVector(colors.astype(np.float64))
|
||||
if normals is not None:
|
||||
cloud.normals = o3d.utility.Vector3dVector(normals.astype(np.float64))
|
||||
return cloud
|
||||
|
||||
|
||||
|
||||
def depth2xyzmap(depth:np.ndarray, K, uvs:np.ndarray=None, zmin=0.1):
|
||||
invalid_mask = (depth<zmin)
|
||||
H,W = depth.shape[:2]
|
||||
if uvs is None:
|
||||
vs,us = np.meshgrid(np.arange(0,H),np.arange(0,W), sparse=False, indexing='ij')
|
||||
vs = vs.reshape(-1)
|
||||
us = us.reshape(-1)
|
||||
else:
|
||||
uvs = uvs.round().astype(int)
|
||||
us = uvs[:,0]
|
||||
vs = uvs[:,1]
|
||||
zs = depth[vs,us]
|
||||
xs = (us-K[0,2])*zs/K[0,0]
|
||||
ys = (vs-K[1,2])*zs/K[1,1]
|
||||
pts = np.stack((xs.reshape(-1),ys.reshape(-1),zs.reshape(-1)), 1) #(N,3)
|
||||
xyz_map = np.zeros((H,W,3), dtype=np.float32)
|
||||
xyz_map[vs,us] = pts
|
||||
if invalid_mask.any():
|
||||
xyz_map[invalid_mask] = 0
|
||||
return xyz_map
|
||||
|
||||
|
||||
def vis_disparity(disp, min_val=None, max_val=None, invalid_thres=np.inf, color_map=cv2.COLORMAP_TURBO, cmap=None, other_output={}):
|
||||
"""
|
||||
@disp: np array (H,W)
|
||||
@invalid_thres: > thres is invalid
|
||||
"""
|
||||
disp = disp.copy()
|
||||
H,W = disp.shape[:2]
|
||||
invalid_mask = disp>=invalid_thres
|
||||
if (invalid_mask==0).sum()==0:
|
||||
other_output['min_val'] = None
|
||||
other_output['max_val'] = None
|
||||
return np.zeros((H,W,3))
|
||||
if min_val is None:
|
||||
min_val = disp[invalid_mask==0].min()
|
||||
if max_val is None:
|
||||
max_val = disp[invalid_mask==0].max()
|
||||
other_output['min_val'] = min_val
|
||||
other_output['max_val'] = max_val
|
||||
vis = ((disp-min_val)/(max_val-min_val)).clip(0,1) * 255
|
||||
if cmap is None:
|
||||
vis = cv2.applyColorMap(vis.clip(0, 255).astype(np.uint8), color_map)[...,::-1]
|
||||
else:
|
||||
vis = cmap(vis.astype(np.uint8))[...,:3]*255
|
||||
if invalid_mask.any():
|
||||
vis[invalid_mask] = 0
|
||||
return vis.astype(np.uint8)
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 115 KiB |
Executable
BIN
Binary file not shown.
|
After Width: | Height: | Size: 397 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 2.9 MiB |
Binary file not shown.
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user