Добавлены пропсы конвейера и стереодвижки, задействованные в прогоне

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>

.gitignore

@@ -5,9 +5,8 @@
 *.ckpt
 *.safetensors
 
-# Пропсы конвейера NVIDIA (274 МБ): не наши и скачиваются штатным скриптом
-assets/conveyors/
-# восстановить: python3 scripts/fetch_assets.py
+# Наборы кадров прежних прогонов CV (1.26 ГБ) - выход, не исходники
+cv/flow_*/
 
 # Выход прогонов - воспроизводится, в истории не нужен
 control_test/captures/

assets/conveyors/Material Library/Metal/Misc/Metal_Rough_A.mdl

@@ -0,0 +1,16 @@
+mdl 1.4;
+
+using ::OmniPBR import OmniPBR;
+import ::tex::gamma_mode;
+import ::state::normal;
+import ::base::mono_mode;
+
+export material Metal_Rough_A(*)
+ = OmniPBR(
+	diffuse_texture: texture_2d("../Textures/T_Metal_Rough_A_Albedo.png", ::tex::gamma_srgb),
+	enable_ORM_texture: true,
+	metallic_texture_influence: 1.0,
+	normalmap_texture: texture_2d("../Textures/T_Metal_Rough_A_Normal.png", ::tex::gamma_linear),
+	ORM_texture: texture_2d("../Textures/T_Metal_Rough_A_ORM.png", ::tex::gamma_linear),
+	reflection_roughness_texture_influence: 1.0,
+	texture_scale: float2(2, 2));

assets/conveyors/Material Library/Metal/Painted/MetalPainted_White_Glossy_A.mdl

@@ -0,0 +1,15 @@
+mdl 1.4;
+
+using ::OmniPBR import OmniPBR;
+import ::tex::gamma_mode;
+import ::state::normal;
+import ::base::mono_mode;
+
+export material MetalPainted_White_Glossy_A(*)
+ = OmniPBR(
+	diffuse_texture: texture_2d("./T_MetalPainted_White_Glossy_A_Albedo.png", ::tex::gamma_srgb),
+	enable_ORM_texture: true,
+	metallic_texture_influence: 1.0,
+	normalmap_texture: texture_2d("./T_MetalPainted_White_Glossy_A_Normal.png", ::tex::gamma_linear),
+	ORM_texture: texture_2d("./T_MetalPainted_White_Glossy_A_ORM.png", ::tex::gamma_linear),
+	reflection_roughness_texture_influence: 1.0);

assets/conveyors/Material Library/Metal/Steel/Steel_A.mdl

@@ -0,0 +1,15 @@
+mdl 1.4;
+
+using ::OmniPBR import OmniPBR;
+import ::tex::gamma_mode;
+import ::state::normal;
+import ::base::mono_mode;
+
+export material Steel_A(*)
+ = OmniPBR(
+	diffuse_texture: texture_2d("../Textures/T_Steel_A1_Albedo.png", ::tex::gamma_srgb),
+	enable_ORM_texture: true,
+	metallic_texture_influence: 1.0,
+	normalmap_texture: texture_2d("../Textures/T_Steel_A1_Normal.png", ::tex::gamma_linear),
+	ORM_texture: texture_2d("../Textures/T_Steel_A1_ORM.png", ::tex::gamma_linear),
+	reflection_roughness_texture_influence: 1.0);

assets/conveyors/Material Library/Plastic/Misc/Plastic_Orange_A.mdl

@@ -0,0 +1,52 @@
+mdl 1.4;
+
+using ::OmniPBR import OmniPBR;
+import ::tex::gamma_mode;
+import ::state::normal;
+import ::base::mono_mode;
+
+export material Plastic_Orange_A(*)
+ = OmniPBR(
+	diffuse_color_constant: color(0.2, 0.2, 0.2),
+	diffuse_texture: texture_2d("../Textures/T_Plastic_Orange_A_Albedo.png", ::tex::gamma_srgb),
+	albedo_desaturation: 0.0,
+	albedo_add: 0.0,
+	albedo_brightness: 1.0,
+	diffuse_tint: color(1, 1, 1),
+	reflection_roughness_constant: 0.5,
+	reflection_roughness_texture_influence: 1.0,
+	reflectionroughness_texture: texture_2d(),
+	metallic_constant: 0.0,
+	metallic_texture_influence: 1.0,
+	metallic_texture: texture_2d(),
+	specular_level: 0.5,
+	enable_ORM_texture: true,
+	ORM_texture: texture_2d("../Textures/T_Plastic_Orange_A_ORM.png", ::tex::gamma_linear),
+	ao_to_diffuse: 0.0,
+	ao_texture: texture_2d(),
+	enable_emission: false,
+	emissive_color: color(1, 0.1, 0.1),
+	emissive_color_texture: texture_2d(),
+	emissive_mask_texture: texture_2d(),
+	emissive_intensity: 40.0,
+	enable_opacity: false,
+	enable_opacity_texture: false,
+	opacity_constant: 1.0,
+	opacity_texture: texture_2d(),
+	opacity_mode: ::base::mono_average,
+	opacity_threshold: 0.0,
+	bump_factor: 1.0,
+	normalmap_texture: texture_2d("../Textures/T_Plastic_Orange_A_Normal.png", ::tex::gamma_linear),
+	detail_bump_factor: 0.30000001192092896,
+	detail_normalmap_texture: texture_2d(),
+	flip_tangent_u: false,
+	flip_tangent_v: true,
+	project_uvw: false,
+	world_or_object: false,
+	uv_space_index: 0,
+	texture_translate: float2(0, 0),
+	texture_rotate: 0.0,
+	texture_scale: float2(3, 3),
+	detail_texture_translate: float2(0, 0),
+	detail_texture_rotate: 0.0,
+	detail_texture_scale: float2(1, 1));

assets/conveyors/Material Library/Plastic/Misc/Plastic_Red_A.mdl

@@ -0,0 +1,20 @@
+mdl 1.4;
+
+using ::OmniPBR import OmniPBR;
+import ::tex::gamma_mode;
+import ::state::normal;
+import ::base::mono_mode;
+
+export material Plastic_Red_A(*)
+ = OmniPBR(
+	albedo_add: 0.4099999964237213,
+	albedo_brightness: 0.9099999666213989,
+	diffuse_texture: texture_2d("../Textures/T_Plastic_Red_A_Albedo.png", ::tex::gamma_srgb),
+	diffuse_tint: color(0.2680115, 0.053293347, 0.053293347),
+	enable_ORM_texture: true,
+	metallic_constant: 0.0,
+	metallic_texture_influence: 1.0,
+	normalmap_texture: texture_2d("../Textures/T_Plastic_Red_A_Normal.png", ::tex::gamma_linear),
+	ORM_texture: texture_2d("../Textures/T_Plastic_Red_A_ORM.png", ::tex::gamma_linear),
+	reflection_roughness_texture_influence: 1.0,
+	texture_scale: float2(3, 3));

assets/conveyors/Material Library/Plastic/Misc/Plastic_Rough_Black_A.mdl

@@ -0,0 +1,52 @@
+mdl 1.4;
+
+using ::OmniPBR import OmniPBR;
+import ::tex::gamma_mode;
+import ::state::normal;
+import ::base::mono_mode;
+
+export material Plastic_Rough_Black_A(*)
+ = OmniPBR(
+	diffuse_color_constant: color(0.2, 0.2, 0.2),
+	diffuse_texture: texture_2d("../Textures/T_Plastic_Rough_Black_A_Albedo.png", ::tex::gamma_srgb),
+	albedo_desaturation: 0.0,
+	albedo_add: 0.0,
+	albedo_brightness: 1.0,
+	diffuse_tint: color(1, 1, 1),
+	reflection_roughness_constant: 0.5,
+	reflection_roughness_texture_influence: 1.0,
+	reflectionroughness_texture: texture_2d(),
+	metallic_constant: 0.0,
+	metallic_texture_influence: 1.0,
+	metallic_texture: texture_2d(),
+	specular_level: 0.5,
+	enable_ORM_texture: true,
+	ORM_texture: texture_2d("../Textures/T_Plastic_Rough_Black_A_ORM.png", ::tex::gamma_linear),
+	ao_to_diffuse: 0.0,
+	ao_texture: texture_2d(),
+	enable_emission: false,
+	emissive_color: color(1, 0.1, 0.1),
+	emissive_color_texture: texture_2d(),
+	emissive_mask_texture: texture_2d(),
+	emissive_intensity: 40.0,
+	enable_opacity: false,
+	enable_opacity_texture: false,
+	opacity_constant: 1.0,
+	opacity_texture: texture_2d(),
+	opacity_mode: ::base::mono_average,
+	opacity_threshold: 0.0,
+	bump_factor: 2.5,
+	normalmap_texture: texture_2d("../Textures/T_Plastic_Rough_Black_A_Normal.png", ::tex::gamma_linear),
+	detail_bump_factor: 0.30000001192092896,
+	detail_normalmap_texture: texture_2d(),
+	flip_tangent_u: false,
+	flip_tangent_v: true,
+	project_uvw: false,
+	world_or_object: false,
+	uv_space_index: 0,
+	texture_translate: float2(0, 0),
+	texture_rotate: 0.0,
+	texture_scale: float2(1, 1),
+	detail_texture_translate: float2(0, 0),
+	detail_texture_rotate: 0.0,
+	detail_texture_scale: float2(1, 1));

cv/circular_section.py

@@ -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)

cv/circular_section_results.json

@@ -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
+  ]
+]

cv/crestereo/.gitattributes

@@ -0,0 +1,2 @@
+# Auto detect text files and perform LF normalization
+* text=auto

cv/crestereo/.gitignore

@@ -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/*

cv/crestereo/README.md

@@ -0,0 +1,46 @@
+# CREStereo-Pytorch
+ Non-official Pytorch implementation of the CREStereo (CVPR 2022 Oral) model converted from the original MegEngine implementation.
+
+![!CREStereo-Pytorch stereo detph estimation](https://github.com/ibaiGorordo/CREStereo-Pytorch/blob/main/doc/img/output.jpg)
+
+**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

cv/crestereo/cfgs/train.yaml

@@ -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

cv/crestereo/convert_to_onnx.py

@@ -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'])
+
+
+
+

cv/crestereo/convert_weights.py

@@ -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}")

cv/crestereo/dataset.py

@@ -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)

cv/crestereo/function_convertion_tests/test_bilinear_sampler.py

@@ -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()

cv/crestereo/function_convertion_tests/test_coords_grid.py

@@ -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()

cv/crestereo/function_convertion_tests/test_manual_pad.py

@@ -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()

cv/crestereo/function_convertion_tests/test_meshgrid.py

@@ -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()

cv/crestereo/function_convertion_tests/test_offset.py

@@ -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()

cv/crestereo/function_convertion_tests/test_split.py

@@ -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()

cv/crestereo/nets/__init__.py

@@ -0,0 +1 @@
+from .crestereo import CREStereo as Model

cv/crestereo/nets/attention/__init__.py

@@ -0,0 +1,2 @@
+from .transformer import LocalFeatureTransformer
+from .position_encoding import PositionEncodingSine

cv/crestereo/nets/attention/linear_attention.py

@@ -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()

cv/crestereo/nets/attention/position_encoding.py

@@ -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)

cv/crestereo/nets/attention/transformer.py

@@ -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

cv/crestereo/nets/corr.py

@@ -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

cv/crestereo/nets/crestereo.py

@@ -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

cv/crestereo/nets/extractor.py

@@ -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

cv/crestereo/nets/update.py

@@ -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

cv/crestereo/nets/utils/__init__.py

@@ -0,0 +1 @@
+from .utils import bilinear_sampler, coords_grid, manual_pad

cv/crestereo/nets/utils/utils.py

@@ -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)

cv/crestereo/test_model.py

@@ -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)
+
+
+

cv/crestereo/test_onnx_model.py

@@ -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)

cv/crestereo/train.py

@@ -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)

cv/defom-stereo/LICENSE

@@ -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.

cv/defom-stereo/README.md

@@ -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}
+}
+```

cv/defom-stereo/core/corr.py

@@ -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())

cv/defom-stereo/core/defom_stereo.py

@@ -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

cv/defom-stereo/core/extractor.py

@@ -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

cv/defom-stereo/core/stereo_datasets.py

@@ -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
+

cv/defom-stereo/core/update.py

@@ -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

cv/defom-stereo/core/utils/augmentor.py

@@ -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

cv/defom-stereo/core/utils/dist_utils.py

@@ -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

cv/defom-stereo/core/utils/frame_utils.py

@@ -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 []

cv/defom-stereo/core/utils/utils.py

@@ -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

cv/defom-stereo/demo.py

@@ -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)

cv/defom-stereo/depth_anything_v2/dinov2.py

@@ -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
+    )

cv/defom-stereo/depth_anything_v2/dinov2_layers/__init__.py

@@ -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